Hello TVM community,
I am working on a problem where I would like to split the host and device
functionality depending on whether target and target_host are same. If target
is same as target_host, I would like to avoid splitting the IR. If target is
different from target_host, then the device functionality should be different
from host functionality. While doing so, I would like to remove or retain an
AttrStmt for binding the pipeline.
I've created a small test case below to explain the scenario.
n = tvm.tir.const(128, "int32")
a = te.placeholder((n, ), name="a")
b = te.placeholder((n, ), name="b")
c = te.compute((n, ), lambda i: a[i] + b[i], name='c')
s = te.create_schedule(c.op)
px, x = s[c].split(s[c].op.axis[0], nparts=1)
s[c].bind(px, te.thread_axis("pipeline"))
ir = tvm.lower(s, [a, b, c], simple_mode=True)
print(ir)
The ir is printed as :
produce c {
// attr [iter_var(pipeline, , pipeline)] pipeline_exec_scope = 1
for (i.inner, 0, 128) {
c[i.inner] = (a[i.inner] + b[i.inner])
}
}
I would like to remove / attr [iter_var(pipeline, , pipeline)]
pipeline_exec_scope = 1 from the IR. I found a way to accomplish the same as
mentioned in the below snippet
attr = []
def find_pipeline(op):
if isinstance(op, tvm.tir.AttrStmt):
if op.attr_key == "pipeline_exec_scope":
attr.append(op)
def remove_attr(op):
if op in attr:
return op.body
tvm.tir.ir_pass.PostOrderVisit(ir, find_pipeline)
stmt = tvm.tir.ir_pass.IRTransform(ir, None, remove_attr, ['AttrStmt'])
print(stmt)
The stmt is printed as
produce c {
for (i.inner, 0, 128) {
c[i.inner] = (a[i.inner] + b[i.inner])
}
}
I would like to know if this is the correct way to remove the attribute or is
there another better way to do that using some functionality in tvm.tir ?
---
[Visit Topic](https://discuss.tvm.ai/t/remove-an-attribute-from-ir/6526/1) to
respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click
here](https://discuss.tvm.ai/email/unsubscribe/36cb40726b0037f8435582e9a75f445502c6ea8219828ba28ed0529ae9d32f2c).
