def test_schedule_pragmas():
input = te.placeholder((12, 12), dtype="uint8", name="input")
out = te.compute(
(12, 12),
lambda i, j: input[i, j],
attrs={
"op": "unity",
"info": 1,
},
)
sch = te.create_schedule([out.op])
sch[out].split(out.op.axis[0], 3)
schedule_pragmas(sch)
iter_var = sch[out].leaf_iter_vars[1]
assert list(sch[out].iter_var_attrs[iter_var].pragma_keys) == ["op", "info"]
assert list(sch[out].iter_var_attrs[iter_var].pragma_values) == ["unity", 1]
def test_schedule_pragmas_for_const():
input = te.placeholder((12, 12), dtype="uint8", name="input")
const = te.compute((), lambda: 2)
add = topi.add(input, const)
sch = te.create_schedule([add.op])
schedule_pragmas(sch)
