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)
hi @sanirudh ,
Your suggestions helped in analysing the IR and getting to know the bufferLoad and BufferStore nodes ,
however i also want to see these nodes as instruction , so is there a way to have these instructions embeded in the whole schedule itself .?
This would help me have some more clarity when i would replace these load/store inst with custom intrin.
Thanks
Hi @yogeesh, I’m not sure I understood what you meant by instruction. Could you give an example of what you need?
sure @sanirudh ,
basically , lets the program is such :
a = te.placeholder()
b = te.placeholder()
c = a+b
so here i want the schedule to be as :
BufferLoad(a)
BufferLoad(b)
compute(c)
BufferStore(c)
i mean this is a most abstract version , but what i mean to say is that , once i replace these BufferLoad/Store inst with external intrisic , is there a way to display them in the schedule .
Currently , there is no information about the BufferLoad and BufferStore operation as such , so wanted to check if its possible to display these buffer commands also in the schedule ?
I’m not sure I understand what you mean by displaying them in the schedule? The schedule is normally just a sequence of primitives right? After you replace the intrinsics, they’ll certainly show up in the PrimFunc I guess
Okay i understand the ambiguity in the question i asked. Let me inspect the schedule generated and wil get back to you once i have some more clarity . Thank you replying so fast. 