[Codegen][CPU] Drop the ACC stride from the inner_tiled ukernel ABI. (#24652)
New bottom of the LLVMCPU C-ukernel stack.
The ACC stride threaded to the C ukernels is always a compile-time
constant
equal to one intrinsic's ACC fragment size (the data-tiled ACC tile is
contiguous), so it folds to a constant after inlining and carries no
information the intrinsic-specific ukernel doesn't already encode. This
drops
it:
- `IREECPUAttrs.cpp`: remove `getAccInnermostCrossIntrinsicDim`; build
the
`ukernel.generic` with `num_strided_outer_dims=0`, so all shaped
operands
(LHS, RHS, ACC) lower to `(base, offset)` only.
- Both seed ukernel `.c` signatures lose the `acc_stride` parameter.
This addresses review feedback on the now-merged #24569 / #24571 (why
pass a
per-operand `strided_dims` array, and why pass an ACC stride that's
statically
known). It makes the three operands uniform and removes the
`getAccInnermostCrossIntrinsicDim` swizzle machinery.
Verified: the lowered `llvm.call` to the ukernel now has no stride
argument,
and the ukernel lit tests (`lower_inner_tiled_to_bitcode_ukernel*`,
`e2e_inner_tiled_pipeline`, `select_ukernel`) pass.
Progress towards https://github.com/iree-org/iree/issues/24574.
Signed-off-by: Benoit Jacob <jacob.benoit.1@gmail.com>
diff --git a/compiler/plugins/target/LLVMCPU/builtins/ukernel/iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16.c b/compiler/plugins/target/LLVMCPU/builtins/ukernel/iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16.c
index 978c478..c7a5728 100644
--- a/compiler/plugins/target/LLVMCPU/builtins/ukernel/iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16.c
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16.c
@@ -22,8 +22,9 @@
// template.
//
// ABI: each shaped operand is passed as (base pointer, element offset) so the
-// caller doesn't need a GEP before the call; the accumulator additionally
-// gets the element stride of its innermost cross-intrinsic (N) dimension.
+// caller doesn't need a GEP before the call. No strides are passed: the ACC
+// tile is contiguous, so the ukernel addresses each intrinsic's fragment from
+// `intrinsics_{m,n}` and its own fixed fragment size.
//
// NOTE (seed scaffolding): this initial seed has a stub body. It exists so
// that the surrounding *framework* -- bitcode build, embedding,
@@ -33,16 +34,14 @@
IREE_UK_ALWAYS_INLINE
void iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16(
const uint16_t *lhs_base, int64_t lhs_offset, const uint16_t *rhs_base,
- int64_t rhs_offset, float *acc_base, int64_t acc_offset, int64_t acc_stride,
- int32_t k_outer, int32_t intrinsics_m, int32_t intrinsics_n,
- int32_t intrinsics_k) {
+ int64_t rhs_offset, float *acc_base, int64_t acc_offset, int32_t k_outer,
+ int32_t intrinsics_m, int32_t intrinsics_n, int32_t intrinsics_k) {
(void)lhs_base;
(void)lhs_offset;
(void)rhs_base;
(void)rhs_offset;
(void)acc_base;
(void)acc_offset;
- (void)acc_stride;
(void)k_outer;
(void)intrinsics_m;
(void)intrinsics_n;
diff --git a/compiler/plugins/target/LLVMCPU/builtins/ukernel/iree_uk_mma_x86_avx512vnni_16x16x2_i32_i8_casti16.c b/compiler/plugins/target/LLVMCPU/builtins/ukernel/iree_uk_mma_x86_avx512vnni_16x16x2_i32_i8_casti16.c
index 478ea1a..9e27639 100644
--- a/compiler/plugins/target/LLVMCPU/builtins/ukernel/iree_uk_mma_x86_avx512vnni_16x16x2_i32_i8_casti16.c
+++ b/compiler/plugins/target/LLVMCPU/builtins/ukernel/iree_uk_mma_x86_avx512vnni_16x16x2_i32_i8_casti16.c
@@ -36,21 +36,19 @@
// on this case.
// ABI matches the inner_tiled -> ukernel.generic lowering (see
// `iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16`): each shaped operand passed
-// as (base, element offset), ACC additionally as its innermost
-// cross-intrinsic stride, then the scalar `k_outer` / `intrinsics_{m,n,k}`.
+// as (base, element offset) only (no strides; the ACC tile is contiguous),
+// then the scalar `k_outer` / `intrinsics_{m,n,k}`.
IREE_UK_ALWAYS_INLINE
void iree_uk_mma_x86_avx512vnni_16x16x2_i32_i8_casti16(
const void *lhs_base, int64_t lhs_offset, const void *rhs_base,
- int64_t rhs_offset, void *acc_base, int64_t acc_offset, int64_t acc_stride,
- int32_t k_outer, int32_t intrinsics_m, int32_t intrinsics_n,
- int32_t intrinsics_k) {
+ int64_t rhs_offset, void *acc_base, int64_t acc_offset, int32_t k_outer,
+ int32_t intrinsics_m, int32_t intrinsics_n, int32_t intrinsics_k) {
(void)lhs_base;
(void)lhs_offset;
(void)rhs_base;
(void)rhs_offset;
(void)acc_base;
(void)acc_offset;
- (void)acc_stride;
(void)k_outer;
(void)intrinsics_m;
(void)intrinsics_n;
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
index 2eb1f43..8997f88 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
@@ -1255,68 +1255,19 @@
return true;
}
-// Returns the index, in the ACC operand's shape, of the innermost
-// CrossIntrinsic dimension (the N cross-intrinsic dim for our layouts), or
-// nullopt if it is dynamic / absent. The ukernel needs this dim's stride to
-// address each unrolled intrinsic's ACC fragment.
-//
-// Example: with `MMA_X86_AVX512BF16_1x16x2` (each intrinsic produces a 1x16
-// f32 fragment) and intrinsics_m = intrinsics_n = 2, the ACC tile holds a 2x2
-// grid of such fragments. The two CrossIntrinsic dims are that grid's M and N;
-// the innermost is N, so this returns the index of the N-grid dim in the ACC
-// result shape, whose stride is the element distance from fragment (m, n) to
-// (m, n+1).
-static std::optional<unsigned>
-getAccInnermostCrossIntrinsicDim(IREE::Codegen::InnerTiledOp op,
- DataTiledMMAAttr mma) {
- auto outputType = dyn_cast<ShapedType>(op.getResultTypes()[0]);
- if (!outputType) {
- return std::nullopt;
- }
- Codegen::TileSwizzle accSwizzle = getSwizzle(mma, /*operandIdx=*/2);
- SmallVector<Codegen::TileSwizzle::Dim> swizzleDims;
- for (const Codegen::TileSwizzle::ExpandShapeDimVectorType &group :
- accSwizzle.expandShape()) {
- swizzleDims.append(group.begin(), group.end());
- }
- applyPermutationToVector(swizzleDims, accSwizzle.permutation());
- int rankDiff = outputType.getRank() - static_cast<int>(swizzleDims.size());
- auto crossIntrinsic = Codegen::TileSwizzle::Dim::Kind::CrossIntrinsic;
- for (size_t i = swizzleDims.size(); i-- > 0;) {
- if (swizzleDims[i].kind() != crossIntrinsic) {
- continue;
- }
- int outputIdx = i + rankDiff;
- if (outputType.isDynamicDim(outputIdx)) {
- return std::nullopt;
- }
- return outputIdx;
- }
- // No CrossIntrinsic dims (intrinsics_m == intrinsics_n == 1): the single
- // fragment sits at the start of the inner tile.
- if (!swizzleDims.empty()) {
- return rankDiff;
- }
- return std::nullopt;
-}
-
// Rewrites an `inner_tiled` carrying a CPU `DataTiledMMAAttr` to a
// `ukernel.generic`, threading the data-tiled-MMA scalar parameters as
// operands so the ukernel can loop over arbitrary `intrinsics_{m,n,k}`:
// ins(lhs, rhs) outs(acc) (k_outer, intrinsics_m, intrinsics_n, intrinsics_k)
-// plus a `strided_dims` entry giving the ACC's innermost cross-intrinsic
-// stride.
+// All shaped operands are passed as `(base, offset)` only, no strides: the
+// ACC tile is contiguous, so the ukernel derives each intrinsic's ACC
+// fragment offset from the compile-time `intrinsics_{m,n}` and its own
+// (fixed, per-intrinsic) fragment size.
static LogicalResult
handleInnerTiledMmaUkernel(RewriterBase &rewriter, StringRef name,
IREE::Codegen::InnerTiledOp op, DataTiledMMAAttr mma,
ArrayRef<Value> inputs, ArrayRef<Value> outputs,
DictionaryAttr fnDefAttrs) {
- std::optional<unsigned> accInnerDim =
- getAccInnermostCrossIntrinsicDim(op, mma);
- if (!accInnerDim) {
- return rewriter.notifyMatchFailure(
- op, "ACC innermost cross-intrinsic dim is dynamic or absent");
- }
Location loc = op.getLoc();
Type i32 = rewriter.getI32Type();
auto constI32 = [&](int64_t v) {
@@ -1327,23 +1278,15 @@
Value kOuter = arith::IndexCastOp::create(
rewriter, loc, i32,
tensor::DimOp::create(rewriter, loc, op.getInputs()[0], 1));
- // `strided_dims` is the `ukernel.generic` ABI knob for which dims' strides
- // are passed to the C function: a list per shaped operand (here LHS, RHS,
- // ACC), each naming the dims whose stride follows that operand's
- // `(base, offset)` in the call. An empty list passes `base, offset` only; a
- // null attribute (the default, used by simpler ukernels) passes all strides.
- // Here only ACC needs one — the innermost cross-intrinsic (N) dim — so the
- // ukernel can address each unrolled intrinsic's ACC fragment. LHS/RHS are
- // contiguous in the data-tiled layout, so they take no stride argument.
- SmallVector<SmallVector<int64_t>> stridedDims(3, {});
- stridedDims[2].push_back(*accInnerDim);
+ // All shaped operands (LHS, RHS, ACC) are passed as `(base, offset)` only;
+ // `num_strided_outer_dims=0` means no stride arguments are appended.
DictionaryAttr discardableAttrs = op->getDiscardableAttrDictionary();
auto newOp = rewriter.replaceOpWithNewOp<IREE::Codegen::UKernelGenericOp>(
op, op.getOutputs().getTypes(), name, inputs, outputs,
ValueRange{kOuter, constI32(mma.getIntrinsicsM()),
constI32(mma.getIntrinsicsN()),
constI32(mma.getIntrinsicsK())},
- fnDefAttrs, stridedDims);
+ fnDefAttrs, /*num_strided_outer_dims=*/0);
newOp->setDiscardableAttrs(discardableAttrs);
return success();
}
