| 54 | def test_lower_view_get(): |
| 55 | @T.prim_func(private=True) |
| 56 | def before1(in_buf: T.Buffer(64, "float32"), out: T.Buffer(64, "float32")) -> None: |
| 57 | T.device_entry() |
| 58 | bx, by, bz = T.cta_id([1, 1, 1]) |
| 59 | T.warp_id([1]) |
| 60 | lane_id = T.lane_id([32]) |
| 61 | A = T.alloc_buffer([2], dtype="float16", scope="local", layout=T.TileLayout(T.S[2:1])) |
| 62 | B_layout = A.layout.tile(L_LANE, (32,), (2,)) |
| 63 | B = A.view(64, layout=B_layout) |
| 64 | A_local = B.local(2) |
| 65 | for i in T.vectorized(2): |
| 66 | A_local[i] = T.float32(in_buf[lane_id * 2 + i]) |
| 67 | B_1 = A.view(64, layout=B_layout) |
| 68 | A_local_1 = B_1.local(2) |
| 69 | for i in T.vectorized(2): |
| 70 | out[lane_id * 2 + i] = T.float32(A_local_1[i]) |
| 71 | |
| 72 | @T.prim_func(private=True) |
| 73 | def after1(in_buf_handle: T.handle, out_handle: T.handle): |