[Mlir-commits] [mlir] [mlir][linalg] Scaled contraction op (PR #203958)
Adam Siemieniuk
llvmlistbot at llvm.org
Mon Jun 15 10:54:39 PDT 2026
https://github.com/adam-smnk created https://github.com/llvm/llvm-project/pull/203958
A new contraction op with extended semantics to represent input scaling using additional scale values as described by their corresponding indexing maps.
RFC:
Assisted-by: Claude
Co-authored-by: Md Asghar Ahmad Shahid <md.asghar.ahmad.shahid at intel.com>
>From 3d3788dfa9b5fc20c348d16ddc41360c63083022 Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk at intel.com>
Date: Thu, 7 May 2026 14:44:03 +0200
Subject: [PATCH] [mlir][linalg] Scaled contraction op
A new contraction op with extended semantics to represent input
scaling using additional scale value as described by their
corresponding indexing maps.
RFC:
Assisted-by: Claude
Co-authored-by: Md Asghar Ahmad Shahid <md.asghar.ahmad.shahid at intel.com>
---
.../Dialect/Linalg/IR/LinalgStructuredOps.td | 144 ++++++
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 414 +++++++++++++++---
.../Linalg/Transforms/TilingInterfaceImpl.cpp | 102 ++++-
.../Dialect/Linalg/generalize-named-ops.mlir | 49 +++
mlir/test/Dialect/Linalg/invalid.mlir | 187 ++++++++
mlir/test/Dialect/Linalg/loops.mlir | 58 +++
mlir/test/Dialect/Linalg/named-ops.mlir | 162 +++++++
mlir/test/Dialect/Linalg/roundtrip.mlir | 28 ++
.../Dialect/Linalg/transform-op-fuse.mlir | 177 +++++++-
.../Dialect/Linalg/transform-op-tile.mlir | 196 +++++++++
10 files changed, 1435 insertions(+), 82 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index 5998f736ced34..c7dbc4cc1d6dd 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -1213,6 +1213,150 @@ def BatchReduceMatmulOp : LinalgStructuredBase_Op<"batch_reduce_matmul", [
}];
}
+//===----------------------------------------------------------------------===//
+// Scaled Contract op.
+//===----------------------------------------------------------------------===//
+
+def ScaledContractOp : LinalgStructuredBase_Op<"scaled_contract", [
+ AttrSizedOperandSegments]> {
+ let summary = [{
+ Perform a scaled contraction on two inputs where their scaling is described
+ by affine maps of the two corresponding scales, accumulating into the output.
+ }];
+ let description = [{
+ Extends the semantics of `linalg.contract` by including extra scaling factors
+ for the inputs `A` and `B` described by the scales' affine maps.
+ The data types of the inputs are chosen independently from the scales.
+
+ The semantics of contracting inputs `A` and `B` with scales `scale_A` and
+ `scale_B` on top of `C` to produce output `D` is given by:
+
+ `D[H] = (SUM_{(I ∪ J) \ H} (A[I] * scale_A[I]) * (B[J] * scale_B[J])) + C[H]`
+
+ The iteration type of each dim is inferred.
+
+ The input operands `A`, `B`, and `C` follow the standard contraction semantics
+ together with broadcasting and transposition rules.
+ See `linalg.contract` for further details.
+
+ **Affine maps for scales**
+
+ Each scale indexing map describes the scaling scheme for the corresponding
+ input `A` and `B`. The following scaling schemes are supported:
+
+ - Tensor scaling - no dimensions: a scalar value used to scale the whole input.
+ - Dimension scaling - dimension `d`: one scale value per element along `d`.
+ - Block scaling - binary expression `d floordiv B`: one scale value for every
+ `B` elements along `d` where `B` corresponds to the block size.
+
+ Scale maps may only use dimensions present in the corresponding input's map.
+ A dimension absent from the scale map indicates that a single scale value is
+ reused across the whole input dimension.
+
+ **Dim sizes for scales**
+
+ - Tensor scaling: a unit-size scale.
+ - Dimension scaling (`d`): the scale dim size equals the input dim size.
+ - Block scaling (`d floordiv B`): the scale dim size equals
+ `ceil(input_dim_size / B)`.
+
+ For dynamic shapes, all sizes are assumed to be correct at runtime.
+
+ Example scaled matmul with transposed B:
+ ```mlir
+ %D = linalg.scaled_contract
+ indexing_maps = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>, // A - 32x128 block scale
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>, // B - row-wise scale
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %scale_A, %B, %scale_B
+ : tensor<200x512xi8>, tensor<7x4xf8E8M0FNU>,
+ tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C: tensor<200x128xf32>) -> tensor<200x128xf32>
+ ```
+
+ Numeric casting is performed on the inputs and scales, promoting/truncating
+ them to the same data type as the accumulator/output.
+
+ TODO: Allow control over the combining/accumulating op and possibly the
+ multiplication op.
+ }];
+
+ let arguments = (ins
+ Variadic<AnyType>:$inputs,
+ Variadic<AnyShaped>:$outputs,
+ AffineMapArrayAttr:$indexing_maps,
+ DefaultValuedOptionalAttr<TypeFnAttr, "TypeFn::cast_signed">:$cast
+ );
+ let results = (outs Variadic<AnyShaped>:$result_tensors);
+ let regions = (region SizedRegion<1>:$combiner);
+
+ let skipDefaultBuilders = 1;
+ let builders = [
+ OpBuilder<
+ (ins "ValueRange":$inputs, "ValueRange":$outputs,
+ CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+ [{
+ buildStructuredOp($_builder, $_state, std::nullopt, inputs, outputs,
+ attributes, regionBuilder);
+ }]>,
+ OpBuilder<(ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
+ "ValueRange":$outputs, "ArrayAttr":$indexingMaps,
+ CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+ [{
+ $_state.addAttribute("indexing_maps", indexingMaps);
+ buildStructuredOp($_builder, $_state, resultTensorTypes, inputs,
+ outputs, attributes, regionBuilder);
+ }]>,
+ OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$outputs,
+ "ArrayAttr":$indexingMaps,
+ CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+ [{
+ $_state.addAttribute("indexing_maps", indexingMaps);
+ buildStructuredOp($_builder, $_state, std::nullopt, inputs, outputs,
+ attributes, regionBuilder);
+ }]>
+ ];
+ let hasCustomAssemblyFormat = 1;
+ let hasFolder = 1;
+ let hasVerifier = 1;
+
+ let extraClassDeclaration = structuredOpsBaseDecls # [{
+ // Declare/implement functions necessary for LinalgStructuredInterface.
+
+ /// Infer iterator types for each dim in the domain of IndexingMaps.
+ SmallVector<utils::IteratorType> getIteratorTypesArray();
+
+ /// IndexingMaps always depends on attr associated to current Op instance.
+ bool hasDynamicIndexingMaps() { return true; };
+ bool hasUserDefinedMaps() { return true; };
+
+ static unsigned getNumRegionArgs();
+
+ static void regionBuilder(ImplicitLocOpBuilder &b,
+ Block &block, ArrayRef<NamedAttribute> attrs,
+ function_ref<InFlightDiagnostic()> emitError);
+
+ static std::function<void(ImplicitLocOpBuilder &,
+ Block &, ArrayRef<NamedAttribute>,
+ function_ref<InFlightDiagnostic()>)>
+ getRegionBuilder() {
+ return regionBuilder;
+ }
+
+ std::string getLibraryCallName() {
+ return "op_has_no_registered_library_name";
+ }
+
+ // Implement function necessary for DestinationStyleOpInterface.
+ ::mlir::MutableOperandRange getDpsInitsMutable() {
+ return getOutputsMutable();
+ }
+ }];
+}
+
//===----------------------------------------------------------------------===//
// Named Linalg ops, implemented as a declarative configurations of generic ops.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 1a56c5a483e73..f205290ab54e8 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -4544,65 +4544,67 @@ void ContractOp::print(OpAsmPrinter &p) {
/*elidedAttrs=*/{"indexing_maps", "operandSegmentSizes"});
}
-LogicalResult ContractOp::verify() {
- int iterationSpaceDims = -1;
- // Map iter space dims to #occurrences in inputs' and output's affine_maps:
- // e.g., inOccurrences[0] will hold #times that dim (with index) 0 is used to
- // access an input operand (so occurrence count can be at most 2) and
- // outOccurrences[1] will indicate whether dim 1 occurred in the output, etc.
- SmallVector<size_t> inOccurrences;
- SmallVector<size_t> outOccurrences;
-
- // A helper so that for each operand's affine_map and type we check that ...
- auto checkAffineMapAndType = [&](AffineMap affineMap, Type operandType,
- bool isInput) -> LogicalResult {
- // ... the affine_map is a projected permutation;
- if (!affineMap.isProjectedPermutation())
- return emitError("provided affine_map is not a projected permutation");
-
- // ... the rank of the affine_map's results and corresponding type match;
- if (auto shapedType = dyn_cast<ShapedType>(operandType)) {
- if (affineMap.getNumResults() != shapedType.getRank())
- return emitError("ranks of shaped operand and results of corresponding "
- "affine_map differ");
- } else if (affineMap.getNumResults() != 0) {
- return emitError("affine_map specifies shaped access while operand has "
- "non-shaped type");
- }
-
- // ... the rank of the affine_map's domain is the same as those seen prior;
- if (iterationSpaceDims == -1) {
- iterationSpaceDims = affineMap.getNumDims();
- inOccurrences = SmallVector<size_t>(iterationSpaceDims, 0);
- outOccurrences = SmallVector<size_t>(iterationSpaceDims, 0);
- } else if (iterationSpaceDims != (int)affineMap.getNumDims()) {
- return emitError("iteration spaces of provided affine_maps differ");
- }
-
- // ... update counts of dims used to access either an input or the output.
- for (AffineExpr affineExpr : affineMap.getResults()) {
- auto affineDimExpr = dyn_cast<AffineDimExpr>(affineExpr);
- if (!affineDimExpr)
- llvm_unreachable("affine_map is a projected permutation");
-
- if (isInput)
- inOccurrences[affineDimExpr.getPosition()] += 1;
- else
- outOccurrences[affineDimExpr.getPosition()] += 1;
- }
-
- return success();
- };
-
- for (auto &&[affineMap, operandType, isInput] :
- llvm::zip(getIndexingMapsArray(), getOperandTypes(),
- SmallVector<bool>{true, true, false})) {
- if (failed(checkAffineMapAndType(affineMap, operandType, isInput)))
- return failure(); // NB: checkAffineMapAndType will emit relevant error.
+/// Validate contraction operands indexing maps and shapes.
+/// For a given affine_map and type, checks that:
+/// - the affine_map is a projected permutation;
+/// - the rank of the affine_map's results and the corresponding type match;
+/// - the rank of the affine_map's domain is consistent with prior maps.
+/// Also updates the per-dim input/output occurrence counts.
+static LogicalResult
+checkContractionAffineMapAndType(AffineMap affineMap, Type operandType,
+ bool isInput, int &iterationSpaceDims,
+ SmallVector<size_t> &inOccurrences,
+ SmallVector<size_t> &outOccurrences,
+ function_ref<InFlightDiagnostic()> emitError) {
+ if (!affineMap.isProjectedPermutation())
+ return emitError() << "provided affine_map is not a projected permutation";
+
+ if (auto shapedType = dyn_cast<ShapedType>(operandType)) {
+ if (affineMap.getNumResults() != shapedType.getRank())
+ return emitError()
+ << "ranks of shaped operand and results of corresponding "
+ "affine_map differ";
+ } else if (affineMap.getNumResults() != 0) {
+ return emitError()
+ << "affine_map specifies shaped access while operand has "
+ "non-shaped type";
+ }
+
+ if (iterationSpaceDims == -1) {
+ iterationSpaceDims = affineMap.getNumDims();
+ inOccurrences = SmallVector<size_t>(iterationSpaceDims, 0);
+ outOccurrences = SmallVector<size_t>(iterationSpaceDims, 0);
+ } else if (iterationSpaceDims != (int)affineMap.getNumDims()) {
+ return emitError() << "iteration spaces of provided affine_maps differ";
+ }
+
+ // Update counts of dims used to access either an input or the output.
+ for (AffineExpr affineExpr : affineMap.getResults()) {
+ auto affineDimExpr = dyn_cast<AffineDimExpr>(affineExpr);
+ if (!affineDimExpr)
+ llvm_unreachable("affine_map is a projected permutation");
+
+ if (isInput)
+ inOccurrences[affineDimExpr.getPosition()] += 1;
+ else
+ outOccurrences[affineDimExpr.getPosition()] += 1;
}
+ return success();
+}
+
+/// Validates the contracting dimension constraints given the per-dim
+/// occurrence counts. Checks that:
+/// - every iteration-space dimension is used by at least one operand;
+/// - every dimension is either contracting (appears in both inputs, not in
+/// output) or parallel (appears in exactly one input and in the output);
+/// - at least one contracting dimension exists.
+static LogicalResult
+verifyContractionDims(size_t iterationSpaceDims, ArrayRef<size_t> inOccurrences,
+ ArrayRef<size_t> outOccurrences,
+ function_ref<InFlightDiagnostic()> emitError) {
bool hasContractingDim = false;
- for (size_t dimIndex = 0; dimIndex < (size_t)iterationSpaceDims; dimIndex++) {
+ for (size_t dimIndex = 0; dimIndex < iterationSpaceDims; dimIndex++) {
size_t inOccCount = inOccurrences[dimIndex];
size_t outOccCount = outOccurrences[dimIndex];
@@ -4629,11 +4631,35 @@ LogicalResult ContractOp::verify() {
}
if (!hasContractingDim)
- return emitError("'indexing_maps' do not specify a contracting dimension");
+ return emitError()
+ << "'indexing_maps' do not specify a contracting dimension";
return success();
}
+LogicalResult ContractOp::verify() {
+ int iterationSpaceDims = -1;
+ // Map iter space dims to #occurrences in inputs' and output's affine_maps:
+ // e.g., inOccurrences[0] will hold #times that dim (with index) 0 is used to
+ // access an input operand (so occurrence count can be at most 2) and
+ // outOccurrences[1] will indicate whether dim 1 occurred in the output, etc.
+ SmallVector<size_t> inOccurrences;
+ SmallVector<size_t> outOccurrences;
+
+ for (auto &&[affineMap, operandType, isInput] :
+ llvm::zip(getIndexingMapsArray(), getOperandTypes(),
+ SmallVector<bool>{true, true, false})) {
+ if (failed(checkContractionAffineMapAndType(
+ affineMap, operandType, isInput, iterationSpaceDims, inOccurrences,
+ outOccurrences, [&]() { return emitError(); })))
+ return failure(); // NB: Validation helper emits relevant error.
+ }
+
+ return verifyContractionDims(static_cast<size_t>(iterationSpaceDims),
+ inOccurrences, outOccurrences,
+ [&]() { return emitError(); });
+}
+
LogicalResult ContractOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) {
return memref::foldMemRefCast(*this);
}
@@ -6802,6 +6828,280 @@ Speculation::Speculatability BatchReduceMatmulOp::getSpeculatability() {
return getGenericSpeculatabilityImpl(cast<LinalgOp>(getOperation()));
}
+//===----------------------------------------------------------------------===//
+// ScaledContractOp
+//===----------------------------------------------------------------------===//
+
+SmallVector<utils::IteratorType> ScaledContractOp::getIteratorTypesArray() {
+ AffineMap outAffineMap = getIndexingMapsArray().pop_back_val();
+ // Infer iterator types based on the output.
+ SmallVector<bool> dimsInOutput(outAffineMap.getNumDims(), false);
+ for (auto result : outAffineMap.getResults()) {
+ auto dimExpr = dyn_cast<AffineDimExpr>(result);
+ assert(dimExpr && "affine_map is a projected permutation");
+ dimsInOutput[dimExpr.getPosition()] = true;
+ }
+
+ SmallVector<utils::IteratorType> iteratorTypes;
+ for (auto dimOccursInOutput : dimsInOutput)
+ iteratorTypes.push_back(dimOccursInOutput ? utils::IteratorType::parallel
+ : utils::IteratorType::reduction);
+
+ return iteratorTypes;
+}
+
+unsigned ScaledContractOp::getNumRegionArgs() { return 5; }
+
+/// Implement block region builder, which is called by 'fillStructuredOpRegion'.
+void ScaledContractOp::regionBuilder(
+ ImplicitLocOpBuilder &b, Block &block, ArrayRef<NamedAttribute> attrs,
+ function_ref<InFlightDiagnostic()> emitError) {
+ if (emitError && block.getNumArguments() != 5) {
+ emitError() << "ScaledContractOp regionBuilder expects 5 args, got "
+ << block.getNumArguments();
+ return;
+ }
+ assert(block.getNumArguments() == 5 &&
+ "ScaledContractOp regionBuilder expects 5 args");
+ RegionBuilderHelper helper(b, block);
+
+ TypeFn castSignedness = TypeFn::cast_signed;
+ auto castIter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {
+ return attr.getName() == "cast";
+ });
+ if (castIter != attrs.end()) {
+ if (auto attr = llvm::dyn_cast<TypeFnAttr>(castIter->getValue()))
+ castSignedness = attr.getValue();
+ }
+
+ // TODO: Support fields with operators besides mult & add.
+ Type outType = block.getArgument(4).getType();
+ Value lhsAtOutType =
+ helper.buildTypeFn(castSignedness, outType, block.getArgument(0));
+ Value lhsScaleAtOutType =
+ helper.buildTypeFn(castSignedness, outType, block.getArgument(1));
+ Value scaledLhs = helper.buildBinaryFn(BinaryFn::mul, lhsAtOutType,
+ lhsScaleAtOutType, emitError);
+ if (!scaledLhs)
+ return;
+ Value rhsAtOutType =
+ helper.buildTypeFn(castSignedness, outType, block.getArgument(2));
+ Value rhsScaleAtOutType =
+ helper.buildTypeFn(castSignedness, outType, block.getArgument(3));
+ Value scaledRhs = helper.buildBinaryFn(BinaryFn::mul, rhsAtOutType,
+ rhsScaleAtOutType, emitError);
+ if (!scaledRhs)
+ return;
+ Value productAtOutType =
+ helper.buildBinaryFn(BinaryFn::mul, scaledLhs, scaledRhs, emitError);
+ if (!productAtOutType)
+ return;
+ Value result = helper.buildBinaryFn(BinaryFn::add, block.getArgument(4),
+ productAtOutType, emitError);
+ if (!result)
+ return;
+ helper.yieldOutputs({result});
+}
+
+ParseResult ScaledContractOp::parse(OpAsmParser &parser,
+ OperationState &result) {
+ FailureOr<ArrayAttr> indexingMapsAttr = parseIndexingMapsAttr(parser);
+ if (failed(indexingMapsAttr) || *indexingMapsAttr == nullptr)
+ return parser.emitError(parser.getCurrentLocation(),
+ "expected 'indexing_maps' attribute");
+ result.addAttribute("indexing_maps", *indexingMapsAttr);
+
+ return parseNamedStructuredOp(parser, result, getNumRegionArgs(),
+ regionBuilder);
+}
+
+void ScaledContractOp::print(OpAsmPrinter &p) {
+ p << " indexing_maps = " << llvm::interleaved_array(getIndexingMaps());
+ printNamedStructuredOp(
+ p, getOperation(), getInputs(), getOutputs(),
+ /*elidedAttrs=*/{"indexing_maps", "operandSegmentSizes"});
+}
+
+LogicalResult ScaledContractOp::verify() {
+ int iterationSpaceDims = -1;
+ // Map iter space dims to #occurrences in inputs' and output's affine_maps:
+ // e.g., inOccurrences[0] will hold #times that dim (with index) 0 is used to
+ // access an input operand (so occurrence count can be at most 2) and
+ // outOccurrences[1] will indicate whether dim 1 occurred in the output, etc.
+ SmallVector<size_t> inOccurrences;
+ SmallVector<size_t> outOccurrences;
+
+ // Validate inputs and contraction semantics.
+ SmallVector<AffineMap, 5> maps = getIndexingMapsArray();
+ SmallVector<Type, 5> types = llvm::to_vector(getOperandTypes());
+ for (auto &&[affineMap, operandType, isInput] :
+ llvm::zip(SmallVector<AffineMap>{maps[0], maps[2], maps[4]},
+ SmallVector<Type>{types[0], types[2], types[4]},
+ SmallVector<bool>{true, true, false})) {
+ if (failed(checkContractionAffineMapAndType(
+ affineMap, operandType, isInput, iterationSpaceDims, inOccurrences,
+ outOccurrences, [&]() { return emitError(); })))
+ return failure(); // NB: Validation helper emits relevant error.
+ }
+
+ if (failed(verifyContractionDims(static_cast<size_t>(iterationSpaceDims),
+ inOccurrences, outOccurrences,
+ [&]() { return emitError(); })))
+ return failure(); // NB: Validation helper emits relevant error.
+
+ // Validate scales and scaling semantics.
+ auto checkScaleAffineMapAndType = [&](AffineMap affineMap, Type operandType,
+ bool isInput) -> LogicalResult {
+ // If scale's map is not a projected permutation, then it must follow
+ // specific scaling scheme semantics.
+ if (!affineMap.isProjectedPermutation()) {
+ if (affineMap.getNumSymbols() > 0)
+ return emitError("scale affine_map must not contain symbols");
+ if (affineMap.getNumResults() > affineMap.getNumInputs())
+ return emitError(
+ "scale affine_map must not have more results than inputs");
+
+ SmallVector<bool, 8> seen(affineMap.getNumInputs(), false);
+ // Allow, at most, only one instance of each input dimension in the result
+ // expressions.
+ for (auto expr : affineMap.getResults()) {
+ AffineDimExpr dim = nullptr;
+ if (isa<AffineDimExpr>(expr)) {
+ // Scaling over whole dimesion.
+ dim = dyn_cast<AffineDimExpr>(expr);
+ } else if (auto binExpr = dyn_cast<AffineBinaryOpExpr>(expr)) {
+ // Scaling over a part of the dimension.
+ // Note: Currently support limited to block scaling i.e.,
+ // one scale per a fixed number of contiguous scalar elements
+ // in a given dimension.
+ if (binExpr.getKind() != AffineExprKind::FloorDiv)
+ return emitError(
+ "only block scale with floordiv is supported for now");
+ auto scaleDim = dyn_cast<AffineDimExpr>(binExpr.getLHS());
+ if (!scaleDim)
+ return emitError("block scale LHS must be dim");
+ auto scaleFactor = dyn_cast<AffineConstantExpr>(binExpr.getRHS());
+ if (!scaleFactor)
+ return emitError("block scale RHS must be constant");
+ if (scaleFactor.getValue() <= 0)
+ return emitError("block scale factor must be positive");
+ dim = scaleDim;
+ } else {
+ return emitError("unsupported scaling variant");
+ }
+
+ if (!dim)
+ return emitError("invalid scale affine_map result expression");
+ if (seen[dim.getPosition()])
+ return emitError(
+ "scale affine_map must not have duplicate result dimensions");
+ seen[dim.getPosition()] = true;
+ }
+ }
+
+ if (auto shapedType = dyn_cast<ShapedType>(operandType)) {
+ if (affineMap.getNumResults() != shapedType.getRank())
+ return emitError(
+ "scale ranks of shaped operand and results of corresponding "
+ "affine_map differ");
+ } else if (affineMap.getNumResults() != 0) {
+ return emitError(
+ "scale affine_map specifies shaped access while operand has "
+ "non-shaped type");
+ }
+
+ return success();
+ };
+
+ // Validate scales' maps.
+ for (auto &&[affineMap, operandType] :
+ llvm::zip(SmallVector<AffineMap>{maps[1], maps[3]},
+ SmallVector<Type>{types[1], types[3]})) {
+ if (failed(checkScaleAffineMapAndType(affineMap, operandType,
+ /*isInput=*/true)))
+ return failure(); // NB: Validation helper emits relevant error.
+ }
+
+ // Cross-validate maps of operand and their scale.
+ for (auto &&[inputMap, inputType, scaleMap, scaleType] :
+ llvm::zip(SmallVector<AffineMap>{maps[0], maps[2]},
+ SmallVector<Type>{types[0], types[2]},
+ SmallVector<AffineMap>{maps[1], maps[3]},
+ SmallVector<Type>{types[1], types[3]})) {
+ if (inputMap.getNumResults() < scaleMap.getNumResults())
+ return emitError("scale must have at most the same rank as input");
+ if (scaleMap.getNumResults() == 0)
+ continue;
+
+ auto inputShape = dyn_cast<ShapedType>(inputType).getShape();
+ auto scaleShape = dyn_cast<ShapedType>(scaleType).getShape();
+
+ // Each scale dim must reference a dim present in the input map:
+ // - a missing scale dim indicates scaling over the whole input dimension
+ // - a scale dim with floordiv indicates reusing the scaling factor over
+ // parts of the input dimension; the factor must match the ratio of
+ // input dim and scale dim sizes.
+ for (auto [scaleIdx, scaleExpr] : llvm::enumerate(scaleMap.getResults())) {
+ AffineDimExpr scaleDimExpr = nullptr;
+ std::optional<int64_t> scaleFactor;
+ if (auto dimExpr = dyn_cast<AffineDimExpr>(scaleExpr)) {
+ // Scaling over the whole dimension.
+ scaleDimExpr = dimExpr;
+ } else if (auto scaleBinExpr = dyn_cast<AffineBinaryOpExpr>(scaleExpr)) {
+ // Block scaling over a part of the dimension.
+ assert(scaleBinExpr.getKind() == AffineExprKind::FloorDiv &&
+ "only floordiv is supported for now");
+ auto scaleDim = dyn_cast<AffineDimExpr>(scaleBinExpr.getLHS());
+ assert(scaleDim && "block scale LHS is a dim expression");
+ scaleDimExpr = scaleDim;
+ scaleFactor =
+ dyn_cast<AffineConstantExpr>(scaleBinExpr.getRHS()).getValue();
+ } else {
+ llvm_unreachable("unknown scale expression");
+ }
+ assert(scaleDimExpr && "failed to find scale dim expression");
+
+ std::optional<unsigned> inputIdx =
+ inputMap.getResultPosition(scaleDimExpr);
+ if (!inputIdx)
+ return emitError(
+ "scale map must contain corresponding input dimensions only");
+
+ // Validate block scaling factor for static shapes.
+ // For dynamic shapes, it is assumed that all sizes are correct.
+ if (scaleFactor && inputShape[*inputIdx] != ShapedType::kDynamic &&
+ scaleShape[scaleIdx] != ShapedType::kDynamic &&
+ llvm::divideCeilSigned(inputShape[*inputIdx], *scaleFactor) !=
+ static_cast<int64_t>(scaleShape[scaleIdx])) {
+ return emitError() << "Invalid scale shape at dim " << *inputIdx
+ << ", expected "
+ << llvm::divideCeilSigned(inputShape[*inputIdx],
+ *scaleFactor)
+ << " but got " << scaleShape[scaleIdx];
+ }
+ }
+ }
+
+ return success();
+}
+
+LogicalResult ScaledContractOp::fold(FoldAdaptor,
+ SmallVectorImpl<OpFoldResult> &) {
+ return memref::foldMemRefCast(*this);
+}
+
+void ScaledContractOp::getEffects(
+ SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+ &effects) {
+ if (hasPureTensorSemantics())
+ return;
+ getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()));
+}
+
+Speculation::Speculatability ScaledContractOp::getSpeculatability() {
+ return getGenericSpeculatabilityImpl(cast<LinalgOp>(getOperation()));
+}
+
} // namespace linalg
} // namespace mlir
diff --git a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
index 4eaa7bf0233c6..22e1a9f12341b 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
@@ -87,6 +87,36 @@ static LogicalResult inlinePayload(OpBuilder &b, LinalgOp linalgOp,
//===----------------------------------------------------------------------===//
namespace {
+
+/// Helper implementing the generalized LinalgOp tiled implementation.
+static FailureOr<TilingResult>
+getTiledImplForLinalgOp(Operation *op, OpBuilder &b,
+ ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes) {
+ Location loc = op->getLoc();
+ LinalgOp linalgOp = cast<LinalgOp>(op);
+ SmallVector<Value> valuesToTile = linalgOp->getOperands();
+ SmallVector<Value> tiledOperands =
+ makeTiledShapes(b, loc, linalgOp, valuesToTile, offsets, sizes, {}, true);
+ SmallVector<Operation *> generatedSlices = llvm::map_to_vector(
+ llvm::make_filter_range(
+ tiledOperands,
+ [](Value v) -> bool {
+ return isa_and_nonnull<tensor::ExtractSliceOp, memref::SubViewOp>(
+ v.getDefiningOp());
+ }),
+ [](Value v) -> Operation * { return v.getDefiningOp(); });
+
+ SmallVector<Type> resultTensorTypes =
+ getTensorOutputTypes(linalgOp, tiledOperands);
+
+ Operation *tiledOp = clone(b, linalgOp, resultTensorTypes, tiledOperands);
+ offsetIndices(b, cast<LinalgOp>(tiledOp), offsets);
+
+ return TilingResult{
+ {tiledOp}, SmallVector<Value>(tiledOp->getResults()), generatedSlices};
+}
+
/// External model implementation of TilingInterface for LinalgOps. An external
/// model implementation is used for now till the use of `TilingInterface` is
/// on-par with the current Linalg tiling + fusion patterns. Once it is
@@ -124,30 +154,7 @@ struct LinalgOpTilingInterface
getTiledImplementation(Operation *op, OpBuilder &b,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
- // Leave the `sizeBounds` value empty. That is only needed when the `sizes`
- // specified could lead to out of bounds accesses.
- Location loc = op->getLoc();
- LinalgOp linalgOp = cast<LinalgOp>(op);
- SmallVector<Value> valuesToTile = linalgOp->getOperands();
- SmallVector<Value> tiledOperands = makeTiledShapes(
- b, loc, linalgOp, valuesToTile, offsets, sizes, {}, true);
- SmallVector<Operation *> generatedSlices = llvm::map_to_vector(
- llvm::make_filter_range(
- tiledOperands,
- [](Value v) -> bool {
- return isa_and_nonnull<tensor::ExtractSliceOp, memref::SubViewOp>(
- v.getDefiningOp());
- }),
- [](Value v) -> Operation * { return v.getDefiningOp(); });
-
- SmallVector<Type> resultTensorTypes =
- getTensorOutputTypes(linalgOp, tiledOperands);
-
- Operation *tiledOp = clone(b, linalgOp, resultTensorTypes, tiledOperands);
- offsetIndices(b, cast<LinalgOp>(tiledOp), offsets);
-
- return TilingResult{
- {tiledOp}, SmallVector<Value>(tiledOp->getResults()), generatedSlices};
+ return getTiledImplForLinalgOp(op, b, offsets, sizes);
}
/// Utility to fetch the offsets and sizes when applied as per the indexing
@@ -386,6 +393,53 @@ struct LinalgOpTilingInterface
}
};
+/// Validate tile sizes against `ScaledContractOp` constraints.
+static LogicalResult
+verifyScaledContractTileSizes(linalg::ScaledContractOp scaledContractOp,
+ ArrayRef<OpFoldResult> sizes) {
+ SmallVector<AffineMap> maps = scaledContractOp.getIndexingMapsArray();
+ for (AffineMap scaleMap : {maps[1], maps[3]}) {
+ for (AffineExpr expr : scaleMap.getResults()) {
+ auto binExpr = dyn_cast<AffineBinaryOpExpr>(expr);
+ if (!binExpr)
+ continue;
+ assert(binExpr.getKind() == AffineExprKind::FloorDiv &&
+ "expected floordiv in scale expression");
+ // For each dimension that uses block-scaling, the corresponding tile size
+ // must be divisible by the scale factor to prevent tiles with dynamic
+ // scale ranges. These cannot be currently expressed by indexing maps.
+ unsigned dimPos = cast<AffineDimExpr>(binExpr.getLHS()).getPosition();
+ int64_t scaleFactor =
+ cast<AffineConstantExpr>(binExpr.getRHS()).getValue();
+ FailureOr<int64_t> tileSize =
+ ValueBoundsConstraintSet::computeConstantBound(
+ presburger::BoundType::UB, sizes[dimPos],
+ /*stopCondition=*/nullptr, ValueBoundsOptions{/*closedUB=*/true});
+ if (succeeded(tileSize) &&
+ !(*tileSize % scaleFactor == 0 || scaleFactor % *tileSize == 0)) {
+ return scaledContractOp.emitOpError()
+ << "tile size " << *tileSize << " for dim " << dimPos
+ << " must divide or be divisible by scale factor "
+ << scaleFactor;
+ }
+ }
+ }
+ return success();
+}
+
+/// Specialization for ScaledContractOp with extra validation to ensure scaling
+/// schemes are compatible with tile sizes.
+template <>
+FailureOr<TilingResult>
+LinalgOpTilingInterface<linalg::ScaledContractOp>::getTiledImplementation(
+ Operation *op, OpBuilder &b, ArrayRef<OpFoldResult> offsets,
+ ArrayRef<OpFoldResult> sizes) const {
+ auto scaledContractOp = cast<linalg::ScaledContractOp>(op);
+ if (failed(verifyScaledContractTileSizes(scaledContractOp, sizes)))
+ return failure();
+ return getTiledImplForLinalgOp(op, b, offsets, sizes);
+}
+
//===----------------------------------------------------------------------===//
// External Model for implementing `PartialReductionInterface` for `LinalgOp`s.
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Linalg/generalize-named-ops.mlir b/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
index e346bee901f1d..b6ec5c4d1f29b 100644
--- a/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/generalize-named-ops.mlir
@@ -1279,3 +1279,52 @@ func.func @preserve_discardable_attrs(%A : tensor<16x8xf32>,
// CHECK-LABEL: func @preserve_discardable_attrs
// CHECK: linalg.generic
// CHECK-SAME: attrs = {another_attr = 42 : i64, my_custom_attr = "preserved"}
+
+// -----
+
+func.func @generalize_scaled_contract(
+ %A: tensor<8x16xf8E5M2>, %sA: tensor<8xf8E8M0FNU>,
+ %B: tensor<4x16xf8E5M2>, %sB: tensor<4xf8E8M0FNU>,
+ %C: tensor<8x4xf32>) -> tensor<8x4xf32> {
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB
+ : tensor<8x16xf8E5M2>, tensor<8xf8E8M0FNU>,
+ tensor<4x16xf8E5M2>, tensor<4xf8E8M0FNU>)
+ outs(%C : tensor<8x4xf32>) -> tensor<8x4xf32>
+ return %D : tensor<8x4xf32>
+}
+
+// CHECK-DAG: #[[$A_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$SA_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0)>
+// CHECK-DAG: #[[$B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK-DAG: #[[$SB_MAP:.+]] = affine_map<(d0, d1, d2) -> (d1)>
+// CHECK-DAG: #[[$C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: func @generalize_scaled_contract
+// CHECK-SAME: %[[A:.+]]: tensor<8x16xf8E5M2>
+// CHECK-SAME: %[[SA:.+]]: tensor<8xf8E8M0FNU>
+// CHECK-SAME: %[[B:.+]]: tensor<4x16xf8E5M2>
+// CHECK-SAME: %[[SB:.+]]: tensor<4xf8E8M0FNU>
+// CHECK-SAME: %[[C:.+]]: tensor<8x4xf32>
+
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$A_MAP]], #[[$SA_MAP]], #[[$B_MAP]], #[[$SB_MAP]], #[[$C_MAP]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK-SAME: ins(%[[A]], %[[SA]], %[[B]], %[[SB]]
+// CHECK-SAME: : tensor<8x16xf8E5M2>, tensor<8xf8E8M0FNU>, tensor<4x16xf8E5M2>, tensor<4xf8E8M0FNU>)
+// CHECK-SAME: outs(%[[C]] : tensor<8x4xf32>)
+// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f8E5M2, %[[SA_ARG:.+]]: f8E8M0FNU, %[[B_ARG:.+]]: f8E5M2, %[[SB_ARG:.+]]: f8E8M0FNU, %[[C_ARG:.+]]: f32)
+// CHECK: %[[A_EXT:.+]] = arith.extf %[[A_ARG]] : f8E5M2 to f32
+// CHECK: %[[SA_EXT:.+]] = arith.extf %[[SA_ARG]] : f8E8M0FNU to f32
+// CHECK: %[[SLHS:.+]] = arith.mulf %[[A_EXT]], %[[SA_EXT]] : f32
+// CHECK: %[[B_EXT:.+]] = arith.extf %[[B_ARG]] : f8E5M2 to f32
+// CHECK: %[[SB_EXT:.+]] = arith.extf %[[SB_ARG]] : f8E8M0FNU to f32
+// CHECK: %[[SRHS:.+]] = arith.mulf %[[B_EXT]], %[[SB_EXT]] : f32
+// CHECK: %[[PROD:.+]] = arith.mulf %[[SLHS]], %[[SRHS]] : f32
+// CHECK: %[[ACC:.+]] = arith.addf %[[C_ARG]], %[[PROD]] : f32
+// CHECK: linalg.yield %[[ACC]] : f32
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index a446cfcc4eec1..b71e9f819e240 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/mlir/test/Dialect/Linalg/invalid.mlir
@@ -742,6 +742,193 @@ func.func @unused_iteration_space_dim_contraction(
// -----
+func.func @scaled_contract_scale_with_symbols(
+ %A: tensor<256x512xi8>, %sA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>, %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @+1 {{scale affine_map must not contain symbols}}
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k)[s0] -> (m floordiv 32, s0)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+func.func @scaled_contract_scale_too_many_results(
+ %A: tensor<256x512xi8>, %sA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>, %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @+1 {{scale affine_map must not have more results than inputs}}
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128, n, m)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+func.func @scaled_contract_scale_non_floordiv(
+ %A: tensor<256x512xi8>, %sA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>, %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @+1 {{only block scale with floordiv is supported for now}}
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k mod 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+func.func @scaled_contract_block_scale_lhs_not_dim(
+ %A: tensor<256x512xi8>, %sA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>, %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @+1 {{block scale LHS must be dim}}
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> ((m + n) floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+func.func @scaled_contract_block_scale_nonpositive_factor(
+ %A: tensor<256x512xi8>, %sA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>, %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @+1 {{block scale factor must be positive}}
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv -32)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+func.func @scaled_contract_scale_duplicate_dims(
+ %A: tensor<256x512xi8>, %sA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>, %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @+1 {{scale affine_map must not have duplicate result dimensions}}
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, m floordiv 16)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+func.func @scaled_contract_scale_rank_mismatch(
+ %A: tensor<256x512xi8>, %sA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>, %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @+1 {{scale ranks of shaped operand and results of corresponding affine_map differ}}
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+func.func @scaled_contract_scale_non_shaped_operand(
+ %A: tensor<256x512xi8>, %sA: f8E8M0FNU,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>, %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @+1 {{scale affine_map specifies shaped access while operand has non-shaped type}}
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, f8E8M0FNU, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+func.func @scaled_contract_scale_rank_exceeds_input(
+ %A: tensor<256x512xi8>, %sA: tensor<8x4x2xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>, %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @+1 {{scale must have at most the same rank as input}}
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m, k, n)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<8x4x2xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+func.func @scaled_contract_invalid_block_scale_shape(
+ %A: tensor<256x512xi8>, %sA: tensor<8x8xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>, %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @+1 {{Invalid scale shape at dim 1, expected 4 but got 8}}
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<8x8xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+func.func @scaled_contract_scale_dim_not_in_input(
+ %A: tensor<256x512xi8>, %sA: tensor<128xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>, %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @+1 {{scale map must contain corresponding input dimensions only}}
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<128xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
func.func @invalid_static_2d_conv(%input : memref<1x3x4x2xf32>, %filter: memref<3x2x2x1xf32>, %output: memref<1x2x3x1xf32>) {
// expected-error @+1 {{inferred input/output operand #0 has shape's dimension #1 to be greater than or equal to 4, but found 3}}
linalg.conv_2d_nhwc_hwcf
diff --git a/mlir/test/Dialect/Linalg/loops.mlir b/mlir/test/Dialect/Linalg/loops.mlir
index efe8010cffc91..da8beaf4b2212 100644
--- a/mlir/test/Dialect/Linalg/loops.mlir
+++ b/mlir/test/Dialect/Linalg/loops.mlir
@@ -674,6 +674,64 @@ func.func @batch_reduce_matmul_as_contract(
// CHECKPARALLEL: %[[res:.*]] = arith.addf %[[vc]], %[[inc]] : f32
// CHECKPARALLEL: store %[[res]], %[[mC]][%[[m]], %[[n]]] : memref<?x?xf32>
+func.func @scaled_contract_as_loops(
+ %A: memref<?x?xf32>, %sA: memref<?xf32>,
+ %B: memref<?x?xf32>, %sB: memref<?xf32>, %C: memref<?x?xf32>) {
+ linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : memref<?x?xf32>, memref<?xf32>, memref<?x?xf32>, memref<?xf32>)
+ outs(%C : memref<?x?xf32>)
+ return
+}
+// CHECK-LABEL: @scaled_contract_as_loops
+// CHECK-SAME: %[[mA:[a-zA-Z0-9]+]]: memref<?x?xf32>
+// CHECK-SAME: %[[msA:[a-zA-Z0-9]+]]: memref<?xf32>
+// CHECK-SAME: %[[mB:[a-zA-Z0-9]+]]: memref<?x?xf32>
+// CHECK-SAME: %[[msB:[a-zA-Z0-9]+]]: memref<?xf32>
+// CHECK-SAME: %[[mC:[a-zA-Z0-9]+]]: memref<?x?xf32>
+// CHECK: %[[M:.*]] = memref.dim %[[mA]], %c0 : memref<?x?xf32>
+// CHECK: %[[K:.*]] = memref.dim %[[mA]], %c1 : memref<?x?xf32>
+// CHECK: %[[N:.*]] = memref.dim %[[mB]], %c0 : memref<?x?xf32>
+// CHECK: scf.for %[[m:.*]] = %{{.*}} to %[[M]]
+// CHECK: scf.for %[[n:.*]] = %{{.*}} to %[[N]]
+// CHECK: scf.for %[[k:.*]] = %{{.*}} to %[[K]]
+// CHECK: %[[va:.*]] = memref.load %[[mA]][%[[m]], %[[k]]] : memref<?x?xf32>
+// CHECK: %[[vsa:.*]] = memref.load %[[msA]][%[[m]]] : memref<?xf32>
+// CHECK: %[[vb:.*]] = memref.load %[[mB]][%[[n]], %[[k]]] : memref<?x?xf32>
+// CHECK: %[[vsb:.*]] = memref.load %[[msB]][%[[n]]] : memref<?xf32>
+// CHECK: %[[vc:.*]] = memref.load %[[mC]][%[[m]], %[[n]]] : memref<?x?xf32>
+// CHECK: %[[sa:.*]] = arith.mulf %[[va]], %[[vsa]] : f32
+// CHECK: %[[sb:.*]] = arith.mulf %[[vb]], %[[vsb]] : f32
+// CHECK: %[[prod:.*]] = arith.mulf %[[sa]], %[[sb]] : f32
+// CHECK: %[[res:.*]] = arith.addf %[[vc]], %[[prod]] : f32
+// CHECK: store %[[res]], %[[mC]][%[[m]], %[[n]]] : memref<?x?xf32>
+
+// CHECKPARALLEL-LABEL: @scaled_contract_as_loops
+// CHECKPARALLEL-SAME: %[[mA:[a-zA-Z0-9]+]]: memref<?x?xf32>
+// CHECKPARALLEL-SAME: %[[msA:[a-zA-Z0-9]+]]: memref<?xf32>
+// CHECKPARALLEL-SAME: %[[mB:[a-zA-Z0-9]+]]: memref<?x?xf32>
+// CHECKPARALLEL-SAME: %[[msB:[a-zA-Z0-9]+]]: memref<?xf32>
+// CHECKPARALLEL-SAME: %[[mC:[a-zA-Z0-9]+]]: memref<?x?xf32>
+// CHECKPARALLEL: %[[M:.*]] = memref.dim %[[mA]], %c0 : memref<?x?xf32>
+// CHECKPARALLEL: %[[K:.*]] = memref.dim %[[mA]], %c1 : memref<?x?xf32>
+// CHECKPARALLEL: %[[N:.*]] = memref.dim %[[mB]], %c0 : memref<?x?xf32>
+// CHECKPARALLEL: scf.parallel (%[[m:.*]], %[[n:.*]]) = ({{.*}}) to (%[[M]], %[[N]]) step ({{.*}}) {
+// CHECKPARALLEL: scf.for %[[k:.*]] = %{{.*}} to %[[K]]
+// CHECKPARALLEL: %[[va:.*]] = memref.load %[[mA]][%[[m]], %[[k]]] : memref<?x?xf32>
+// CHECKPARALLEL: %[[vsa:.*]] = memref.load %[[msA]][%[[m]]] : memref<?xf32>
+// CHECKPARALLEL: %[[vb:.*]] = memref.load %[[mB]][%[[n]], %[[k]]] : memref<?x?xf32>
+// CHECKPARALLEL: %[[vsb:.*]] = memref.load %[[msB]][%[[n]]] : memref<?xf32>
+// CHECKPARALLEL: %[[vc:.*]] = memref.load %[[mC]][%[[m]], %[[n]]] : memref<?x?xf32>
+// CHECKPARALLEL: %[[sa:.*]] = arith.mulf %[[va]], %[[vsa]] : f32
+// CHECKPARALLEL: %[[sb:.*]] = arith.mulf %[[vb]], %[[vsb]] : f32
+// CHECKPARALLEL: %[[prod:.*]] = arith.mulf %[[sa]], %[[sb]] : f32
+// CHECKPARALLEL: %[[res:.*]] = arith.addf %[[vc]], %[[prod]] : f32
+// CHECKPARALLEL: store %[[res]], %[[mC]][%[[m]], %[[n]]] : memref<?x?xf32>
+
func.func @named_batch_matmul(%A: memref<?x?x?xf32>, %B: memref<?x?x?xf32>, %C: memref<?x?x?xf32>) {
linalg.batch_matmul ins(%A, %B : memref<?x?x?xf32>, memref<?x?x?xf32>)
outs(%C : memref<?x?x?xf32>)
diff --git a/mlir/test/Dialect/Linalg/named-ops.mlir b/mlir/test/Dialect/Linalg/named-ops.mlir
index 8068c23a4a0fd..b9cb8e012e651 100644
--- a/mlir/test/Dialect/Linalg/named-ops.mlir
+++ b/mlir/test/Dialect/Linalg/named-ops.mlir
@@ -1978,6 +1978,168 @@ func.func @contract_matmul_bcast_b_transpose_a(%A: memref<5x3xf32>, %B: memref<5
// -----
+// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK: #[[$ACCESS_SCALE_A:.+]] = affine_map<(d0, d1, d2) -> (d0 floordiv 32, d2 floordiv 128)>
+// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK: #[[$ACCESS_SCALE_B:.+]] = affine_map<(d0, d1, d2) -> (d1)>
+// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL: func @scaled_contract
+func.func @scaled_contract(
+ %A: tensor<256x512xi8>, %sA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>,
+ %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+// CHECK: linalg.scaled_contract
+// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_SCALE_A]], #[[$ACCESS_B]], #[[$ACCESS_SCALE_B]], #[[$ACCESS_C]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}} : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+// CHECK-SAME: outs(%{{.+}} : tensor<256x128xf32>) -> tensor<256x128xf32>
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK: #[[$ACCESS_SCALE:.+]] = affine_map<(d0, d1, d2) -> ()>
+// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL: func @scaled_contract_tensor_scale
+func.func @scaled_contract_tensor_scale(
+ %A: memref<256x512xi8>, %sA: memref<f8E8M0FNU>,
+ %B: memref<128x512xi8>, %sB: memref<f8E8M0FNU>,
+ %C: memref<256x128xf32>) {
+// CHECK: linalg.scaled_contract
+// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_SCALE]], #[[$ACCESS_B]], #[[$ACCESS_SCALE]], #[[$ACCESS_C]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}} : memref<256x512xi8>, memref<f8E8M0FNU>, memref<128x512xi8>, memref<f8E8M0FNU>)
+// CHECK-SAME: outs(%{{.+}} : memref<256x128xf32>)
+ linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> ()>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> ()>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : memref<256x512xi8>, memref<f8E8M0FNU>, memref<128x512xi8>, memref<f8E8M0FNU>)
+ outs(%C : memref<256x128xf32>)
+ return
+}
+
+// -----
+
+// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK: #[[$ACCESS_SCALE_A:.+]] = affine_map<(d0, d1, d2) -> (d0)>
+// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK: #[[$ACCESS_SCALE_B:.+]] = affine_map<(d0, d1, d2) -> (d1)>
+// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL: func @scaled_contract_dim_scale
+func.func @scaled_contract_dim_scale(
+ %A: tensor<256x512xi8>, %sA: tensor<256xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>,
+ %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+// CHECK: linalg.scaled_contract
+// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_SCALE_A]], #[[$ACCESS_B]], #[[$ACCESS_SCALE_B]], #[[$ACCESS_C]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}} : tensor<256x512xi8>, tensor<256xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+// CHECK-SAME: outs(%{{.+}} : tensor<256x128xf32>) -> tensor<256x128xf32>
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<256xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK: #[[$ACCESS_SCALE_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2 floordiv 32)>
+// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK: #[[$ACCESS_SCALE_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2 floordiv 32)>
+// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL: func @scaled_contract_block_scale
+func.func @scaled_contract_block_scale(
+ %A: tensor<?x?xi8>, %sA: tensor<?x?xf8E8M0FNU>,
+ %B: tensor<?x?xi8>, %sB: tensor<?x?xf8E8M0FNU>,
+ %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
+// CHECK: linalg.scaled_contract
+// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_SCALE_A]], #[[$ACCESS_B]], #[[$ACCESS_SCALE_B]], #[[$ACCESS_C]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}} : tensor<?x?xi8>, tensor<?x?xf8E8M0FNU>, tensor<?x?xi8>, tensor<?x?xf8E8M0FNU>)
+// CHECK-SAME: outs(%{{.+}} : tensor<?x?xf32>) -> tensor<?x?xf32>
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m, k floordiv 32)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n, k floordiv 32)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<?x?xi8>, tensor<?x?xf8E8M0FNU>, tensor<?x?xi8>, tensor<?x?xf8E8M0FNU>)
+ outs(%C : tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %D : tensor<?x?xf32>
+}
+
+// -----
+
+// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK: #[[$ACCESS_SCALE_A:.+]] = affine_map<(d0, d1, d2) -> (d0 floordiv 32, d2 floordiv 128)>
+// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK: #[[$ACCESS_SCALE_B:.+]] = affine_map<(d0, d1, d2) -> (d1)>
+// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL: func @scaled_contract_cast_unsigned
+func.func @scaled_contract_cast_unsigned(
+ %A: tensor<256x512xi8>, %sA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>,
+ %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+// CHECK: linalg.scaled_contract
+// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_SCALE_A]], #[[$ACCESS_B]], #[[$ACCESS_SCALE_B]], #[[$ACCESS_C]]]
+// CHECK-SAME: {cast = #linalg.type_fn<cast_unsigned>}
+// CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}} : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+// CHECK-SAME: outs(%{{.+}} : tensor<256x128xf32>) -> tensor<256x128xf32>
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) {cast = #linalg.type_fn<cast_unsigned>} -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
+// CHECK: #[[$ACCESS_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK: #[[$ACCESS_SCALE_A:.+]] = affine_map<(d0, d1, d2) -> (d2 floordiv 128, d0 floordiv 32)>
+// CHECK: #[[$ACCESS_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK: #[[$ACCESS_SCALE_B:.+]] = affine_map<(d0, d1, d2) -> (d1)>
+// CHECK: #[[$ACCESS_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL: func @scaled_contract_permuted_scale
+func.func @scaled_contract_permuted_scale(
+ %A: tensor<256x512xi8>, %sA: tensor<4x8xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %sB: tensor<128xf8E8M0FNU>,
+ %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+// CHECK: linalg.scaled_contract
+// CHECK-SAME: indexing_maps = [#[[$ACCESS_A]], #[[$ACCESS_SCALE_A]], #[[$ACCESS_B]], #[[$ACCESS_SCALE_B]], #[[$ACCESS_C]]]
+// CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}} : tensor<256x512xi8>, tensor<4x8xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+// CHECK-SAME: outs(%{{.+}} : tensor<256x128xf32>) -> tensor<256x128xf32>
+ %D = linalg.scaled_contract
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (k floordiv 128, m floordiv 32)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %sA, %B, %sB : tensor<256x512xi8>, tensor<4x8xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+
+// -----
+
// CHECK-LABEL: func @mmt4d
func.func @mmt4d(%A: tensor<10x32x8x1xf32>, %B: tensor<80x32x4x1xf32>, %C: tensor<10x80x8x4xf32>) -> tensor<10x80x8x4xf32> {
// CHECK: %{{.+}} = linalg.mmt4d
diff --git a/mlir/test/Dialect/Linalg/roundtrip.mlir b/mlir/test/Dialect/Linalg/roundtrip.mlir
index bfb92c3289a49..fc39da9499d6f 100644
--- a/mlir/test/Dialect/Linalg/roundtrip.mlir
+++ b/mlir/test/Dialect/Linalg/roundtrip.mlir
@@ -308,6 +308,34 @@ func.func @named_ops(%a3: memref<?x?x?xf32>, %b3: memref<?x?x?xf32>, %c3: memref
// -----
+#scaledA = affine_map<(m, n, k) -> (m, k)>
+#scaledScaleA = affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>
+#scaledB = affine_map<(m, n, k) -> (n, k)>
+#scaledScaleB = affine_map<(m, n, k) -> (n)>
+#scaledC = affine_map<(m, n, k) -> (m, n)>
+func.func @scaled_contract_named_op(
+ %a: memref<256x512xi8>, %sa: memref<8x4xf8E8M0FNU>, %b: memref<128x512xi8>,
+ %sb: memref<128xf8E8M0FNU>, %c: memref<256x128xf32>,
+ %ta: tensor<256x512xi8>, %tsa: tensor<8x4xf8E8M0FNU>, %tb: tensor<128x512xi8>,
+ %tsb: tensor<128xf8E8M0FNU>, %tc: tensor<256x128xf32>)
+ -> tensor<256x128xf32>
+{
+ linalg.scaled_contract
+ indexing_maps = [#scaledA, #scaledScaleA, #scaledB, #scaledScaleB, #scaledC]
+ ins(%a, %sa, %b, %sb : memref<256x512xi8>, memref<8x4xf8E8M0FNU>, memref<128x512xi8>, memref<128xf8E8M0FNU>)
+ outs(%c : memref<256x128xf32>)
+ %res = linalg.scaled_contract
+ indexing_maps = [#scaledA, #scaledScaleA, #scaledB, #scaledScaleB, #scaledC]
+ ins(%ta, %tsa, %tb, %tsb : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%tc : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %res : tensor<256x128xf32>
+}
+// CHECK-LABEL: func @scaled_contract_named_op
+// CHECK: linalg.scaled_contract
+// CHECK: linalg.scaled_contract
+
+// -----
+
func.func @fill_tensor(%arg0 : index, %arg1 : index, %arg2 : f32) -> tensor<?x?xf32> {
%0 = tensor.empty(%arg0, %arg1) : tensor<?x?xf32>
%1 = linalg.fill ins(%arg2 : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
diff --git a/mlir/test/Dialect/Linalg/transform-op-fuse.mlir b/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
index dab8491708104..e8ca29b5fa6cd 100644
--- a/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-fuse.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s --transform-interpreter --split-input-file -canonicalize | FileCheck %s
+// RUN: mlir-opt %s --transform-interpreter --split-input-file -canonicalize -verify-diagnostics | FileCheck %s
// CHECK-LABEL: func.func @fuse_unary
func.func @fuse_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
@@ -667,3 +667,178 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+// Tiling and fusing `linalg.scaled_contract` must respect the block-scaling
+// scheme encoded in the scale indexing maps: along every block-scaled iteration
+// dimension the tile size must divide or be divisible by the block factor.
+
+// CHECK-DAG: #[[$MAP_M:.+]] = affine_map<(d0) -> (d0 floordiv 32)>
+// CHECK-LABEL: func.func @fuse_scaled_contract_with_producer(
+// CHECK-SAME: %[[A:[a-zA-Z0-9]+]]: tensor<256x512xi8>
+// CHECK-SAME: %[[SCALE_A:[a-zA-Z0-9]+]]: tensor<8x4xf8E8M0FNU>
+// CHECK-SAME: %[[B:[a-zA-Z0-9]+]]: tensor<128x512xi8>
+// CHECK-SAME: %[[SCALE_B:[a-zA-Z0-9]+]]: tensor<128xf8E8M0FNU>
+// CHECK-SAME: %[[INIT:[a-zA-Z0-9]+]]: tensor<256x128xf32>
+func.func @fuse_scaled_contract_with_producer(%A: tensor<256x512xi8>, %scaleA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %scaleB: tensor<128xf8E8M0FNU>,
+ %init: tensor<256x128xf32>) -> tensor<256x128xf32> {
+// CHECK-DAG: %[[CST:.+]] = arith.constant 0.0{{.*}} : f32
+// CHECK: scf.for %[[IV_M:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}} iter_args(%[[INIT_M:.+]] = %[[INIT]])
+// CHECK: scf.for %[[IV_N:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}} iter_args(%[[INIT_N:.+]] = %[[INIT_M]])
+// CHECK-DAG: %[[OFF_M:.+]] = affine.apply #[[$MAP_M]](%[[IV_M]])
+// CHECK-DAG: %[[A_TILE:.+]] = tensor.extract_slice %[[A]][%[[IV_M]], 0] [32, 512] [1, 1]
+// CHECK-DAG: %[[SA_TILE:.+]] = tensor.extract_slice %[[SCALE_A]][%[[OFF_M]], 0] [1, 4] [1, 1]
+// CHECK-DAG: %[[B_TILE:.+]] = tensor.extract_slice %[[B]][%[[IV_N]], 0] [16, 512] [1, 1]
+// CHECK-DAG: %[[SB_TILE:.+]] = tensor.extract_slice %[[SCALE_B]][%[[IV_N]]] [16] [1]
+// CHECK-DAG: %[[C_TILE:.+]] = tensor.extract_slice %[[INIT_N]][%[[IV_M]], %[[IV_N]]] [32, 16] [1, 1]
+// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[C_TILE]] : tensor<32x16xf32>)
+// CHECK: %[[RES:.+]] = linalg.scaled_contract
+// CHECK-SAME: ins(%[[A_TILE]], %[[SA_TILE]], %[[B_TILE]], %[[SB_TILE]] :
+// CHECK-SAME: outs(%[[FILL]] : tensor<32x16xf32>) -> tensor<32x16xf32>
+// CHECK: tensor.insert_slice %[[RES]] into %[[INIT_N]][%[[IV_M]], %[[IV_N]]] [32, 16] [1, 1]
+ %cst = arith.constant 0.0 : f32
+ %C = linalg.fill ins(%cst : f32) outs(%init : tensor<256x128xf32>) -> tensor<256x128xf32>
+ %D = linalg.scaled_contract
+ indexing_maps = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %scaleA, %B, %scaleB
+ : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.scaled_contract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.fuse %0 tile_sizes [32, 16, 0]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-DAG: #[[$MAP_M:.+]] = affine_map<(d0) -> (d0 floordiv 32)>
+// CHECK-LABEL: func.func @fuse_scaled_contract_with_consumer(
+// CHECK-SAME: %[[A:[a-zA-Z0-9]+]]: tensor<256x512xi8>
+// CHECK-SAME: %[[SCALE_A:[a-zA-Z0-9]+]]: tensor<8x4xf8E8M0FNU>
+// CHECK-SAME: %[[B:[a-zA-Z0-9]+]]: tensor<128x512xi8>
+// CHECK-SAME: %[[SCALE_B:[a-zA-Z0-9]+]]: tensor<128xf8E8M0FNU>
+// CHECK-SAME: %[[C:[a-zA-Z0-9]+]]: tensor<256x128xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<256x128xf32>
+func.func @fuse_scaled_contract_with_consumer(%A: tensor<256x512xi8>, %scaleA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %scaleB: tensor<128xf8E8M0FNU>,
+ %C: tensor<256x128xf32>, %out: tensor<256x128xf32>) -> tensor<256x128xf32> {
+// CHECK: scf.for %[[IV_M:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}} iter_args(%[[INIT_M:.+]] = %[[OUT]])
+// CHECK: scf.for %[[IV_N:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}} iter_args(%[[INIT_N:.+]] = %[[INIT_M]])
+// CHECK-DAG: %[[OFF_M:.+]] = affine.apply #[[$MAP_M]](%[[IV_M]])
+// CHECK-DAG: %[[A_TILE:.+]] = tensor.extract_slice %[[A]][%[[IV_M]], 0] [32, 512] [1, 1]
+// CHECK-DAG: %[[SA_TILE:.+]] = tensor.extract_slice %[[SCALE_A]][%[[OFF_M]], 0] [1, 4] [1, 1]
+// CHECK-DAG: %[[B_TILE:.+]] = tensor.extract_slice %[[B]][%[[IV_N]], 0] [16, 512] [1, 1]
+// CHECK-DAG: %[[SB_TILE:.+]] = tensor.extract_slice %[[SCALE_B]][%[[IV_N]]] [16] [1]
+// CHECK-DAG: %[[C_TILE:.+]] = tensor.extract_slice %[[C]][%[[IV_M]], %[[IV_N]]] [32, 16] [1, 1]
+// CHECK: %[[SC:.+]] = linalg.scaled_contract
+// CHECK-SAME: ins(%[[A_TILE]], %[[SA_TILE]], %[[B_TILE]], %[[SB_TILE]] :
+// CHECK-SAME: outs(%[[C_TILE]] : tensor<32x16xf32>) -> tensor<32x16xf32>
+// CHECK: %[[OUT_TILE:.+]] = tensor.extract_slice %[[INIT_N]][%[[IV_M]], %[[IV_N]]] [32, 16] [1, 1]
+// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[SC]] : tensor<32x16xf32>) outs(%[[OUT_TILE]] : tensor<32x16xf32>)
+// CHECK: tensor.insert_slice %[[COPY]] into %[[INIT_N]][%[[IV_M]], %[[IV_N]]] [32, 16] [1, 1]
+ %D = linalg.scaled_contract
+ indexing_maps = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %scaleA, %B, %scaleB
+ : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ %E = linalg.copy ins(%D : tensor<256x128xf32>) outs(%out : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %E : tensor<256x128xf32>
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.fuse %0 tile_sizes [32, 16]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+func.func @negative_fuse_scaled_contract_partial_scaling_block(%A: tensor<256x512xi8>, %scaleA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %scaleB: tensor<128xf8E8M0FNU>,
+ %init: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ %cst = arith.constant 0.0 : f32
+ %C = linalg.fill ins(%cst : f32) outs(%init : tensor<256x128xf32>) -> tensor<256x128xf32>
+ // expected-error @below {{'linalg.scaled_contract' op tile size 48 for dim 0 must divide or be divisible by scale factor 32}}
+ // expected-error @below {{'linalg.scaled_contract' op failed to tile operation}}
+ // expected-error @below {{'linalg.scaled_contract' op failed to generate tiling loops}}
+ %D = linalg.scaled_contract
+ indexing_maps = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %scaleA, %B, %scaleB
+ : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.scaled_contract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.fuse %0 tile_sizes [48, 16, 0]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-LABEL: func.func @negative_fuse_scaled_contract_partial_consumer_scaling_block(
+// CHECK-SAME: %[[A:[a-zA-Z0-9]+]]: tensor<256x512xi8>
+// CHECK-SAME: %[[SCALE_A:[a-zA-Z0-9]+]]: tensor<8x4xf8E8M0FNU>
+// CHECK-SAME: %[[B:[a-zA-Z0-9]+]]: tensor<128x512xi8>
+// CHECK-SAME: %[[SCALE_B:[a-zA-Z0-9]+]]: tensor<128xf8E8M0FNU>
+// CHECK-SAME: %[[C:[a-zA-Z0-9]+]]: tensor<256x128xf32>
+// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<256x128xf32>
+func.func @negative_fuse_scaled_contract_partial_consumer_scaling_block(%A: tensor<256x512xi8>, %scaleA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %scaleB: tensor<128xf8E8M0FNU>,
+ %C: tensor<256x128xf32>, %out: tensor<256x128xf32>) -> tensor<256x128xf32> {
+// CHECK: %[[FULL:.+]] = linalg.scaled_contract
+// CHECK-SAME: outs(%[[C]] : tensor<256x128xf32>) -> tensor<256x128xf32>
+// CHECK: scf.for %[[IV_M:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}}
+// CHECK: scf.for %[[IV_N:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}}
+// CHECK: %[[D_TILE:.+]] = tensor.extract_slice %[[FULL]][%[[IV_M]], %[[IV_N]]] [%{{.+}}, 16] [1, 1]
+// CHECK: %[[COPY:.+]] = linalg.copy ins(%[[D_TILE]] :
+// CHECK: tensor.insert_slice %[[COPY]]
+ // expected-error @below {{'linalg.scaled_contract' op tile size 48 for dim 0 must divide or be divisible by scale factor 32}}
+ %D = linalg.scaled_contract
+ indexing_maps = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %scaleA, %B, %scaleB
+ : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ %E = linalg.copy ins(%D : tensor<256x128xf32>) outs(%out : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %E : tensor<256x128xf32>
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.copy"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:2 = transform.structured.fuse %0 tile_sizes [48, 16]
+ : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/Linalg/transform-op-tile.mlir b/mlir/test/Dialect/Linalg/transform-op-tile.mlir
index 3e442f311a806..a8924b8b61215 100644
--- a/mlir/test/Dialect/Linalg/transform-op-tile.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-tile.mlir
@@ -287,3 +287,199 @@ func.func @tile_linalg_matmul(
-> tensor<128x128xf32>
return %0 : tensor<128x128xf32>
}
+
+// -----
+
+// Tiling `linalg.scaled_contract` must respect the block-scaling scheme encoded in the scale
+// indexing maps: along every iteration dimension that is block-scaled (`d floordiv B`),
+// the tile size must divide or be divisible by the block factor `B`.
+
+// CHECK-LABEL: func.func @tile_scaled_contract_multiple_of_block_scale(
+// CHECK-SAME: %[[A:[a-zA-Z0-9]+]]: tensor<256x512xi8>
+// CHECK-SAME: %[[SCALE_A:[a-zA-Z0-9]+]]: tensor<8x4xf8E8M0FNU>
+// CHECK-SAME: %[[B:[a-zA-Z0-9]+]]: tensor<128x512xi8>
+// CHECK-SAME: %[[SCALE_B:[a-zA-Z0-9]+]]: tensor<128xf8E8M0FNU>
+// CHECK-SAME: %[[C:[a-zA-Z0-9]+]]: tensor<256x128xf32>
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[C32:.+]] = arith.constant 32 : index
+// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
+// CHECK: scf.for %[[IV_M:[a-zA-Z0-9]+]] = %{{.+}} to %{{.+}} step %[[C32]] iter_args(%[[INIT_M:.+]] = %[[C]])
+// CHECK: scf.for %[[IV_N:[a-zA-Z0-9]+]] = %{{.+}} to %{{.+}} step %[[C16]] iter_args(%[[INIT_N:.+]] = %[[INIT_M]])
+// CHECK: scf.for %[[IV_K:[a-zA-Z0-9]+]] = %{{.+}} to %{{.+}} step %[[C128]] iter_args(%[[INIT_K:.+]] = %[[INIT_N]])
+// CHECK-DAG: %[[OFF_M:.+]] = affine.apply affine_map<(d0) -> (d0 floordiv 32)>(%[[IV_M]])
+// CHECK-DAG: %[[OFF_K:.+]] = affine.apply affine_map<(d0) -> (d0 floordiv 128)>(%[[IV_K]])
+// CHECK-DAG: %[[A_TILE:.+]] = tensor.extract_slice %[[A]][%[[IV_M]], %[[IV_K]]] [32, 128] [1, 1]
+// CHECK-DAG: %[[SA_TILE:.+]] = tensor.extract_slice %[[SCALE_A]][%[[OFF_M]], %[[OFF_K]]] [1, 1] [1, 1]
+// CHECK-DAG: %[[B_TILE:.+]] = tensor.extract_slice %[[B]][%[[IV_N]], %[[IV_K]]] [16, 128] [1, 1]
+// CHECK-DAG: %[[SB_TILE:.+]] = tensor.extract_slice %[[SCALE_B]][%[[IV_N]]] [16] [1]
+// CHECK-DAG: %[[C_TILE:.+]] = tensor.extract_slice %[[INIT_K]][%[[IV_M]], %[[IV_N]]] [32, 16] [1, 1]
+// CHECK: %[[RES:.+]] = linalg.scaled_contract
+// CHECK-SAME: ins(%[[A_TILE]], %[[SA_TILE]], %[[B_TILE]], %[[SB_TILE]] :
+// CHECK-SAME: outs(%[[C_TILE]] :
+// CHECK-SAME: -> tensor<32x16xf32>
+// CHECK: tensor.insert_slice %[[RES]] into %[[INIT_K]][%[[IV_M]], %[[IV_N]]] [32, 16] [1, 1]
+func.func @tile_scaled_contract_multiple_of_block_scale(
+ %A: tensor<256x512xi8>, %scaleA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %scaleB: tensor<128xf8E8M0FNU>,
+ %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ %D = linalg.scaled_contract
+ indexing_maps = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %scaleA, %B, %scaleB
+ : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.scaled_contract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:3 = transform.structured.tile_using_for %0 tile_sizes [32, 16, 128] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-LABEL: func.func @tile_scaled_contract_divisor_of_block_scale(
+// CHECK-SAME: %[[A:[a-zA-Z0-9]+]]: tensor<256x512xi8>
+// CHECK-SAME: %[[SCALE_A:[a-zA-Z0-9]+]]: tensor<8x4xf8E8M0FNU>
+// CHECK: scf.for %[[IV_M:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}}
+// CHECK: scf.for %[[IV_N:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}}
+// CHECK: scf.for %[[IV_K:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}}
+// CHECK-DAG: %[[OFF_M:.+]] = affine.apply affine_map<(d0) -> (d0 floordiv 32)>(%[[IV_M]])
+// CHECK-DAG: %[[OFF_K:.+]] = affine.apply affine_map<(d0) -> (d0 floordiv 128)>(%[[IV_K]])
+// CHECK-DAG: %[[A_TILE:.+]] = tensor.extract_slice %[[A]][%[[IV_M]], %[[IV_K]]] [32, 64] [1, 1]
+// CHECK-DAG: %[[SA_TILE:.+]] = tensor.extract_slice %[[SCALE_A]][%[[OFF_M]], %[[OFF_K]]] [1, 1] [1, 1]
+// CHECK: %[[RES:.+]] = linalg.scaled_contract
+// CHECK-SAME: -> tensor<32x16xf32>
+func.func @tile_scaled_contract_divisor_of_block_scale(
+ %A: tensor<256x512xi8>, %scaleA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %scaleB: tensor<128xf8E8M0FNU>,
+ %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ %D = linalg.scaled_contract
+ indexing_maps = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %scaleA, %B, %scaleB
+ : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.scaled_contract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:3 = transform.structured.tile_using_for %0 tile_sizes [32, 16, 64] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Scalar (whole-tensor) scale on the LHS and per-dimension scale on the RHS
+// impose no additional constraint, so the tile sizes can be arbitrary (and
+// need not divide the operand shapes).
+
+// CHECK-LABEL: func.func @tile_scaled_contract_scalar_and_dimension_scale(
+// CHECK-SAME: %[[A:[a-zA-Z0-9]+]]: tensor<256x100xi8>
+// CHECK-SAME: %[[SCALE_A:[a-zA-Z0-9]+]]: tensor<f8E8M0FNU>
+// CHECK: scf.for %[[IV_M:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}}
+// CHECK: scf.for %[[IV_N:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}}
+// CHECK: scf.for %[[IV_K:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}}
+// CHECK: %[[RES:.+]] = linalg.scaled_contract
+// CHECK-SAME: ins(%{{.+}}, %[[SCALE_A]], %{{.+}}, %{{.+}} :
+// CHECK-SAME: -> tensor<?x?xf32>
+func.func @tile_scaled_contract_scalar_and_dimension_scale(
+ %A: tensor<256x100xi8>, %scaleA: tensor<f8E8M0FNU>,
+ %B: tensor<100x128xi8>, %scaleB: tensor<128xf8E8M0FNU>,
+ %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ %D = linalg.scaled_contract
+ indexing_maps = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> ()>,
+ affine_map<(m, n, k) -> (k, n)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %scaleA, %B, %scaleB
+ : tensor<256x100xi8>, tensor<f8E8M0FNU>, tensor<100x128xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.scaled_contract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:3 = transform.structured.tile_using_for %0 tile_sizes [30, 17, 25] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+// Dynamic operand shapes with static, valid tile sizes. The constraint is
+// checked against the (constant) upper bound of the tile size, so block-scaled
+// tiling of dynamic tensors is allowed as long as the tile sizes divide or are
+// divisible by the block factors.
+// CHECK-LABEL: func.func @tile_scaled_contract_dynamic_shapes(
+// CHECK: scf.for %[[IV_M:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}}
+// CHECK: scf.for %[[IV_N:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}}
+// CHECK: scf.for %[[IV_K:[a-zA-Z0-9]+]] = %{{.+}} step %{{.+}}
+// CHECK: %[[RES:.+]] = linalg.scaled_contract
+// CHECK-SAME: -> tensor<?x?xf32>
+func.func @tile_scaled_contract_dynamic_shapes(
+ %A: tensor<?x?xi8>, %scaleA: tensor<?x?xf8E8M0FNU>,
+ %B: tensor<?x?xi8>, %scaleB: tensor<?xf8E8M0FNU>,
+ %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
+ %D = linalg.scaled_contract
+ indexing_maps = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %scaleA, %B, %scaleB
+ : tensor<?x?xi8>, tensor<?x?xf8E8M0FNU>, tensor<?x?xi8>, tensor<?xf8E8M0FNU>)
+ outs(%C : tensor<?x?xf32>) -> tensor<?x?xf32>
+ return %D : tensor<?x?xf32>
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.scaled_contract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:3 = transform.structured.tile_using_for %0 tile_sizes [32, 16, 128] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
+
+// -----
+
+func.func @negative_tile_scaled_contract_non_divisible_block_scale(
+ %A: tensor<256x512xi8>, %scaleA: tensor<8x4xf8E8M0FNU>,
+ %B: tensor<128x512xi8>, %scaleB: tensor<128xf8E8M0FNU>,
+ %C: tensor<256x128xf32>) -> tensor<256x128xf32> {
+ // expected-error @below {{'linalg.scaled_contract' op tile size 48 for dim 0 must divide or be divisible by scale factor 32}}
+ // expected-error @below {{'linalg.scaled_contract' op failed to tile operation}}
+ // expected-error @below {{'linalg.scaled_contract' op failed to generate tiling loops}}
+ %D = linalg.scaled_contract
+ indexing_maps = [
+ affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (m floordiv 32, k floordiv 128)>,
+ affine_map<(m, n, k) -> (n, k)>,
+ affine_map<(m, n, k) -> (n)>,
+ affine_map<(m, n, k) -> (m, n)>]
+ ins(%A, %scaleA, %B, %scaleB
+ : tensor<256x512xi8>, tensor<8x4xf8E8M0FNU>, tensor<128x512xi8>, tensor<128xf8E8M0FNU>)
+ outs(%C : tensor<256x128xf32>) -> tensor<256x128xf32>
+ return %D : tensor<256x128xf32>
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.scaled_contract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1, %loops:3 = transform.structured.tile_using_for %0 tile_sizes [48, 16, 128] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+ transform.yield
+ }
+}
More information about the Mlir-commits
mailing list