feat(codegen): cube gemv 1-row lhs on a2a3 + gemv/gemv_bias/gemv_acc ST#1847
feat(codegen): cube gemv 1-row lhs on a2a3 + gemv/gemv_bias/gemv_acc ST#1847Little-oil wants to merge 1 commit into
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for single-row Mat operands (cube gemv lhs / bias) in the layout deduction logic and introduces comprehensive runtime tests for the gemv family (gemv, gemv_bias, and gemv_acc). Feedback on the changes highlights two issues: first, the row dimension deduction in memory.cpp does not correctly handle batched gemv operations where the shape rank is greater than 2; second, the GemvProgram test definition expects a bias parameter even when bias=False and the tensor is omitted from define_tensors(), which will cause compilation or runtime errors.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| const auto& shp = shapes_tuple->elements_; | ||
| ExprPtr row_dim = (transpose && shp.size() == 2) ? shp[1] : shp[0]; |
There was a problem hiding this comment.
For batched gemv operations where shp.size() > 2 (e.g., shape [B, 1, K]), the matrix dimensions are the last two dimensions. Using shp[0] as the row dimension incorrectly targets the batch dimension B instead of the actual matrix row dimension shp[shp.size() - 2]. This will cause the layout deduction to fail to detect single-row operands for batched gemv, leading to compilation failures.
We should determine the row dimension dynamically based on the shape rank and whether transpose is enabled.
| const auto& shp = shapes_tuple->elements_; | |
| ExprPtr row_dim = (transpose && shp.size() == 2) ? shp[1] : shp[0]; | |
| const auto& shp = shapes_tuple->elements_; | |
| ExprPtr row_dim = shp.size() >= 2 ? (transpose ? shp.back() : shp[shp.size() - 2]) : shp[0]; |
References
- Avoid assuming tile shape dimensions are always ConstInt and immediately dereferencing them with ->value_. Instead, forward the existing shape expressions directly.
| def get_program(self) -> Any: | ||
| k, n = self._k, self._n | ||
| ab = _PL_DT[self._ab] | ||
| use_bias = self._bias | ||
| a_v = [1, VALID_K] if self._narrow == "K" else [1, k] | ||
| b_v = [VALID_K, n] if self._narrow == "K" else [k, n] | ||
|
|
||
| @pl.program | ||
| class GemvProgram: | ||
| @pl.function(type=pl.FunctionType.InCore) | ||
| def kernel( | ||
| self, | ||
| a: pl.Tensor[[1, k], ab], | ||
| b: pl.Tensor[[k, n], ab], | ||
| bias: pl.Tensor[[1, n], pl.FP32], | ||
| out: pl.Out[pl.Tensor[[1, n], pl.FP32]], | ||
| ) -> pl.Tensor[[1, n], pl.FP32]: | ||
| tile_a = pl.load(a, [0, 0], [1, k], valid_shapes=a_v, target_memory=pl.MemorySpace.Mat) | ||
| tile_b = pl.load(b, [0, 0], [k, n], valid_shapes=b_v, target_memory=pl.MemorySpace.Mat) | ||
| if use_bias: | ||
| tile_bias = pl.load( | ||
| bias, [0, 0], [1, n], valid_shapes=[1, n], target_memory=pl.MemorySpace.Mat | ||
| ) | ||
| res = pl.tile.gemv_bias(tile_a, tile_b, tile_bias) | ||
| else: | ||
| res = pl.tile.gemv(tile_a, tile_b) | ||
| out = pl.store(res, [0, 0], out) | ||
| return out | ||
|
|
||
| @pl.function(type=pl.FunctionType.Orchestration) | ||
| def orchestrator( | ||
| self, | ||
| a: pl.Tensor[[1, k], ab], | ||
| b: pl.Tensor[[k, n], ab], | ||
| bias: pl.Tensor[[1, n], pl.FP32], | ||
| out: pl.Out[pl.Tensor[[1, n], pl.FP32]], | ||
| ) -> pl.Tensor[[1, n], pl.FP32]: | ||
| out = self.kernel(a, b, bias, out) | ||
| return out | ||
|
|
||
| return GemvProgram |
There was a problem hiding this comment.
When bias=False, the bias tensor is omitted from define_tensors(). However, the kernel and orchestrator function signatures in GemvProgram still expect bias as a parameter. This mismatch will cause a compilation or runtime error during test execution because the test runner will not pass the bias argument.
To resolve this, we should conditionally define the GemvProgram class based on whether self._bias is enabled.
def get_program(self) -> Any:
k, n = self._k, self._n
ab = _PL_DT[self._ab]
a_v = [1, VALID_K] if self._narrow == "K" else [1, k]
b_v = [VALID_K, n] if self._narrow == "K" else [k, n]
if self._bias:
@pl.program
class GemvProgram:
@pl.function(type=pl.FunctionType.InCore)
def kernel(
self,
a: pl.Tensor[[1, k], ab],
b: pl.Tensor[[k, n], ab],
bias: pl.Tensor[[1, n], pl.FP32],
out: pl.Out[pl.Tensor[[1, n], pl.FP32]],
) -> pl.Tensor[[1, n], pl.FP32]:
tile_a = pl.load(a, [0, 0], [1, k], valid_shapes=a_v, target_memory=pl.MemorySpace.Mat)
tile_b = pl.load(b, [0, 0], [k, n], valid_shapes=b_v, target_memory=pl.MemorySpace.Mat)
tile_bias = pl.load(
bias, [0, 0], [1, n], valid_shapes=[1, n], target_memory=pl.MemorySpace.Mat
)
res = pl.tile.gemv_bias(tile_a, tile_b, tile_bias)
out = pl.store(res, [0, 0], out)
return out
@pl.function(type=pl.FunctionType.Orchestration)
def orchestrator(
self,
a: pl.Tensor[[1, k], ab],
b: pl.Tensor[[k, n], ab],
bias: pl.Tensor[[1, n], pl.FP32],
out: pl.Out[pl.Tensor[[1, n], pl.FP32]],
) -> pl.Tensor[[1, n], pl.FP32]:
out = self.kernel(a, b, bias, out)
return out
return GemvProgram
else:
@pl.program
class GemvProgram:
@pl.function(type=pl.FunctionType.InCore)
def kernel(
self,
a: pl.Tensor[[1, k], ab],
b: pl.Tensor[[k, n], ab],
out: pl.Out[pl.Tensor[[1, n], pl.FP32]],
) -> pl.Tensor[[1, n], pl.FP32]:
tile_a = pl.load(a, [0, 0], [1, k], valid_shapes=a_v, target_memory=pl.MemorySpace.Mat)
tile_b = pl.load(b, [0, 0], [k, n], valid_shapes=b_v, target_memory=pl.MemorySpace.Mat)
res = pl.tile.gemv(tile_a, tile_b)
out = pl.store(res, [0, 0], out)
return out
@pl.function(type=pl.FunctionType.Orchestration)
def orchestrator(
self,
a: pl.Tensor[[1, k], ab],
b: pl.Tensor[[k, n], ab],
out: pl.Out[pl.Tensor[[1, n], pl.FP32]],
) -> pl.Tensor[[1, n], pl.FP32]:
out = self.kernel(a, b, out)
return out
return GemvProgramReferences
- When writing unit tests for IR passes that target specific function types (e.g., InCore), ensure the test functions are explicitly annotated with that type.
Little-oil
force-pushed
the
fix-pr1823-op-isa-gaps
branch
2 times, most recently
from
7ed48ad to
5c10633
Compare
The cube gemv / gemv_bias / gemv_acc family takes a 1-row lhs [1, K]. PyPTO loaded it to Mat as the NZ fractal (col_major blayout + row_major slayout) that multi-row matmul operands use, so the Mat->Left move lowered through the pto-isa TExtractToA path, whose srcRow/dstRow % 16 == 0 static_assert rejects a single row. pto-isa's own RunTGEMV reference declares the gemv lhs as an ND row vector: Tile<Mat, 1, K, BLayout::RowMajor, ..., SLayout::NoneBox>. That row_major + none_box pair routes TMovToLeft to the rows==1 vector path (TExtractToAVector), which carries no row-alignment constraint. Tag a single-row Mat tile row_major/none_box at load so the gemv lhs (and the layout-agnostic 1-row bias) take that path. Add a2a3 ST for gemv, gemv_bias and gemv_acc (>=2 cases each, including a narrowed valid_shape). The vector path needs K aligned to 512/sizeof (K % 128 fp32, K % 256 bf16), and gemv keeps the whole rhs resident in L0B (no K-split), so the shapes honour both limits.
1 participant
What
gemv/gemv_bias/gemv_accfor their 1-row lhs[1, K]on a2a3.Why
The gemv lhs is a single row. PyPTO loaded it to Mat as the NZ fractal
(
col_majorblayout +row_majorslayout) that multi-row matmul operands use, sothe
Mat -> Leftmove lowered through the pto-isaTExtractToApath, whosesrcRow/dstRow % 16 == 0static_assert rejects a single row. This is whyPR #1823 deferred the gemv family.
pto-isa's own
RunTGEMVreference (tests/npu/a2a3/.../tmatmul_kernel.cpp)declares the gemv lhs as an ND row vector —
Tile<Mat, 1, K, BLayout::RowMajor, ..., SLayout::NoneBox>. Thatrow_major + none_boxpair routesTMovToLeftto therows==1vector path(
TExtractToAVector), which has no row-alignment constraint. So the fix tagsa single-row Mat tile
row_major/none_boxat load. The 1-row bias takes the samepath (the
Mat -> Biasmove is layout-agnostic — it only assertsRows == 1).Shapes / ISA constraints honoured by the tests
512 / sizeof(K % 128fp32,K % 256bf16); tests use K=128/256.
[K, N]resident in L0B (no K-split), soK * N * sizeof <= 64 KiB.valid_shape(physical K aligned,valid K narrowed).
Validation
PTOAS + a2a3 pto-isa toolchain (every prior
static_assertis gone).test_tile_ops.py,test_pto_codegen.py).worktree runtime env was unable to execute any a2a3 kernel (even CI-green
matmul_biasreturned507018), so I could not numerically validate locally.Keeping this as a draft until the a2a3 ST is green in CI.
Relates to #1846 (tracking issue for the remaining a2a3 tile-op gaps deferred in
PR #1823).