[TensorIR] CreatePrimFunc from TE

[Hzfengsy]

This PR introduces a method to create PrimFunc (and further TensorIR schedule) from Tensor Expression. (e.g. te.compute)

cc @tqchen @junrushao1994 @comaniac @jcf94 @xqdan

[tkonolige]

Would it be possible to add a name field to the create_prim_func function? That way when we print out primfuncs for debugging, they aren't all called "main".

[tqchen]

Create prim_func returns a PrimFunc, the name is attached in module creation process, we can likely recompose by

mod = IRModule({"mycoolname": create_prim_func())

Note such name is different from the global_symbol(name of generated function), which can be set by create_prim_func().with_attr("tir.global_symbol", global_symbol)

[tkonolige]

Instead of create_prim_func().with_attr("tir.global_symbol", global_symbol), I think we should just make create_prim_func take a name parameter. I was working on autoscheduler earlier and all the primfuncs it had were called "main". Having a name parameter encourages the user to name their functions and makes debugging easier.

[tqchen]

If we really want to encourage the name, and given IRModule is the first class item for running exchange between passes. Let us consider make the function return an IRModule instead, in this case, perhaps we should rename to create_tir (coming back to another old name)

[junrushao]

@tkonolige The mechanism causing the issue you mentioned is that when constructing a tir::Schedule class, and the input is a PrimFunc, then it will put that PrimFunc into an IRModule with the default name main, so that the schedule is always IRModule-based.

[junrushao]

@tqchen The problem of creating an IRModule instead is that we might want to an IRModule that contains multiple PrimFuncs, e.g.:

mod = IRModule({
  "func_1": create_prim_func(...),
  "func_2": create_prim_func(...),
})

[tqchen]

I see, i now agree that we can go with the compositional style API. While I agree with @tkonolige that naming functions is going to be useful. We can do that through the compositional API and when creating example tutorials for scheduling. Thanks @junrushao1994 @Hzfengsy

[junrushao]

Thanks for the great discussion @tkonolige @tqchen @Hzfengsy! I am going to merge it in given there is no objection :-)