Inserting config instructions in a TIR schedule

Questions
#1

With TE there was the option to mark certain loops with a pragma which can then be matched to inject a function call. Is there something similar with TIR?

#2

Hi, I suspect annotate is what you’re looking for? See an example here: https://github.com/apache/tvm/blob/bd67d2e5ebde1aec18bcfa74c087516579bda1ae/tests/python/meta_schedule/test_meta_schedule_trace_apply.py#L2957

#3

Thanks for the hint! I tried annotating a loop with one of the commands, but that doesn’t seem to change anything about the generated C code. Here is the code snippet I am using to test my stuff in C:

sch.annotate(mvin_data_ax1, ann_key="pragma_unroll_explicit", ann_val=2)

#################################
##  Lower to C for inspection  ##
#################################
inp_buf = tvm.tir.decl_buffer(inp.shape, inp.dtype, name="inp_buf")
wght_buf = tvm.tir.decl_buffer(wght.shape, wght.dtype, name="wght_buf")
bias_buf = tvm.tir.decl_buffer(bias.shape, bias.dtype, name="bias_buf")
res_buf = tvm.tir.decl_buffer(res.shape, res.dtype, name="res_buf")

# Set stm32F7 as the target device, Use the AOT executor rather than graph or vm executors.
RUNTIME = Runtime("crt", {"system-lib": False})
TARGET = tvm.target.target.stm32("stm32F7xx")
EXECUTOR = Executor("aot", options={"interface-api": "c", "unpacked-api": 1})

# Generate Lowering Passes
pass_list = [
        (0, tvm.tir.transform.StorageFlatten(16)),
        (1, MvinIntrin()),
    ]
config = {"tir.add_lower_pass": pass_list,
            "tir.disable_vectorize": True,}
ctx = tvm.transform.PassContext(config=config, opt_level=3)
with ctx:
    matmul = tvm.build(sch.mod, 
                [inp, wght, bias,res],
                name="matmul_test",
                binds={inp: inp_buf, wght:wght_buf, bias:bias_buf, res: res_buf},
                target=TARGET)

# Generate the Project
BUILD_DIR = pathlib.Path(pathlib.Path.cwd())
BUILD_DIR.mkdir(exist_ok=True)

# Now, we export the model into a tar file:
TAR_PATH = pathlib.Path(BUILD_DIR) / "model.tar"
export_model_library_format(matmul, TAR_PATH)

I think I am missing something about the way I am supposed to use annotations, but I can’t figure out what exactly they are doing.

#4

Apologies that example may have been a bad one as the annotation relates to the meta scheduler, so I wouldn’t anticipate any changes in codegen using your example above.

It might be better to try something like “pragma_import_c” which the c codegen uses here: https://github.com/apache/tvm/blob/d42fb60e927c3c081f47de0fa798735888c551dc/src/target/source/codegen_c.cc#L987. There should be a few examples of its usage around the codebase e.g. https://github.com/apache/tvm/blob/958c27123a45a9629e57cee20dbca28263c836bd/python/tvm/topi/arm_cpu/qnn.py#L525

#5

Thank you! Using pragma_import_c does indeed result in the code showing up. It’s still not quite what I want, I would like the code to show up directly on the loop, while this statement seems to be more useful for bringing functions into scope or importing headers. Do you know if there is a way to emit C directly on a loop while operating on a TIR schedule? The examples I can find all do this directly on when defining a T.prim_func.