| 40 | # fmt: off |
| 41 | @T.prim_func |
| 42 | def func(A: T.Buffer((128, 512)), B: T.Buffer((128, 512))): |
| 43 | T.func_attr({"num_inputs": 1}) |
| 44 | T.device_entry() |
| 45 | A_sbuf = T.alloc_buffer((128, 512), "float32", scope="trn.sbuf",) |
| 46 | B_sbuf = T.alloc_buffer((128, 512), "float32", scope="trn.sbuf",) |
| 47 | with T.attr(0, "tensorized_nki_instruction", 1): |
| 48 | for i in range(0, 128): |
| 49 | for j in range(0, 512): |
| 50 | T.nki.load(A_sbuf[i, j], A[i, j]) |
| 51 | with T.attr(0, "tensorized_nki_instruction", 1): |
| 52 | for i in range(0, 128): |
| 53 | for j in range(0, 512): |
| 54 | T.nki.tensorscalar(B_sbuf[i, j], A_sbuf[i, j], T.float32(1.0), "add") |
| 55 | with T.attr(0, "tensorized_nki_instruction", 1): |
| 56 | for i in range(0, 128): |
| 57 | for j in range(0, 512): |
| 58 | T.nki.store(B[i, j], B_sbuf[i, j]) |
| 59 | # fmt: on |
| 60 | src = lower_and_get_source(func) |
| 61 | print(src) |
| 62 | expected = """# Function: func_kernel |