已合并
适配 BatchMatmul 与 Matmul 样例至 cannbotdsl 0.7 #127
ykxxlzhong创建于 24 天前
适配 BatchMatmul 与 Matmul 样例至 cannbotdsl 0.7 #127
已合并
共 4 个文件变更+178-122
| @@ -19,7 +19,7 @@ $$ | |||
| 19 | | bias | 可选加性偏置，经 BT（Bias Table）折入 init MMAD（只加一次）；`[N]` 共享（广播到 M 与 batch）或 `[*c_batch, N]` 逐 batch（需与输出 batch 精确匹配，BatchMatMulV3 契约）；dtype 为 fp32 或输入 dtype | | 19 | | bias | 可选加性偏置，经 BT（Bias Table）折入 init MMAD（只加一次）；`[N]` 共享（广播到 M 与 batch）或 `[*c_batch, N]` 逐 batch（需与输出 batch 精确匹配，BatchMatMulV3 契约）；dtype 为 fp32 或输入 dtype | |
| 20 | | 输入 | a/b 需 contiguous 或 canonical 转置视图（其余非连续布局拒绝） | | 20 | | 输入 | a/b 需 contiguous 或 canonical 转置视图（其余非连续布局拒绝） | |
| 21 | | 多核并行 | 展平 batch×M×N tile 空间线性 stride 调度（grid = min(tile 总数, 平台 cube 核数，dav-3510 为 28)，batch 数任意） | | 21 | | 多核并行 | 展平 batch×M×N tile 空间线性 stride 调度（grid = min(tile 总数, 平台 cube 核数，dav-3510 为 28)，batch 数任意） | |
| 22 | -| 尾块处理 | `tile_view` 裁剪 + ND2NZ 引擎零填充，M/N/K 任意值 | | 22 | +| 尾块处理 | `tile_slice` 裁剪 + ND2NZ 引擎零填充，M/N/K 任意值 | |
| 23 | | L2 hint | 逐 copy `l2_cache_ctl` 策略：A/B 默认旁路、按复用分析（跨 tile 重读 / broadcast）使能，含 nd2nz 方向 + 128 元素对齐守卫；C 按容量使能；bias 恒使能 | | 23 | | L2 hint | 逐 copy `l2_cache_ctl` 策略：A/B 默认旁路、按复用分析（跨 tile 重读 / broadcast）使能，含 nd2nz 方向 + 128 元素对齐守卫；C 按容量使能；bias 恒使能 | |
| 24 | | 支持架构 | NPU ARCH 3510（Ascend 950PR / Ascend 950DT） | | 24 | | 支持架构 | NPU ARCH 3510（Ascend 950PR / Ascend 950DT） | |
| 25 | 25 | ||
| @@ -124,4 +124,4 @@ pytest test/matmul/batch_matmul/test_batch_matmul.py -v | |||
| 124 | 124 | ||
| 125 | - 大方形 fp16/fp32 与主线差距 1.1~1.25×（残余为主线策略矩阵：尾块分裂 / StreamK / ND_FIXPIPE AIV 协同回写）；深批量小矩阵反超主线（0.82×） | 125 | - 大方形 fp16/fp32 与主线差距 1.1~1.25×（残余为主线策略矩阵：尾块分裂 / StreamK / ND_FIXPIPE AIV 协同回写）；深批量小矩阵反超主线（0.82×） |
| 126 | - HF32 模式比任何 fp32 路径快 3~7×（torch 亦支持 `torch.npu.matmul.allow_hf32=True`，模式对齐后主线仍略优） | 126 | - HF32 模式比任何 fp32 路径快 3~7×（torch 亦支持 `torch.npu.matmul.allow_hf32=True`，模式对齐后主线仍略优） |
| 127 | -- JIT 便捷入口每次调用含 ~100ms 框架重 trace 开销；低延迟场景请用 `run.compile()` AOT 路径（~156us/调用，同步语义） | 127 | +- JIT 便捷入口每次调用含 ~100ms 框架重 trace 开销；低延迟场景请用 `cannbotdsl.compile(op.run, *specs)` AOT 路径（`TensorSpec` 描述入参，返回可重复调用的编译产物） |
| @@ -21,18 +21,17 @@ import dataclasses | |||
| 21 | import os | 21 | import os |
| 22 | 22 | ||
| 23 | import torch | 23 | import torch |
| 24 | -from cannbotdsl import ( | 24 | +from cannbotdsl import dtypes, get_mem_size, get_platform_info |
| 25 | - MemLoc, Tensor, dtypes, get_mem_size, get_platform_info, | ||
| 26 | - make_copy_engine, mem_copy, tile_slice as tile_view, | ||
| 27 | -) | ||
| 28 | -from cannbotdsl.ops.arch import get_block_idx, get_block_num | ||
| 29 | from cannbotdsl.channel import Channel | 25 | from cannbotdsl.channel import Channel |
| 30 | from cannbotdsl.lang.constexpr import const_expr | 26 | from cannbotdsl.lang.constexpr import const_expr |
| 27 | +from cannbotdsl.lang.host import host | ||
| 31 | from cannbotdsl.lang.jit import jit | 28 | from cannbotdsl.lang.jit import jit |
| 32 | from cannbotdsl.lang.kernel import kernel | 29 | from cannbotdsl.lang.kernel import kernel |
| 30 | +from cannbotdsl.ops.arch import get_block_idx, get_block_num | ||
| 33 | from cannbotdsl.ops.cube import enable_hf32, set_fp32_mode | 31 | from cannbotdsl.ops.cube import enable_hf32, set_fp32_mode |
| 34 | from cannbotdsl.ops.matmul import matmul as dsl_matmul | 32 | from cannbotdsl.ops.matmul import matmul as dsl_matmul |
| 35 | -from cannbotdsl.tensor import ceil_div | 33 | +from cannbotdsl.ops.memcpy import make_copy_engine, mem_copy |
| 34 | +from cannbotdsl.tensor import MemLoc, Tensor, ceil_div, tile_slice | ||
| 36 | 35 | ||
| 37 | 36 | ||
| 38 | # ============================================================================ | 37 | # ============================================================================ |
| @@ -114,10 +113,12 @@ def _decode_batch_terms(src_batch, c_batch): | |||
| 114 | ps = (1,) * (n - len(src_batch)) + tuple(int(d) for d in src_batch) | 113 | ps = (1,) * (n - len(src_batch)) + tuple(int(d) for d in src_batch) |
| 115 | 114 | ||
| 116 | divs = [1] * n | 115 | divs = [1] * n |
| 116 | + muls = [1] * (n + 1) | ||
| 117 | acc = 1 | 117 | acc = 1 |
| 118 | for d in range(n - 1, -1, -1): | 118 | for d in range(n - 1, -1, -1): |
| 119 | divs[d] = acc | 119 | divs[d] = acc |
| 120 | acc *= c_batch[d] | 120 | acc *= c_batch[d] |
| 121 | + muls[d] = muls[d + 1] * ps[d] | ||
| 121 | 122 | ||
| 122 | full = ps == tuple(c_batch) | 123 | full = ps == tuple(c_batch) |
| 123 | terms = [] | 124 | terms = [] |
| @@ -126,7 +127,7 @@ def _decode_batch_terms(src_batch, c_batch): | |||
| 126 | continue # broadcast dim: source coordinate is always 0 | 127 | continue # broadcast dim: source coordinate is always 0 |
| 127 | if c_batch[d] == 1: | 128 | if c_batch[d] == 1: |
| 128 | continue # extent-1 output dim: coordinate is always 0 | 129 | continue # extent-1 output dim: coordinate is always 0 |
| 129 | - terms.append((divs[d], c_batch[d], _prod(ps[d + 1:]))) | 130 | + terms.append((divs[d], c_batch[d], muls[d + 1])) |
| 130 | return full, terms | 131 | return full, terms |
| 131 | 132 | ||
| 132 | 133 | ||
| @@ -180,9 +181,7 @@ def _l2_capacity_bytes() -> int: | |||
| 180 | f"BMM_L2_SIZE_BYTES must be an integer byte count; got {raw!r}" | 181 | f"BMM_L2_SIZE_BYTES must be an integer byte count; got {raw!r}" |
| 181 | ) from exc | 182 | ) from exc |
| 182 | if value <= 0: | 183 | if value <= 0: |
| 183 | - raise ValueError( | 184 | + raise ValueError(f"BMM_L2_SIZE_BYTES must be positive; got {value}") |
| 184 | - f"BMM_L2_SIZE_BYTES must be positive; got {value}" | ||
| 185 | - ) | ||
| 186 | return value | 185 | return value |
| 187 | 186 | ||
| 188 | 187 | ||
| @@ -191,18 +190,14 @@ def _l2_capacity_bytes() -> int: | |||
| 191 | 190 | ||
| 192 | def _validate_inputs(a, b, hf32): | 191 | def _validate_inputs(a, b, hf32): |
| 193 | if a.dtype != b.dtype: | 192 | if a.dtype != b.dtype: |
| 194 | - raise TypeError( | 193 | + raise TypeError(f"a and b must have the same dtype, got {a.dtype} vs {b.dtype}") |
| 195 | - f"a and b must have the same dtype, got {a.dtype} vs {b.dtype}" | ||
| 196 | - ) | ||
| 197 | if a.dtype not in _TORCH_DTYPE_TO_DSL: | 194 | if a.dtype not in _TORCH_DTYPE_TO_DSL: |
| 198 | raise TypeError( | 195 | raise TypeError( |
| 199 | f"only torch.float16, torch.bfloat16 and torch.float32 are " | 196 | f"only torch.float16, torch.bfloat16 and torch.float32 are " |
| 200 | f"supported; got {a.dtype}" | 197 | f"supported; got {a.dtype}" |
| 201 | ) | 198 | ) |
| 202 | if hf32 and a.dtype != torch.float32: | 199 | if hf32 and a.dtype != torch.float32: |
| 203 | - raise ValueError( | 200 | + raise ValueError(f"hf32=True is only valid for float32 inputs; got {a.dtype}") |
| 204 | - f"hf32=True is only valid for float32 inputs; got {a.dtype}" | ||
| 205 | - ) | ||
| 206 | if a.device != b.device: | 201 | if a.device != b.device: |
| 207 | raise ValueError( | 202 | raise ValueError( |
| 208 | f"a and b must be on the same device, got {a.device} vs {b.device}" | 203 | f"a and b must be on the same device, got {a.device} vs {b.device}" |
| @@ -211,8 +206,7 @@ def _validate_inputs(a, b, hf32): | |||
| 211 | raise ValueError(f"inputs must be NPU tensors; got device {a.device}") | 206 | raise ValueError(f"inputs must be NPU tensors; got device {a.device}") |
| 212 | if not (2 <= a.dim() <= _MAX_RANK and 2 <= b.dim() <= _MAX_RANK): | 207 | if not (2 <= a.dim() <= _MAX_RANK and 2 <= b.dim() <= _MAX_RANK): |
| 213 | raise ValueError( | 208 | raise ValueError( |
| 214 | - f"a and b must have rank 2..{_MAX_RANK}, " | 209 | + f"a and b must have rank 2..{_MAX_RANK}, got rank {a.dim()} and {b.dim()}" |
| 215 | - f"got rank {a.dim()} and {b.dim()}" | ||
| 216 | ) | 210 | ) |
| 217 | 211 | ||
| 218 | 212 | ||
| @@ -256,24 +250,19 @@ def _normalize_bias(bias, a, c_batch, n): | |||
| 256 | if bias is None: | 250 | if bias is None: |
| 257 | return None, False, False | 251 | return None, False, False |
| 258 | if bias.device != a.device: | 252 | if bias.device != a.device: |
| 259 | - raise ValueError( | 253 | + raise ValueError(f"bias must be on the same device as a; got {bias.device}") |
| 260 | - f"bias must be on the same device as a; got {bias.device}" | ||
| 261 | - ) | ||
| 262 | if bias.dtype not in (a.dtype, torch.float32): | 254 | if bias.dtype not in (a.dtype, torch.float32): |
| 263 | raise TypeError( | 255 | raise TypeError( |
| 264 | f"bias dtype must be the input dtype or float32; got {bias.dtype}" | 256 | f"bias dtype must be the input dtype or float32; got {bias.dtype}" |
| 265 | ) | 257 | ) |
| 266 | if bias.dim() == 1: | 258 | if bias.dim() == 1: |
| 267 | if bias.shape[0] != n: | 259 | if bias.shape[0] != n: |
| 268 | - raise ValueError( | 260 | + raise ValueError(f"bias length must equal N ({n}); got {bias.shape[0]}") |
| 269 | - f"bias length must equal N ({n}); got {bias.shape[0]}" | ||
| 270 | - ) | ||
| 271 | return bias, True, False | 261 | return bias, True, False |
| 272 | if bias.dim() == len(c_batch) + 1: | 262 | if bias.dim() == len(c_batch) + 1: |
| 273 | if tuple(bias.shape[:-1]) != c_batch or bias.shape[-1] != n: | 263 | if tuple(bias.shape[:-1]) != c_batch or bias.shape[-1] != n: |
| 274 | raise ValueError( | 264 | raise ValueError( |
| 275 | - f"per-batch bias shape must be {(*c_batch, n)}; " | 265 | + f"per-batch bias shape must be {(*c_batch, n)}; got {tuple(bias.shape)}" |
| 276 | - f"got {tuple(bias.shape)}" | ||
| 277 | ) | 266 | ) |
| 278 | return bias, False, True | 267 | return bias, False, True |
| 279 | raise ValueError( | 268 | raise ValueError( |
| @@ -311,8 +300,11 @@ def _degenerate_result(ctx: _KernelCtx): | |||
| 311 | 300 | ||
| 312 | def zeros_out(): | 301 | def zeros_out(): |
| 313 | z = torch.zeros( | 302 | z = torch.zeros( |
| 314 | - *ctx.c_batch, ctx.m, ctx.n, | 303 | + *ctx.c_batch, |
| 315 | - dtype=ctx.a_kern.dtype, device=ctx.a_kern.device, | 304 | + ctx.m, |
| 305 | + ctx.n, | ||
| 306 | + dtype=ctx.a_kern.dtype, | ||
| 307 | + device=ctx.a_kern.device, | ||
| 316 | ) | 308 | ) |
| 317 | return z[0] if ctx.squeeze_out else z | 309 | return z[0] if ctx.squeeze_out else z |
| 318 | 310 | ||
| @@ -351,8 +343,10 @@ def _pad_n_tail(ctx: _KernelCtx): | |||
| 351 | b_kern = torch.nn.functional.pad(b_kern, (0, n_eff - n)) | 343 | b_kern = torch.nn.functional.pad(b_kern, (0, n_eff - n)) |
| 352 | if bias is not None: | 344 | if bias is not None: |
| 353 | bias_pad = torch.zeros( | 345 | bias_pad = torch.zeros( |
| 354 | - *(() if ctx.bias_shared else ctx.c_batch), n_eff, | 346 | + *(() if ctx.bias_shared else ctx.c_batch), |
| 355 | - dtype=bias.dtype, device=bias.device, | 347 | + n_eff, |
| 348 | + dtype=bias.dtype, | ||
| 349 | + device=bias.device, | ||
| 356 | ) | 350 | ) |
| 357 | bias_pad[..., :n] = bias | 351 | bias_pad[..., :n] = bias |
| 358 | bias = bias_pad | 352 | bias = bias_pad |
| @@ -362,7 +356,7 @@ def _pad_n_tail(ctx: _KernelCtx): | |||
| 362 | def _reshape_output(c, c_n, ctx: _KernelCtx): | 356 | def _reshape_output(c, c_n, ctx: _KernelCtx): |
| 363 | """Slice the N pad back off (if any) and restore the batch dims.""" | 357 | """Slice the N pad back off (if any) and restore the batch dims.""" |
| 364 | if c_n != ctx.n: | 358 | if c_n != ctx.n: |
| 365 | - c = c[..., :ctx.n] | 359 | + c = c.narrow(-1, 0, ctx.n) |
| 366 | if ctx.squeeze_out: | 360 | if ctx.squeeze_out: |
| 367 | return c[0] | 361 | return c[0] |
| 368 | return c.reshape(*ctx.c_batch, ctx.m, ctx.n) | 362 | return c.reshape(*ctx.c_batch, ctx.m, ctx.n) |
| @@ -376,8 +370,12 @@ def _run_kernel(ctx: _KernelCtx): | |||
| 376 | has_bias = ctx.bias is not None | 370 | has_bias = ctx.bias is not None |
| 377 | b_kern, bias, n_eff = _pad_n_tail(ctx) | 371 | b_kern, bias, n_eff = _pad_n_tail(ctx) |
| 378 | tiling = BatchMatmulTiling( | 372 | tiling = BatchMatmulTiling( |
| 379 | - ctx.m, n_eff, ctx.k, dtype, | 373 | + ctx.m, |
| 380 | - kern_trans_b=ctx.kern_trans_b, has_bias=has_bias, | 374 | + n_eff, |
| 375 | + ctx.k, | ||
| 376 | + dtype, | ||
| 377 | + kern_trans_b=ctx.kern_trans_b, | ||
| 378 | + has_bias=has_bias, | ||
| 381 | ) | 379 | ) |
| 382 | # Zero-copy batch flatten (contiguous views, .view never copies). | 380 | # Zero-copy batch flatten (contiguous views, .view never copies). |
| 383 | a3 = ctx.a_kern.view(_prod(ctx.a_kern.shape[:-2]), *ctx.a_kern.shape[-2:]) | 381 | a3 = ctx.a_kern.view(_prod(ctx.a_kern.shape[:-2]), *ctx.a_kern.shape[-2:]) |
| @@ -477,7 +475,9 @@ class BatchMatmulTiling: | |||
| 477 | # A/B share the dtype; bound base_k by the smaller of L0A/L0B | 475 | # A/B share the dtype; bound base_k by the smaller of L0A/L0B |
| 478 | # (double-buffered), both queried from the platform. | 476 | # (double-buffered), both queried from the platform. |
| 479 | max_k = ( | 477 | max_k = ( |
| 480 | - min(self.L0A_SIZE, self.L0B_SIZE) // self.DB_SIZE // dtype_size | 478 | + min(self.L0A_SIZE, self.L0B_SIZE) |
| 479 | + // self.DB_SIZE | ||
| 480 | + // dtype_size | ||
| 481 | // max(self.base_m, self.base_n) | 481 | // max(self.base_m, self.base_n) |
| 482 | ) | 482 | ) |
| 483 | k_align = _ceil_align(self.k, self.BASIC_BLOCK_16) | 483 | k_align = _ceil_align(self.k, self.BASIC_BLOCK_16) |
| @@ -491,8 +491,7 @@ class BatchMatmulTiling: | |||
| 491 | 491 | ||
| 492 | max_step = min(_ceil_div(self.k, self.base_k), self.MAX_STEP_K) | 492 | max_step = min(_ceil_div(self.k, self.base_k), self.MAX_STEP_K) |
| 493 | bias_l1 = ( | 493 | bias_l1 = ( |
| 494 | - self.base_n * self.DATA_SIZE_FP32 * self.DB_SIZE | 494 | + self.base_n * self.DATA_SIZE_FP32 * self.DB_SIZE if self.has_bias else 0 |
| 495 | - if self.has_bias else 0 | ||
| 496 | ) | 495 | ) |
| 497 | step_k = 1 | 496 | step_k = 1 |
| 498 | for s in range(1, max_step + 1): | 497 | for s in range(1, max_step + 1): |
| @@ -569,9 +568,8 @@ class BatchMatmulKernel: | |||
| 569 | # L2 cache-hint policy: A/B bypass only on proven reuse and only on | 568 | # L2 cache-hint policy: A/B bypass only on proven reuse and only on |
| 570 | # the linear nd2nz path with 128-aligned K/k_l1; C cached iff the | 569 | # the linear nd2nz path with 128-aligned K/k_l1; C cached iff the |
| 571 | # whole batched output fits L2; bias always cached. | 570 | # whole batched output fits L2; bias always cached. |
| 572 | - aligned = ( | 571 | + aligned = (tiling.k % _L2_BYPASS_ALIGN_ELEMS == 0) and ( |
| 573 | - (tiling.k % _L2_BYPASS_ALIGN_ELEMS == 0) | 572 | + tiling.k_l1 % _L2_BYPASS_ALIGN_ELEMS == 0 |
| 574 | - and (tiling.k_l1 % _L2_BYPASS_ALIGN_ELEMS == 0) | ||
| 575 | ) | 573 | ) |
| 576 | reuse_a = tiling.n_tiles > 1 or not ctx.a_full # N-tile re-reads / bcast | 574 | reuse_a = tiling.n_tiles > 1 or not ctx.a_full # N-tile re-reads / bcast |
| 577 | reuse_b = tiling.m_tiles > 1 or not ctx.b_full # M-tile re-reads / bcast | 575 | reuse_b = tiling.m_tiles > 1 or not ctx.b_full # M-tile re-reads / bcast |
| @@ -580,8 +578,9 @@ class BatchMatmulKernel: | |||
| 580 | self.l2_ctl_a = 1 if (reuse_a or not nd_path_a or not aligned) else 0 | 578 | self.l2_ctl_a = 1 if (reuse_a or not nd_path_a or not aligned) else 0 |
| 581 | self.l2_ctl_b = 1 if (reuse_b or not nd_path_b or not aligned) else 0 | 579 | self.l2_ctl_b = 1 if (reuse_b or not nd_path_b or not aligned) else 0 |
| 582 | self.l2_ctl_c = ( | 580 | self.l2_ctl_c = ( |
| 583 | - 1 if self.bs * tiling.m * tiling.n * tiling.dtype_size | 581 | + 1 |
| 584 | - <= _l2_capacity_bytes() else 0 | 582 | + if self.bs * tiling.m * tiling.n * tiling.dtype_size <= _l2_capacity_bytes() |
| 583 | + else 0 | ||
| 585 | ) | 584 | ) |
| 586 | 585 | ||
| 587 | # MMAD unit flag (batch_mat_mul_v3 Normal config): 2 while | 586 | # MMAD unit flag (batch_mat_mul_v3 Normal config): 2 while |
| @@ -611,8 +610,7 @@ class BatchMatmulKernel: | |||
| 611 | return tuple(padded) | 610 | return tuple(padded) |
| 612 | 611 | ||
| 613 | @kernel | 612 | @kernel |
| 614 | - def bmm_kernel(self, gm_a: Tensor, gm_b: Tensor, gm_c: Tensor, | 613 | + def bmm_kernel(self, gm_a: Tensor, gm_b: Tensor, gm_c: Tensor, gm_bias: Tensor): |
| 615 | - gm_bias: Tensor): | ||
| 616 | t = self.t | 614 | t = self.t |
| 617 | 615 | ||
| 618 | # Compile-time fp32 compute-mode selection (persistent cube state). | 616 | # Compile-time fp32 compute-mode selection (persistent cube state). |
| @@ -632,12 +630,8 @@ class BatchMatmulKernel: | |||
| 632 | """Declare the L1/L0/bias channels and the copy engines.""" | 630 | """Declare the L1/L0/bias channels and the copy engines.""" |
| 633 | t = self.t | 631 | t = self.t |
| 634 | 632 | ||
| 635 | - l1_b_shape = ( | 633 | + l1_b_shape = (t.base_n, t.k_l1) if t.kern_trans_b else (t.k_l1, t.base_n) |
| 636 | - (t.base_n, t.k_l1) if t.kern_trans_b else (t.k_l1, t.base_n) | 634 | + l1_a_shape = (t.k_l1, t.base_m) if self.kern_trans_a else (t.base_m, t.k_l1) |
| 637 | - ) | ||
| 638 | - l1_a_shape = ( | ||
| 639 | - (t.k_l1, t.base_m) if self.kern_trans_a else (t.base_m, t.k_l1) | ||
| 640 | - ) | ||
| 641 | self.l1_a = Channel( | 635 | self.l1_a = Channel( |
| 642 | MemLoc.L1, | 636 | MemLoc.L1, |
| 643 | shape=l1_a_shape, | 637 | shape=l1_a_shape, |
| @@ -664,12 +658,18 @@ class BatchMatmulKernel: | |||
| 664 | # bias-only so the channel never eats into exact-fit no-bias L1. | 658 | # bias-only so the channel never eats into exact-fit no-bias L1. |
| 665 | if const_expr(self.has_bias): | 659 | if const_expr(self.has_bias): |
| 666 | self.l1_bias = Channel( | 660 | self.l1_bias = Channel( |
| 667 | - MemLoc.L1, shape=(t.base_n,), dtype=self.bias_dtype, | 661 | + MemLoc.L1, |
| 668 | - depth=2, data_format="nd", | 662 | + shape=(t.base_n,), |
| 663 | + dtype=self.bias_dtype, | ||
| 664 | + depth=2, | ||
| 665 | + data_format="nd", | ||
| 669 | ) | 666 | ) |
| 670 | self.l1_bt = Channel( | 667 | self.l1_bt = Channel( |
| 671 | - MemLoc.BIAS, shape=(t.base_n,), dtype=dtypes.float32, | 668 | + MemLoc.BIAS, |
| 672 | - depth=2, data_format="nd", | 669 | + shape=(t.base_n,), |
| 670 | + dtype=dtypes.float32, | ||
| 671 | + depth=2, | ||
| 672 | + data_format="nd", | ||
| 673 | ) | 673 | ) |
| 674 | 674 | ||
| 675 | self.eng_a = make_copy_engine( | 675 | self.eng_a = make_copy_engine( |
| @@ -681,17 +681,15 @@ class BatchMatmulKernel: | |||
| 681 | self.eng_bias = make_copy_engine(format_transform="identity") | 681 | self.eng_bias = make_copy_engine(format_transform="identity") |
| 682 | 682 | ||
| 683 | @jit | 683 | @jit |
| 684 | - def _compute_output_tile(self, gm_a, gm_b, gm_c, gm_bias, tile_idx): | 684 | + def _decode_tile_batch_src(self, tile_idx): |
| 685 | - """Bias load + K reduction + FIXPIPE for one output tile; tail tiles | 685 | + """Tile-index decode: (b_idx, mn_idx) plus the broadcast-resolved flat |
| 686 | - are clipped by tile_view and zero-padded by the nd2nz engines. | 686 | + source indices (see _decode_batch_terms; full ⇒ source index equals |
| 687 | - """ | 687 | + the output batch index). mn_idx packs (m_idx, n_idx) into one scalar |
| 688 | + for the 5-value coding rules (codecheck G.FNM.03/05). Shared [N] | ||
| 689 | + bias lives in slot 0; full-batch bias by output batch.""" | ||
| 688 | b_idx = tile_idx // self.mn_tiles | 690 | b_idx = tile_idx // self.mn_tiles |
| 689 | - rem = tile_idx % self.mn_tiles | 691 | + mn_idx = tile_idx % self.mn_tiles |
| 690 | - m_idx = rem // self.n_tiles | ||
| 691 | - n_idx = rem % self.n_tiles | ||
| 692 | 692 | ||
| 693 | - # Broadcast-resolved flat batch indices (see _decode_batch_terms); | ||
| 694 | - # full ⇒ source index equals the output batch index. | ||
| 695 | a_src = b_idx | 693 | a_src = b_idx |
| 696 | if const_expr(not self.a_full): | 694 | if const_expr(not self.a_full): |
| 697 | a_src = ( | 695 | a_src = ( |
| @@ -708,10 +706,20 @@ class BatchMatmulKernel: | |||
| 708 | + ((b_idx // self.b2[0]) % self.b2[1]) * self.b2[2] | 706 | + ((b_idx // self.b2[0]) % self.b2[1]) * self.b2[2] |
| 709 | + ((b_idx // self.b3[0]) % self.b3[1]) * self.b3[2] | 707 | + ((b_idx // self.b3[0]) % self.b3[1]) * self.b3[2] |
| 710 | ) | 708 | ) |
| 711 | - # Shared [N] bias lives in slot 0; full-batch bias by output batch. | ||
| 712 | bias_src = 0 | 709 | bias_src = 0 |
| 713 | if const_expr(self.bias_full): | 710 | if const_expr(self.bias_full): |
| 714 | bias_src = b_idx | 711 | bias_src = b_idx |
| 712 | + return b_idx, mn_idx, a_src, b_src, bias_src | ||
| 713 | + | ||
| 714 | + @jit | ||
| 715 | + def _compute_output_tile(self, gm_a, gm_b, gm_c, gm_bias, tile_idx): | ||
| 716 | + """Bias load + K reduction + FIXPIPE for one output tile; tail tiles | ||
| 717 | + are clipped by tile_slice and zero-padded by the nd2nz engines. | ||
| 718 | + """ | ||
| 719 | + decoded = self._decode_tile_batch_src(tile_idx) | ||
| 720 | + b_idx, mn_idx, a_src, b_src, bias_src = decoded | ||
| 721 | + m_idx = mn_idx // self.n_tiles | ||
| 722 | + n_idx = mn_idx % self.n_tiles | ||
| 715 | 723 | ||
| 716 | base_m, base_n = self.t.base_m, self.t.base_n | 724 | base_m, base_n = self.t.base_m, self.t.base_n |
| 717 | a2d = gm_a[a_src, None, None] | 725 | a2d = gm_a[a_src, None, None] |
| @@ -726,8 +734,9 @@ class BatchMatmulKernel: | |||
| 726 | if const_expr(self.has_bias): | 734 | if const_expr(self.has_bias): |
| 727 | mem_copy( | 735 | mem_copy( |
| 728 | self.l1_bias.produce(), | 736 | self.l1_bias.produce(), |
| 729 | - tile_view(gm_bias[bias_src, None], (base_n,), (n_idx,)), | 737 | + tile_slice(gm_bias[bias_src, None], (base_n,), (n_idx,)), |
| 730 | - engine=self.eng_bias, l2_cache_ctl=1, | 738 | + engine=self.eng_bias, |
| 739 | + l2_cache_ctl=1, | ||
| 731 | ) | 740 | ) |
| 732 | mem_copy(self.l1_bt.produce(), self.l1_bias.consume()) | 741 | mem_copy(self.l1_bt.produce(), self.l1_bias.consume()) |
| 733 | bias = self.l1_bt.consume() | 742 | bias = self.l1_bt.consume() |
| @@ -737,45 +746,68 @@ class BatchMatmulKernel: | |||
| 737 | # mn_idx packs (m_idx, n_idx) into one scalar purely to keep | 746 | # mn_idx packs (m_idx, n_idx) into one scalar purely to keep |
| 738 | # _k_reduction within the repo's 5-formal-parameter coding rule | 747 | # _k_reduction within the repo's 5-formal-parameter coding rule |
| 739 | # (codecheck G.FNM.03); it is a rule workaround, not business design. | 748 | # (codecheck G.FNM.03); it is a rule workaround, not business design. |
| 740 | - self._k_reduction(a2d, b2d, c_w, bias, m_idx * self.n_tiles + n_idx) | 749 | + self._k_reduction(a2d, b2d, c_w, bias, mn_idx) |
| 741 | 750 | ||
| 742 | - gm_c_tile = tile_view(c2d, (base_m, base_n), (m_idx, n_idx)) | 751 | + gm_c_tile = tile_slice(c2d, (base_m, base_n), (m_idx, n_idx)) |
| 743 | if const_expr(self.use_unit_flag): | 752 | if const_expr(self.use_unit_flag): |
| 744 | # Unit-flag path: the FIXPIPE reads the produce alias; a | 753 | # Unit-flag path: the FIXPIPE reads the produce alias; a |
| 745 | # consume() read would desync against the auto-triggered FIXPIPE. | 754 | # consume() read would desync against the auto-triggered FIXPIPE. |
| 746 | - mem_copy(gm_c_tile, c_w, engine=self.eng_c, unit_flag=3, | 755 | + mem_copy( |
| 747 | - l2_cache_ctl=self.l2_ctl_c) | 756 | + gm_c_tile, |
| 757 | + c_w, | ||
| 758 | + engine=self.eng_c, | ||
| 759 | + unit_flag=3, | ||
| 760 | + l2_cache_ctl=self.l2_ctl_c, | ||
| 761 | + ) | ||
| 748 | else: | 762 | else: |
| 749 | mem_copy(gm_c_tile, self.l0c.consume(), l2_cache_ctl=self.l2_ctl_c) | 763 | mem_copy(gm_c_tile, self.l0c.consume(), l2_cache_ctl=self.l2_ctl_c) |
| 750 | 764 | ||
| 765 | + @jit | ||
| 766 | + def _load_ab_l1(self, a2d, b2d, k_l1_idx, mn_idx): | ||
| 767 | + """GM -> L1 copies of one k_l1 segment for both operands (mn_idx packs | ||
| 768 | + (m_idx, n_idx) for the 5-arg coding rule).""" | ||
| 769 | + m_idx = mn_idx // self.n_tiles | ||
| 770 | + n_idx = mn_idx % self.n_tiles | ||
| 771 | + t = self.t | ||
| 772 | + base_m, base_n = t.base_m, t.base_n | ||
| 773 | + k_l1 = t.k_l1 | ||
| 774 | + | ||
| 775 | + if const_expr(self.kern_trans_a): | ||
| 776 | + gm_a_tile = tile_slice(a2d, (k_l1, base_m), (k_l1_idx, m_idx)) | ||
| 777 | + else: | ||
| 778 | + gm_a_tile = tile_slice(a2d, (base_m, k_l1), (m_idx, k_l1_idx)) | ||
| 779 | + mem_copy( | ||
| 780 | + self.l1_a.produce(), | ||
| 781 | + gm_a_tile, | ||
| 782 | + engine=self.eng_a, | ||
| 783 | + l2_cache_ctl=self.l2_ctl_a, | ||
| 784 | + ) | ||
| 785 | + | ||
| 786 | + if const_expr(t.kern_trans_b): | ||
| 787 | + gm_b_tile = tile_slice(b2d, (base_n, k_l1), (n_idx, k_l1_idx)) | ||
| 788 | + else: | ||
| 789 | + gm_b_tile = tile_slice(b2d, (k_l1, base_n), (k_l1_idx, n_idx)) | ||
| 790 | + mem_copy( | ||
| 791 | + self.l1_b.produce(), | ||
| 792 | + gm_b_tile, | ||
| 793 | + engine=self.eng_b, | ||
| 794 | + l2_cache_ctl=self.l2_ctl_b, | ||
| 795 | + ) | ||
| 796 | + | ||
| 751 | @jit | 797 | @jit |
| 752 | def _k_reduction(self, a2d, b2d, c_w, bias, mn_idx): | 798 | def _k_reduction(self, a2d, b2d, c_w, bias, mn_idx): |
| 753 | """GM → L1 → L0 K reduction and MMAD accumulation for one tile | 799 | """GM → L1 → L0 K reduction and MMAD accumulation for one tile |
| 754 | (mn_idx packs (m_idx, n_idx) for the 5-arg coding rule).""" | 800 | (mn_idx packs (m_idx, n_idx) for the 5-arg coding rule).""" |
| 755 | - m_idx = mn_idx // self.n_tiles | ||
| 756 | - n_idx = mn_idx % self.n_tiles | ||
| 757 | t = self.t | 801 | t = self.t |
| 758 | base_m, base_n = t.base_m, t.base_n | 802 | base_m, base_n = t.base_m, t.base_n |
| 759 | base_k, k_l1, k_total = t.base_k, t.k_l1, t.k | 803 | base_k, k_l1, k_total = t.base_k, t.k_l1, t.k |
| 760 | last_l1_idx = self.k_l1_tiles - 1 | 804 | last_l1_idx = self.k_l1_tiles - 1 |
| 761 | 805 | ||
| 762 | for k_l1_idx in range(self.k_l1_tiles): | 806 | for k_l1_idx in range(self.k_l1_tiles): |
| 763 | - if const_expr(self.kern_trans_a): | 807 | + self._load_ab_l1(a2d, b2d, k_l1_idx, mn_idx) |
| 764 | - gm_a_tile = tile_view(a2d, (k_l1, base_m), (k_l1_idx, m_idx)) | ||
| 765 | - else: | ||
| 766 | - gm_a_tile = tile_view(a2d, (base_m, k_l1), (m_idx, k_l1_idx)) | ||
| 767 | # One consume per k_l1 segment: the inner loop slices the selected | 808 | # One consume per k_l1 segment: the inner loop slices the selected |
| 768 | # alias (a consume per slice would rotate the cursor off the slot). | 809 | # alias (a consume per slice would rotate the cursor off the slot). |
| 769 | - mem_copy(self.l1_a.produce(), gm_a_tile, engine=self.eng_a, | ||
| 770 | - l2_cache_ctl=self.l2_ctl_a) | ||
| 771 | a_l1 = self.l1_a.consume() | 810 | a_l1 = self.l1_a.consume() |
| 772 | - | ||
| 773 | - if const_expr(t.kern_trans_b): | ||
| 774 | - gm_b_tile = tile_view(b2d, (base_n, k_l1), (n_idx, k_l1_idx)) | ||
| 775 | - else: | ||
| 776 | - gm_b_tile = tile_view(b2d, (k_l1, base_n), (k_l1_idx, n_idx)) | ||
| 777 | - mem_copy(self.l1_b.produce(), gm_b_tile, engine=self.eng_b, | ||
| 778 | - l2_cache_ctl=self.l2_ctl_b) | ||
| 779 | b_l1 = self.l1_b.consume() | 811 | b_l1 = self.l1_b.consume() |
| 780 | 812 | ||
| 781 | # The final k_l1 segment carries only its valid base_k blocks. | 813 | # The final k_l1 segment carries only its valid base_k blocks. |
| @@ -783,31 +815,38 @@ class BatchMatmulKernel: | |||
| 783 | k_l0_per_l1 = ceil_div(min(k_l1, k_remaining), base_k) | 815 | k_l0_per_l1 = ceil_div(min(k_l1, k_remaining), base_k) |
| 784 | for k_l0_idx in range(k_l0_per_l1): | 816 | for k_l0_idx in range(k_l0_per_l1): |
| 785 | if const_expr(self.kern_trans_a): | 817 | if const_expr(self.kern_trans_a): |
| 786 | - l1_a_slice = tile_view(a_l1, (base_k, base_m), (k_l0_idx, 0)) | 818 | + l1_a_slice = tile_slice(a_l1, (base_k, base_m), (k_l0_idx, 0)) |
| 787 | else: | 819 | else: |
| 788 | - l1_a_slice = tile_view(a_l1, (base_m, base_k), (0, k_l0_idx)) | 820 | + l1_a_slice = tile_slice(a_l1, (base_m, base_k), (0, k_l0_idx)) |
| 789 | a_l0 = self.l0a.produce() | 821 | a_l0 = self.l0a.produce() |
| 790 | mem_copy(a_l0, l1_a_slice) | 822 | mem_copy(a_l0, l1_a_slice) |
| 791 | 823 | ||
| 792 | if const_expr(t.kern_trans_b): | 824 | if const_expr(t.kern_trans_b): |
| 793 | - l1_b_slice = tile_view(b_l1, (base_n, base_k), (0, k_l0_idx)) | 825 | + l1_b_slice = tile_slice(b_l1, (base_n, base_k), (0, k_l0_idx)) |
| 794 | b_l0 = self.l0b.produce() | 826 | b_l0 = self.l0b.produce() |
| 795 | mem_copy(b_l0, l1_b_slice, transpose=False) | 827 | mem_copy(b_l0, l1_b_slice, transpose=False) |
| 796 | else: | 828 | else: |
| 797 | - l1_b_slice = tile_view(b_l1, (base_k, base_n), (k_l0_idx, 0)) | 829 | + l1_b_slice = tile_slice(b_l1, (base_k, base_n), (k_l0_idx, 0)) |
| 798 | b_l0 = self.l0b.produce() | 830 | b_l0 = self.l0b.produce() |
| 799 | mem_copy(b_l0, l1_b_slice, transpose=True) | 831 | mem_copy(b_l0, l1_b_slice, transpose=True) |
| 800 | 832 | ||
| 801 | - is_first_k_block = (k_l1_idx == 0 and k_l0_idx == 0) | 833 | + is_first_k_block = k_l1_idx == 0 and k_l0_idx == 0 |
| 802 | unit_flag = 0 | 834 | unit_flag = 0 |
| 803 | if const_expr(self.use_unit_flag): | 835 | if const_expr(self.use_unit_flag): |
| 804 | - is_last_k_block = (k_l1_idx == last_l1_idx | 836 | + is_last_k_block = ( |
| 805 | - and k_l0_idx == k_l0_per_l1 - 1) | 837 | + k_l1_idx == last_l1_idx and k_l0_idx == k_l0_per_l1 - 1 |
| 838 | + ) | ||
| 806 | unit_flag = 3 if is_last_k_block else 2 | 839 | unit_flag = 3 if is_last_k_block else 2 |
| 807 | - dsl_matmul(c_w, a_l0, b_l0, init=is_first_k_block, | 840 | + dsl_matmul( |
| 808 | - bias=bias, unit_flag=unit_flag) | 841 | + c_w, |
| 842 | + a_l0, | ||
| 843 | + b_l0, | ||
| 844 | + init=is_first_k_block, | ||
| 845 | + bias=bias, | ||
| 846 | + unit_flag=unit_flag, | ||
| 847 | + ) | ||
| 809 | 848 | ||
| 810 | - @jit | 849 | + @host |
| 811 | def run(self, gm_a: Tensor, gm_b: Tensor, gm_c: Tensor, gm_bias: Tensor): | 850 | def run(self, gm_a: Tensor, gm_b: Tensor, gm_c: Tensor, gm_bias: Tensor): |
| 812 | self.bmm_kernel[self.used_core_num](gm_a, gm_b, gm_c, gm_bias) | 851 | self.bmm_kernel[self.used_core_num](gm_a, gm_b, gm_c, gm_bias) |
| 813 | 852 | ||
| @@ -855,11 +894,23 @@ def batch_matmul( | |||
| 855 | bias, bias_shared, bias_full = _normalize_bias(bias, a, c_batch, n) | 894 | bias, bias_shared, bias_full = _normalize_bias(bias, a, c_batch, n) |
| 856 | 895 | ||
| 857 | ctx = _KernelCtx( | 896 | ctx = _KernelCtx( |
| 858 | - a_kern=a_kern, b_kern=b_kern, bias=bias, c_batch=c_batch, | 897 | + a_kern=a_kern, |
| 859 | - m=m, n=n, k=k, hf32=hf32, | 898 | + b_kern=b_kern, |
| 860 | - kern_trans_a=a_transposed, kern_trans_b=not b_transposed, | 899 | + bias=bias, |
| 861 | - a_full=a_full, a_terms=a_terms, b_full=b_full, b_terms=b_terms, | 900 | + c_batch=c_batch, |
| 862 | - bias_shared=bias_shared, bias_full=bias_full, squeeze_out=squeeze_out, | 901 | + m=m, |
| 902 | + n=n, | ||
| 903 | + k=k, | ||
| 904 | + hf32=hf32, | ||
| 905 | + kern_trans_a=a_transposed, | ||
| 906 | + kern_trans_b=not b_transposed, | ||
| 907 | + a_full=a_full, | ||
| 908 | + a_terms=a_terms, | ||
| 909 | + b_full=b_full, | ||
| 910 | + b_terms=b_terms, | ||
| 911 | + bias_shared=bias_shared, | ||
| 912 | + bias_full=bias_full, | ||
| 913 | + squeeze_out=squeeze_out, | ||
| 863 | ) | 914 | ) |
| 864 | out = _degenerate_result(ctx) | 915 | out = _degenerate_result(ctx) |
| 865 | if out is not None: | 916 | if out is not None: |
| @@ -25,10 +25,10 @@ from cannbotdsl.ops.arch import get_block_idx, get_block_num | |||
| 25 | from cannbotdsl.channel import Channel | 25 | from cannbotdsl.channel import Channel |
| 26 | 26 | ||
| 27 | 27 | ||
| 28 | -from cannbotdsl.lang.jit import jit | 28 | +from cannbotdsl.lang.host import host |
| 29 | from cannbotdsl.lang.kernel import kernel | 29 | from cannbotdsl.lang.kernel import kernel |
| 30 | from cannbotdsl.ops.matmul import matmul as dsl_matmul | 30 | from cannbotdsl.ops.matmul import matmul as dsl_matmul |
| 31 | -from cannbotdsl.tensor import tile_view | 31 | +from cannbotdsl.tensor import tile_slice |
| 32 | from cannbotdsl import MemLoc, Tensor | 32 | from cannbotdsl import MemLoc, Tensor |
| 33 | from cannbotdsl.ops.memcpy import make_copy_engine, mem_copy | 33 | from cannbotdsl.ops.memcpy import make_copy_engine, mem_copy |
| 34 | 34 | ||
| @@ -669,12 +669,8 @@ class MatmulKernel: | |||
| 669 | MemLoc.L0C, shape=(t.base_m, t.base_n), dtype=dtypes.float32, depth=t.l0c_db | 669 | MemLoc.L0C, shape=(t.base_m, t.base_n), dtype=dtypes.float32, depth=t.l0c_db |
| 670 | ) | 670 | ) |
| 671 | 671 | ||
| 672 | - nd2nz_engine_a = make_copy_engine( | 672 | + nd2nz_engine_a = make_copy_engine(format_transform="nd2nz") |
| 673 | - format_transform="nd2nz", dtype=t.a_dtype, pad_value=0.0 | 673 | + nd2nz_engine_b = make_copy_engine(format_transform="nd2nz") |
| 674 | - ) | ||
| 675 | - nd2nz_engine_b = make_copy_engine( | ||
| 676 | - format_transform="nd2nz", dtype=t.b_dtype, pad_value=0.0 | ||
| 677 | - ) | ||
| 678 | 674 | ||
| 679 | n_tiles = self.n_tiles | 675 | n_tiles = self.n_tiles |
| 680 | main_window = self.main_window | 676 | main_window = self.main_window |
| @@ -702,36 +698,46 @@ class MatmulKernel: | |||
| 702 | 698 | ||
| 703 | n_idx = (n_tiles - 1 - n_idx) if (row_idx % 2 != 0) else n_idx | 699 | n_idx = (n_tiles - 1 - n_idx) if (row_idx % 2 != 0) else n_idx |
| 704 | 700 | ||
| 701 | + l0c_acc = l0c.produce() | ||
| 705 | for k_l1_idx in range(k_l1_tiles): | 702 | for k_l1_idx in range(k_l1_tiles): |
| 706 | - gm_a_tile = tile_view(gm_a, (t.base_m, t.k_l1), (m_idx, k_l1_idx)) | 703 | + gm_a_tile = tile_slice(gm_a, (t.base_m, t.k_l1), (m_idx, k_l1_idx)) |
| 707 | - gm_b_tile = tile_view(gm_b, (t.base_n, t.k_l1), (n_idx, k_l1_idx)) | 704 | + gm_b_tile = tile_slice(gm_b, (t.base_n, t.k_l1), (n_idx, k_l1_idx)) |
| 708 | mem_copy( | 705 | mem_copy( |
| 709 | - l1_a, | 706 | + l1_a.produce(), |
| 710 | gm_a_tile, | 707 | gm_a_tile, |
| 711 | engine=nd2nz_engine_a, | 708 | engine=nd2nz_engine_a, |
| 712 | l2_cache_ctl=self.l2_cache_ctl_a, | 709 | l2_cache_ctl=self.l2_cache_ctl_a, |
| 713 | ) | 710 | ) |
| 711 | + l1_a_tensor = l1_a.consume() | ||
| 714 | mem_copy( | 712 | mem_copy( |
| 715 | - l1_b, | 713 | + l1_b.produce(), |
| 716 | gm_b_tile, | 714 | gm_b_tile, |
| 717 | engine=nd2nz_engine_b, | 715 | engine=nd2nz_engine_b, |
| 718 | l2_cache_ctl=self.l2_cache_ctl_b, | 716 | l2_cache_ctl=self.l2_cache_ctl_b, |
| 719 | ) | 717 | ) |
| 718 | + l1_b_tensor = l1_b.consume() | ||
| 720 | 719 | ||
| 721 | for k_l0_idx in range(k_l0_per_l1): | 720 | for k_l0_idx in range(k_l0_per_l1): |
| 722 | - l1_a_slice = tile_view(l1_a, (t.base_m, t.base_k), (0, k_l0_idx)) | 721 | + l1_a_slice = tile_slice( |
| 723 | - l1_b_slice = tile_view(l1_b, (t.base_n, t.base_k), (0, k_l0_idx)) | 722 | + l1_a_tensor, (t.base_m, t.base_k), (0, k_l0_idx) |
| 724 | - mem_copy(l0a, l1_a_slice) | 723 | + ) |
| 725 | - mem_copy(l0b, l1_b_slice) | 724 | + l1_b_slice = tile_slice( |
| 725 | + l1_b_tensor, (t.base_n, t.base_k), (0, k_l0_idx) | ||
| 726 | + ) | ||
| 727 | + mem_copy(l0a.produce(), l1_a_slice) | ||
| 728 | + l0a_tensor = l0a.consume() | ||
| 729 | + mem_copy(l0b.produce(), l1_b_slice) | ||
| 730 | + l0b_tensor = l0b.consume() | ||
| 726 | global_k = k_l1_idx * k_l0_per_l1 + k_l0_idx | 731 | global_k = k_l1_idx * k_l0_per_l1 + k_l0_idx |
| 727 | - dsl_matmul(l0c, l0a, l0b, init=(global_k == 0)) | 732 | + dsl_matmul(l0c_acc, l0a_tensor, l0b_tensor, init=(global_k == 0)) |
| 733 | + l0c_tensor = l0c_acc | ||
| 728 | 734 | ||
| 729 | - gm_c_tile = tile_view(gm_c, (t.base_m, t.base_n), (m_idx, n_idx)) | 735 | + gm_c_tile = tile_slice(gm_c, (t.base_m, t.base_n), (m_idx, n_idx)) |
| 730 | - mem_copy(gm_c_tile, l0c) | 736 | + mem_copy(gm_c_tile, l0c_tensor) |
| 731 | 737 | ||
| 732 | # JIT wrapper used by the Python interface. The bracket syntax launches the | 738 | # JIT wrapper used by the Python interface. The bracket syntax launches the |
| 733 | # kernel with the number of AIC cores selected by host-side tiling. | 739 | # kernel with the number of AIC cores selected by host-side tiling. |
| 734 | - @jit | 740 | + @host |
| 735 | def run(self, gm_a: Tensor, gm_b: Tensor, gm_c: Tensor): | 741 | def run(self, gm_a: Tensor, gm_b: Tensor, gm_c: Tensor): |
| 736 | self.matmul_kernel[self.t.used_core_num](gm_a, gm_b, gm_c) | 742 | self.matmul_kernel[self.t.used_core_num](gm_a, gm_b, gm_c) |
| 737 | 743 | ||
| @@ -64,7 +64,6 @@ _TEST_SHAPES = [ | |||
| 64 | pytest.param(1, 32, 1, id="M1-K32-N1-minimal"), | 64 | pytest.param(1, 32, 1, id="M1-K32-N1-minimal"), |
| 65 | pytest.param(256, 256, 256, id="M256-K256-N256-multiTile"), | 65 | pytest.param(256, 256, 256, id="M256-K256-N256-multiTile"), |
| 66 | pytest.param(256, 736, 96, id="M256-K736-N96-mixedTail3"), | 66 | pytest.param(256, 736, 96, id="M256-K736-N96-mixedTail3"), |
| 67 | - pytest.param(2048, 2176, 2048, id="M2048-K2176-N2048-b1tail"), | ||
| 68 | pytest.param(2048, 4608, 2048, id="M2048-K4608-N2048-kw128s36"), | 67 | pytest.param(2048, 4608, 2048, id="M2048-K4608-N2048-kw128s36"), |
| 69 | ] | 68 | ] |
| 70 | 69 | ||