| 58 | |
| 59 | |
| 60 | def test_one_fold_addone(): |
| 61 | # put before after in a single module |
| 62 | @tvm.script.ir_module |
| 63 | class Module: |
| 64 | @T.prim_func(s_tir=True) |
| 65 | def addone(A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32")) -> None: |
| 66 | for i, j in T.grid(16, 16): |
| 67 | with T.sblock("addone"): |
| 68 | vi, vj = T.axis.remap("SS", [i, j]) |
| 69 | B[vi, vj] = A[vi, vj] + T.float32(1) |
| 70 | |
| 71 | @R.function |
| 72 | def before(c0: R.Tensor((16, 16), "float32")): |
| 73 | cls = Module |
| 74 | lv0 = relax.call_tir(cls.addone, (c0,), R.Tensor((16, 16), dtype="float32")) |
| 75 | return lv0 |
| 76 | |
| 77 | @R.function |
| 78 | def expected(c1: R.Tensor((16, 16), "float32")): |
| 79 | return c1 |
| 80 | |
| 81 | c0_np = np.arange(16 * 16).astype("float32").reshape(16, 16) |
| 82 | c1_np = c0_np + 1 |
| 83 | before = gen_mod(Module, "before", {"c0": c0_np}) |
| 84 | expected = gen_mod(Module, "expected", {"c1": c1_np}) |
| 85 | |
| 86 | after = relax.transform.FoldConstant()(before) |
| 87 | tvm.ir.assert_structural_equal(after, expected) |
| 88 | |
| 89 | |
| 90 | def test_one_fold_transpose(): |