Check the same variable?

#1

Say, I have a piece of simple program like this:

import tvm

a = tvm.placeholder((5, ))
b = tvm.compute((5, ), lambda x: a[x])

sch = tvm.create_schedule(b.op)

ir = tvm.lower(sch, [a, b], simple_mode=True)

print(ir)

def visitor(op):
    if isinstance(op, tvm.expr.Load):
        print('Load', op.index, id(op.index))
    if isinstance(op, tvm.stmt.Store):
        print('Store', op.index, id(op.index))
    if isinstance(op, tvm.stmt.For):
        print('For', op.loop_var, id(op.loop_var))


tvm.ir_pass.PostOrderVisit(ir, visitor)

print('Axis', b.op.axis[0].var, id(b.op.axis[0].var))

What I want to do is to check if the for-variable is the same as the axis-variable.
However, if you run this piece of program, you may get:

Load x 140018558610888
Store x 140018558610056
For x 140018558610056
Axis x 140018558611016

Say, “For” and “Store” index are the same, but others are all different.

1 Like