Hi all,
I’m new to TVM, please be nice. I just want to ask a question. I am trying to implement a profiler for a Hardware accelerator to estimate certain performance metrics like number of cycles for LLM. What I would like to know is whether it’s possible to read, access loops variable or any variable within a computational block during runtime.
for example:
@T.prim_func(private=True)
def matmul(x: T.Buffer((T.int64(1), T.int64(784)), “float32”), lv: T.Buffer((T.int64(784), T.int64(128)), “float32”), T_matmul_NN: T.Buffer((T.int64(1), T.int64(128)), “float32”)):
T.func_attr({“layout_free_buffers”: [1], “tir.noalias”: T.bool(True)})
# with T.block(“root”):
for i0_0, i1_0, i0_1, i1_1, k in T.grid(T.int64(1), T.int64(2), T.int64(64), T.int64(64), T.int64(784)):
with T.block(“T_matmul_NN”):
v_i0 = T.axis.spatial(T.int64(1), i0_0 * T.int64(64) + i0_1)
v_i1 = T.axis.spatial(T.int64(128), i1_0 * T.int64(64) + i1_1)
v_k = T.axis.reduce(T.int64(784), k)
T.where(i0_0 * T.int64(64) + i0_1 < T.int64(1))
T.reads(x[v_i0, v_k], lv[v_k, v_i1])
T.writes(T_matmul_NN[v_i0, v_i1])
with T.init():
T_matmul_NN[v_i0, v_i1] = T.float32(0)
T_matmul_NN[v_i0, v_i1] = T_matmul_NN[v_i0, v_i1] + x[v_i0, v_k] * lv[v_k, v_i1]
are there some tvm functions with which I can add hooks in the IRModule to capture the loop variables during runtime in relax vm? any reference is welcome, thanks