| 88 | |
| 89 | |
| 90 | def test_one_fold_transpose(): |
| 91 | # put before after in a single module |
| 92 | @tvm.script.ir_module |
| 93 | class Module: |
| 94 | @T.prim_func(s_tir=True) |
| 95 | def func(A: T.Buffer((2, 3), "float32"), B: T.Buffer((3, 2), "float32")) -> None: |
| 96 | for i, j in T.grid(3, 2): |
| 97 | with T.sblock("transpose"): |
| 98 | vi, vj = T.axis.remap("SS", [i, j]) |
| 99 | B[vi, vj] = A[vj, vi] |
| 100 | |
| 101 | @R.function |
| 102 | def before(c0: R.Tensor((2, 3), "float32")): |
| 103 | cls = Module |
| 104 | lv0 = relax.call_tir(cls.func, (c0,), R.Tensor((3, 2), dtype="float32")) |
| 105 | return lv0 |
| 106 | |
| 107 | @R.function |
| 108 | def expected(c1: R.Tensor((3, 2), "float32")): |
| 109 | return c1 |
| 110 | |
| 111 | c0_np = np.arange(2 * 3).astype("float32").reshape(2, 3) |
| 112 | c1_np = c0_np.T |
| 113 | before = gen_mod(Module, "before", {"c0": c0_np}) |
| 114 | expected = gen_mod(Module, "expected", {"c1": c1_np}) |
| 115 | |
| 116 | after = relax.transform.FoldConstant()(before) |
| 117 | tvm.ir.assert_structural_equal(after, expected) |
| 118 | |
| 119 | |
| 120 | def test_two_hop_addone(): |