[Mlir-commits] [mlir] [MLIR][Vector][Linalg] Expose allow_fast_math_flag_loss in transform ops (PR #236856)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 30 20:52:31 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: Md Abdullah Shahneous Bari (mshahneo)
<details>
<summary>Changes</summary>
#<!-- -->233546 added an `allowFastMathFlagLoss` option to `populateVectorReductionToContractPatterns`. It is `true` by default. If it is `false`, a multiply and a `vector.multi_reduction` whose fast-math flags differ are not rewritten to a `vector.contract`, so no flag is lost.
The two transform ops that use this pattern set cannot change the option today. This patch adds an `allow_fast_math_flag_loss` option to both, with the same name and default as the C++ option:
```mlir
transform.apply_patterns.vector.reduction_to_contract allow_fast_math_flag_loss = false
transform.structured.vectorize_children_and_apply_patterns %f <allow_fast_math_flag_loss = false> : (!transform.any_op) -> !transform.any_op
```
The Python wrapper of `VectorizeChildrenAndApplyPatternsOp` takes the new option too.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---
Full diff: https://github.com/llvm/llvm-project/pull/236856.diff
9 Files Affected:
- (modified) mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td (+7-1)
- (modified) mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td (+13-1)
- (modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+2-1)
- (modified) mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp (+2-1)
- (modified) mlir/python/mlir/dialects/transform/structured.py (+2)
- (modified) mlir/test/Dialect/Linalg/vectorization/linalg-ops-with-patterns.mlir (+65)
- (modified) mlir/test/Dialect/Vector/transform-vector.mlir (+50)
- (modified) mlir/test/python/dialects/transform_structured_ext.py (+4)
- (modified) mlir/test/python/dialects/transform_vector_ext.py (+3)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index c31cd399032986..1922b35ce1af0a 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -2560,6 +2560,11 @@ def VectorizeChildrenAndApplyPatternsOp :
- `disable_multi_reduction_to_contract_patterns`: a `UnitAttr` to deactivate
the rewrite of `vector.multi_reduction` to `vector.contract`. This is
intended to be used in tests only.
+ - `allow_fast_math_flag_loss`: a `BoolAttr`, `true` by default. The rewrite
+ of `vector.multi_reduction` to `vector.contract` gives the contract only
+ the fast-math flags that both the multiply and the reduction carry. If
+ `false`, a multiply and a reduction whose flags differ are not rewritten,
+ so no flag is lost.
- `disable_transfer_permutation_map_lowering_patterns`: a `UnitAttr` to
deactivate the rewrite of `vector.transfer` with permutation maps into
explicit `vector.transpose` operations. This is intended to be used in
@@ -2579,7 +2584,8 @@ def VectorizeChildrenAndApplyPatternsOp :
UnitAttr:$vectorize_nd_extract,
UnitAttr:$flatten_1d_depthwise_conv,
UnitAttr:$disable_multi_reduction_to_contract_patterns,
- UnitAttr:$disable_transfer_permutation_map_lowering_patterns);
+ UnitAttr:$disable_transfer_permutation_map_lowering_patterns,
+ DefaultValuedAttr<BoolAttr, "true">:$allow_fast_math_flag_loss);
let results = (outs TransformHandleTypeInterface:$transformed);
let assemblyFormat =
diff --git a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
index af690211b894de..9283b1f4d76972 100644
--- a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
+++ b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
@@ -615,9 +615,21 @@ def ApplyVectorReductionToContractPatternsOp : Op<Transform_Dialect,
These patterns have the effect of rewriting a vector.multi_reduce into a
vector.contract.
+
+ The new contract keeps only the fast-math flags that both the multiply and
+ the reduction carry. With `allow_fast_math_flag_loss = false`, a multiply
+ and a reduction whose flags differ are not rewritten, so no flag is lost.
}];
- let assemblyFormat = "attr-dict";
+ let arguments = (ins
+ DefaultValuedAttr<BoolAttr, "true">:$allow_fast_math_flag_loss
+ );
+
+ let assemblyFormat = [{
+ oilist (
+ `allow_fast_math_flag_loss` `=` $allow_fast_math_flag_loss
+ ) attr-dict
+ }];
}
def ApplySinkVectorPatternsOp : Op<Transform_Dialect,
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 816d7499845789..a2b80315ae4f59 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -4299,7 +4299,8 @@ transform::VectorizeChildrenAndApplyPatternsOp::applyToOne(
vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
if (!getDisableMultiReductionToContractPatterns())
- vector::populateVectorReductionToContractPatterns(patterns);
+ vector::populateVectorReductionToContractPatterns(
+ patterns, /*benefit=*/1, getAllowFastMathFlagLoss());
vector::populateSinkVectorOpsPatterns(patterns);
diff --git a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
index 44d461e4195c4f..257f5276fadf01 100644
--- a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
+++ b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
@@ -107,7 +107,8 @@ void transform::ApplyFoldElementwiseToVectorPatternsOp::populatePatterns(
void transform::ApplyVectorReductionToContractPatternsOp::populatePatterns(
RewritePatternSet &patterns) {
- vector::populateVectorReductionToContractPatterns(patterns);
+ vector::populateVectorReductionToContractPatterns(patterns, /*benefit=*/1,
+ getAllowFastMathFlagLoss());
}
void transform::ApplyLowerCreateMaskPatternsOp::populatePatterns(
diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index 9862447824bc94..af0dd8c5ce649a 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -770,6 +770,7 @@ def __init__(
vectorize_padding: bool = False,
flatten_1d_depthwise_conv: bool = False,
fold_type_extensions_into_contract: bool = False,
+ allow_fast_math_flag_loss: bool = True,
loc=None,
ip=None,
):
@@ -783,6 +784,7 @@ def __init__(
vectorize_nd_extract=vectorize_nd_extract,
vectorize_padding=vectorize_padding,
fold_type_extensions_into_contract=fold_type_extensions_into_contract,
+ allow_fast_math_flag_loss=allow_fast_math_flag_loss,
loc=loc,
ip=ip,
)
diff --git a/mlir/test/Dialect/Linalg/vectorization/linalg-ops-with-patterns.mlir b/mlir/test/Dialect/Linalg/vectorization/linalg-ops-with-patterns.mlir
index 28eaf74f237088..9dec965e9eac36 100644
--- a/mlir/test/Dialect/Linalg/vectorization/linalg-ops-with-patterns.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization/linalg-ops-with-patterns.mlir
@@ -94,6 +94,71 @@ module attributes {transform.with_named_sequence} {
// -----
+// The mulf has no fast-math flags, but the addf has. By default, they are
+// converted to a vector.contract, which keeps the flags that both carry (none).
+
+// CHECK-LABEL: @generic_matmul_uneven_fast_math_flags
+// CHECK-NOT: vector.multi_reduction
+// CHECK: vector.contract {indexing_maps =
+func.func @generic_matmul_uneven_fast_math_flags(%A: memref<8x16xf32>, %B: memref<16x32xf32>, %C: memref<8x32xf32>) {
+ linalg.generic {
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (k, n)>,
+ affine_map<(m, n, k) -> (m, n)>],
+ iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%A, %B : memref<8x16xf32>, memref<16x32xf32>) outs(%C : memref<8x32xf32>) {
+ ^bb0(%a: f32, %b: f32, %c: f32):
+ %mul = arith.mulf %a, %b : f32
+ %add = arith.addf %c, %mul fastmath<reassoc> : f32
+ linalg.yield %add : f32
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.get_parent_op %0 <isolated_from_above> : (!transform.any_op) -> !transform.any_op
+ %2 = transform.structured.vectorize_children_and_apply_patterns %1 : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// Same as above, but with allow_fast_math_flag_loss = false: the mulf and the
+// multi_reduction are not converted, so the reduction keeps its flags.
+
+// CHECK-LABEL: @generic_matmul_keep_fast_math_flags
+// CHECK-NOT: vector.contract
+// CHECK: arith.mulf
+// CHECK-NEXT: vector.multi_reduction <add>, %{{.+}}, %{{.+}} fastmath<reassoc> [2]
+func.func @generic_matmul_keep_fast_math_flags(%A: memref<8x16xf32>, %B: memref<16x32xf32>, %C: memref<8x32xf32>) {
+ linalg.generic {
+ indexing_maps = [affine_map<(m, n, k) -> (m, k)>,
+ affine_map<(m, n, k) -> (k, n)>,
+ affine_map<(m, n, k) -> (m, n)>],
+ iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%A, %B : memref<8x16xf32>, memref<16x32xf32>) outs(%C : memref<8x32xf32>) {
+ ^bb0(%a: f32, %b: f32, %c: f32):
+ %mul = arith.mulf %a, %b : f32
+ %add = arith.addf %c, %mul fastmath<reassoc> : f32
+ linalg.yield %add : f32
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.get_parent_op %0 <isolated_from_above> : (!transform.any_op) -> !transform.any_op
+ %2 = transform.structured.vectorize_children_and_apply_patterns %1 <allow_fast_math_flag_loss = false> : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
// CHECK-LABEL: func @vectorization_test_2
func.func @vectorization_test_2(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
%C: memref<8x32xf32>) {
diff --git a/mlir/test/Dialect/Vector/transform-vector.mlir b/mlir/test/Dialect/Vector/transform-vector.mlir
index 16011e3041eb55..1f74a191d034b6 100644
--- a/mlir/test/Dialect/Vector/transform-vector.mlir
+++ b/mlir/test/Dialect/Vector/transform-vector.mlir
@@ -97,6 +97,56 @@ module attributes {transform.with_named_sequence} {
// -----
+// The mulf has no fast-math flags, but the multi_reduction has. By default,
+// they are converted, and the contract keeps the flags that both carry (none).
+
+// CHECK-LABEL: func.func @reduction_to_contract_uneven_fast_math_flags
+// CHECK-NEXT: %[[R:.+]] = vector.contract {indexing_maps =
+// CHECK-SAME: kind = #vector.kind<add>}
+// CHECK-NEXT: return %[[R]]
+func.func @reduction_to_contract_uneven_fast_math_flags(%arg0: vector<8x32x16xf32>, %arg1: vector<8x32x16xf32>, %acc: vector<8x16xf32>) -> vector<8x16xf32> {
+ %0 = arith.mulf %arg0, %arg1 : vector<8x32x16xf32>
+ %1 = vector.multi_reduction <add>, %0, %acc fastmath<reassoc> [1] : vector<8x32x16xf32> to vector<8x16xf32>
+ return %1 : vector<8x16xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %module_op : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.vector.reduction_to_contract
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// Same as above, but with allow_fast_math_flag_loss = false: the two ops are
+// not converted, so the reduction keeps its flags.
+
+// CHECK-LABEL: func.func @reduction_to_contract_keep_fast_math_flags
+// CHECK-NEXT: %[[M:.+]] = arith.mulf
+// CHECK-NEXT: %[[R:.+]] = vector.multi_reduction <add>, %[[M]], %{{.+}} fastmath<reassoc> [1]
+// CHECK-NEXT: return %[[R]]
+func.func @reduction_to_contract_keep_fast_math_flags(%arg0: vector<8x32x16xf32>, %arg1: vector<8x32x16xf32>, %acc: vector<8x16xf32>) -> vector<8x16xf32> {
+ %0 = arith.mulf %arg0, %arg1 : vector<8x32x16xf32>
+ %1 = vector.multi_reduction <add>, %0, %acc fastmath<reassoc> [1] : vector<8x32x16xf32> to vector<8x16xf32>
+ return %1 : vector<8x16xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %module_op : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.vector.reduction_to_contract allow_fast_math_flag_loss = false
+ } : !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
// CHECK-LABEL: func.func @arith_to_outerproduct_scalable_i32
// CHECK-SAME: %[[LHS:.*]]: vector<[4]xi32>,
// CHECK-SAME: %[[RHS:.*]]: vector<[4]xi32>) -> vector<[4]x[4]xi32> {
diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py
index b96c8edb177c41..7387d2f662fe67 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -726,6 +726,7 @@ def testVectorizeChildrenAndApplyPatternsAllAttrs(target):
vectorize_padding=True,
flatten_1d_depthwise_conv=True,
fold_type_extensions_into_contract=True,
+ allow_fast_math_flag_loss=False,
)
# CHECK-LABEL: TEST: testVectorizeChildrenAndApplyPatternsAllAttrs
# CHECK: transform.sequence
@@ -736,6 +737,7 @@ def testVectorizeChildrenAndApplyPatternsAllAttrs(target):
# CHECK-SAME: flatten_1d_depthwise_conv
# CHECK-SAME: disable_multi_reduction_to_contract_patterns
# CHECK-SAME: disable_transfer_permutation_map_lowering_patterns
+ # CHECK-SAME: allow_fast_math_flag_loss = false
@run
@@ -749,6 +751,7 @@ def testVectorizeChildrenAndApplyPatternsNoAttrs(target):
vectorize_padding=False,
flatten_1d_depthwise_conv=False,
fold_type_extensions_into_contract=False,
+ allow_fast_math_flag_loss=True,
)
# CHECK-LABEL: TEST: testVectorizeChildrenAndApplyPatternsNoAttrs
# CHECK: transform.sequence
@@ -759,6 +762,7 @@ def testVectorizeChildrenAndApplyPatternsNoAttrs(target):
# CHECK-NOT: fold_type_extensions_into_contract
# CHECK-NOT: vectorize_nd_extract
# CHECK-NOT: vectorize_padding
+ # CHECK-NOT: allow_fast_math_flag_loss
@run
diff --git a/mlir/test/python/dialects/transform_vector_ext.py b/mlir/test/python/dialects/transform_vector_ext.py
index 26225e014ed7b2..161591847d381b 100644
--- a/mlir/test/python/dialects/transform_vector_ext.py
+++ b/mlir/test/python/dialects/transform_vector_ext.py
@@ -77,6 +77,9 @@ def configurable_patterns():
# CHECK: transform.apply_patterns.vector.flatten_vector_transfer_ops
# CHECK-SAME: target_vector_bitwidth = 1
vector.ApplyFlattenVectorTransferOpsPatternsOp(target_vector_bitwidth=1)
+ # CHECK: transform.apply_patterns.vector.reduction_to_contract
+ # CHECK-SAME: allow_fast_math_flag_loss = false
+ vector.ApplyVectorReductionToContractPatternsOp(allow_fast_math_flag_loss=False)
@run_apply_patterns
``````````
</details>
https://github.com/llvm/llvm-project/pull/236856
More information about the Mlir-commits
mailing list