Lambda in tvm.compute

Questions
1

How the lambda in tvm.compute func transform into c++ expression? like following code:

B = tvm.compute((n,), lambda i: tvm.sum(A[i, k], axis=k), name="B")
1 Like
2

checkout https://docs.tvm.ai/dev/runtime.html

3

OK, I go to the site and have a look, thank you~

4

I read the article, and execute the code step by step:

B = tvm.compute((n,), lambda i: tvm.sum(A[i, k], axis=k), name="B")

I encounter the following questions in order:

  1. code walked into compute func and execute following two codes:
body = convert(body) // body's type is tvm.container.Array
op_node = _api_internal._ComputeOp(name, tag, attrs, dim_var, body) 
// op_node's type is tvm.tensor.ComputeOp
  1. I found the registered _ComputeOp func in src//api/api_lang.cc:268:TVM_REGISTER_API("_ComputeOp"), which calls following code in ComputeOpNode class:
  static Operation make(std::string name,
                        std::string tag,
                        Map<std::string, NodeRef> attrs,
                        Array<IterVar> axis,
                        Array<Expr> body);

the return value is Operation which mismatches op_node’s type.

two questions:

  1. the article https://docs.tvm.ai/dev/runtime.html talk about calling python func in C++ code using PackedFunc, but the lambda lambda i: tvm.sum(A[i, k], axis=k) defined in tvm.compute is not being called in c++ code, or I don’t know how it’s called in C++ code.
  2. why op_node’s type is not the same with _api_internal._ComputeOp return type

thank you very much~

5

lambda in tvm.compute is only invoked in Python side, see https://github.com/dmlc/tvm/blob/53ac89ede7cddd1649b01b2ff10cc67a963757ce/python/tvm/api.py#L262

tvm.tensor.ComputeOp is defined in python/tvm/tensor.py, you may want to check python/tvm/_ffi/node.py to see how Node in C++ is mapped to the python type