| 29 | |
| 30 | |
| 31 | def get_model(): |
| 32 | pipeline = relax.get_pipeline() |
| 33 | |
| 34 | @tvm.script.ir_module |
| 35 | class Mod: |
| 36 | @R.function |
| 37 | def main(x: R.Tensor([1024], "float32"), y: R.Tensor([1024], "float32")): |
| 38 | lv0 = R.add(x, y) |
| 39 | return lv0 |
| 40 | |
| 41 | mod = pipeline(Mod) |
| 42 | sch = tvm.s_tir.Schedule(mod) |
| 43 | # manually transform loop |
| 44 | sch.work_on("add") |
| 45 | (i,) = sch.get_loops(block=sch.get_sblock("T_add")) |
| 46 | i0, i1 = sch.split(i, [None, 128]) |
| 47 | sch.bind(i0, "blockIdx.x") |
| 48 | sch.bind(i1, "threadIdx.x") |
| 49 | return sch.mod |
| 50 | |
| 51 | |
| 52 | def test_rpc(): |