hi , I am trying to figure out how to print the Read and write instruction explicitly in TVM schedules.
for ex , if we have :
temp = buffer2[idx2]
can we convert this to :
temp = read(buffer2 , idx2) , something in this format.
kind cc @sanirudh for valuable insights.
If you just want to print, I guess you could use post_order_visit to recursively visit all nodes in the IR. I think reads and writes in TVM are BufferRead and BufferWrite nodes, so you could check for those types and print them in any way you want in its callback that you could provide.
hi @sanirudh
i had a small doubt on this thread regarding read/write nodes in a schedule.
Specifically for my use case i would like to replace the read/write nodes with an intrinsic call to hardware which will handle read/write.
hence wanted to ask about its feasibility whether we can do this in tvm ?
Yes you can easily do that by writing a transformation pass. Either directly in a cpp pass, or if you prefer python, a simple way would be to use the ir_transform function.
Something like this should work:
def modify_read_write(node):
if isinstance(node, tvm.tir.BufferStore):
# Return the write intrinsic
elif isinstance(node, tvm.tir.BufferLoad):
# Return the load intrinsic
# You can pass the modify_read_write function to the second arg for preorder visit or this way for postorder visit
modified_func_body = tvm.tir.stmt_functor.ir_transform(tir_func.body, None, modify_read_write)
new_func = tir_func.with_body(modified_func_body)