Remove an attribute from IR

Questions
1

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 ?

2

Hi @FrozenGene, Could you please help in sharing your thoughts on this approach?

3

I think it is the correct way to handle it. This is the same as the doc: https://docs.tvm.ai/tutorials/dev/low_level_custom_pass.html

1 Like
4

Thank you @FrozenGene !

We have few schedules implemented in TOPI which have s[c].bind(px, te.thread_axis("pipeline")) .

I want invoke the below mentioned TVM pass if target and target_host are same while performing relay.build operation

tvm.tir.ir_pass.PostOrderVisit(ir, find_pipeline)
stmt = tvm.tir.ir_pass.IRTransform(ir, None, remove_attr, ['AttrStmt'])

If target and target_host are different, then pipeline_exec_scope should be added to the IR.

Could you please suggest a way to incorporate this ?

5

Hi @FrozenGene,

In the TOPI library the target and target_host information is not available and hence the bind pipeline pass cannot be implemented.

Do you think a relay pass should be implemented to take care of it ? If yes, could you please share an example codebase within tvm that i can refer ?

Any help would be greatly appreciated.

6

In the topi, we could get the target information during schedule using tvm.target.Target.current(). But we don’t have target_host information as far as I know.

But seems that you could do in def _build_for_device (we could add pass inside it like other passes). However, you should double check we maybe go to this function if we use relay build (pure C++ way) and doesn’t enter into def _build_for_device:

runtime::Module build(const Map<Target, IRModule>& inputs,
                      const Target& target_host,
                      const BuildConfig& config)

I haven’t verified it.