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)