[TE] Adding function attributes to TE schedules for use in LLVM codegen

Questions
1

Hi, I want to ensure that when a specific TE schedule is used during compilation, LLVM codegen inserts a function-level attribute on the resulting LLVM function that is generated. Currently I’ve been doing this by inserting a pragma on the outer loop of the schedule where it can then be picked up in LLVM codegen and added to the current function being codegenned. I’ve tried to illustrate this below with a trimmed down (non-runnable) example:

TE schedule:

def my_schedule(outs):
    sched = te.create_schedule([x.op for x in outs])
    
    def _callback(op):
        a = sched[op].op.axis
        sched[op].pragma(a, "my_function_attribute")

     traverse_inline(sched, outs[-1].op, _callback)
     return sched

Scheduled TIR:

@T.prim_func
def tvmgen_default(in: T.Buffer((64, 32), "float32"), out: T.Buffer((64, 128), "float32")):
    T.func_attr({"from_legacy_te_schedule": T.bool(True), "global_symbol": "tvmgen_default", "tir.noalias": T.bool(True)})
    compute = T.allocate([4096], "float32", "global")
    compute_1 = T.Buffer((4096,), data=compute)
    for a in range(4):
        a = T.int32()
        T.attr(T.iter_var(a, None, "DataPar", ""), "pragma_my_function_attribute")
     ...(do something useful)...

This pragma can then be picked up during codegen by visiting each AttrStmtNode and the LLVM function attribute can be added to the function currently being generated:

void CodeGenCPU::VisitStmt_(const AttrStmtNode* op) {
  String attr_key = op->attr_key;
  if (tir::attr::IsPragmaKey(attr_key)) {
    if (attr_key == "pragma_my_function_attribute") {
      function_->addFnAttr("my_function_attribute");
      this->VisitStmt(op->body);
    }
  } else {
    CodeGenLLVM::VisitStmt_(op);
  }
}

This works well and it achieves the desired result, however it seems to me to be more of a hack. My thinking is that pragmas operate at the loop level, while I want to add annotations at the function level. A motivating example for adding LLVM function attributes in the first place can be found here.

Adding the attribute to T.func_attr({"from_legacy_te_schedule": T.bool(True), "global_symbol": "tvmgen_default", "tir.noalias": T.bool(True)}) and picking that up in the LLVM codegen backend makes more sense to me, however, I wasn’t able to find a way to make changes to func_attr from within TE.

Curious if anyone has experienced similar issues or if there are any approaches I’ve not considered? I’m fairly new to TE/TIR so any suggestions are welcome :slight_smile:

2

If you’re okay with generating and working with a PrimFunc, you could use the with_attr attribute of PrimFunc.

From your example, something like the below could work?

# Instead of a te schedule, create a PrimFunc
# sched = te.create_schedule([x.op for x in outs])
func = te.create_prim_func([x.op for x in outs])
func_with_attr = func.with_attr({"some_attr": "attr_value"})
sched = tvm.tir.Schedule(func_with_attr)
# Use TIR schedule primitives instead

If you would prefer to keep schedules in te, there is a tvm.driver.build_module.schedule_to_module() function which can be used to create an IRModule after you schedules and the attribute can be added to this before returning

# Instead of a te schedule, create a PrimFunc
sched = te.create_schedule([x.op for x in outs])
# Apply your te schedules
mod = tvm.driver.build_module.schedule_to_module(sched, [x.op for x in outs])
func = mod["main"]
func_with_attr = func.with_attr({"some_attr": "attr_value"})
mod.update_func(mod.get_global_var("main"), func_with_attr)
# This is the scheduled mod with attribute attached to the function

P.S. The schedule_to_module function is mentioned as only for debugging in the comments, but maybe in this case, its okay to use it as long as we know that the normal optimization passes are not yet run, and this is just the mod with scheduling changes.

I still think the first approach of directly using create_prim_func and using tir schedules is probably better, but if you already have a lot of schedules written in te, this (second approach) might be a temporary alternative.

3

I agree that we should move new workloads to the tir schedule flow :slight_smile: which generally gets better and have everything we need in te schedule.

1 Like
4

Thanks both, I’ll take a look into the suggestions and update when I get something working. Appreciate you taking the time to answer!