[Mlir-commits] [mlir] [MLIR][Linalg] Fix linalg crash during elementwise op fusion (PR #117667)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Ian Wood (IanWood1)

<details>
<summary>Changes</summary>

`isOpOperandCanBeDroppedAfterFusedLinalgs` crashes when `indexingMaps` is empty. This can occur when `producer` only has DPS init operands and `consumer ` only has a single DPS input operand (all operands are ignored and nothing gets added to `indexingMaps`). 
Additionally, `concatAffineMaps` has no way to get the `MLIRContext` when the array is empty. This function was crashing when trying to get the context from the first map in the array. To fix this, I added an early return when the maps are of size zero and changed the impl of `concatAffineMaps` to assert when given an empty map. `concatAffineMaps` should probably return an empty map when given an empty array, but that can be left to a later PR. 

---
Full diff: https://github.com/llvm/llvm-project/pull/117667.diff


4 Files Affected:

- (modified) mlir/include/mlir/IR/AffineMap.h (+4-1) 
- (modified) mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp (+2-1) 
- (modified) mlir/lib/IR/AffineMap.cpp (+1) 
- (modified) mlir/test/Dialect/Linalg/fusion-elementwise.mlir (+31) 


``````````diff
diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h
index e30950bbf292d6..4cfc538a2fe72f 100644
--- a/mlir/include/mlir/IR/AffineMap.h
+++ b/mlir/include/mlir/IR/AffineMap.h
@@ -595,7 +595,10 @@ AffineMap inverseAndBroadcastProjectedPermutation(AffineMap map);
 /// potentially empty maps. Assumes each of the underlying map has 0 symbols.
 /// The resulting map has a number of dims equal to the max of `maps`' dims and
 /// the concatenated results as its results.
-/// Returns an empty map if all input `maps` are empty.
+///
+/// This method asserts when `maps` is empty.
+/// TODO: this should return an empty map when `maps` is empty but there is no
+/// way to get the MLIRContext needed to construct it.
 ///
 /// Example:
 /// When applied to the following list of 3 affine maps,
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index a934e47794051c..6ddea53dfb997a 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -93,7 +93,8 @@ static bool isOpOperandCanBeDroppedAfterFusedLinalgs(
   // the bounds of the op can be still computed after dropping the selected
   // operand. inversePermutation returns an empty AffineMap in case the
   // concatanated indexing maps are not invertible.
-  return inversePermutation(concatAffineMaps(indexingMaps)) != AffineMap();
+  return !indexingMaps.empty() &&
+         inversePermutation(concatAffineMaps(indexingMaps)) != AffineMap();
 }
 
 /// Returns a set of indices of the producer's results which would
diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index ea3c0723b07759..8056058511ebe6 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -834,6 +834,7 @@ AffineMap mlir::inverseAndBroadcastProjectedPermutation(AffineMap map) {
 }
 
 AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
+  assert(maps.size());
   unsigned numResults = 0, numDims = 0, numSymbols = 0;
   for (auto m : maps)
     numResults += m.getNumResults();
diff --git a/mlir/test/Dialect/Linalg/fusion-elementwise.mlir b/mlir/test/Dialect/Linalg/fusion-elementwise.mlir
index 8131e4054cc6b0..bd9977f1410b91 100644
--- a/mlir/test/Dialect/Linalg/fusion-elementwise.mlir
+++ b/mlir/test/Dialect/Linalg/fusion-elementwise.mlir
@@ -28,3 +28,34 @@ func.func @drop_unused_producer_result(%arg0 : tensor<?x?xf32>,
 //       CHECK:   %[[FUSED_OP:.+]] = linalg.generic
 //  CHECK-SAME:       ins(%[[ARG0]], %[[ARG1]] :
 //       CHECK:   return %[[FUSED_OP]]
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+func.func @handle_unused_operands(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<8xf32> {
+  %cst_0 = arith.constant 0.000000e+00 : f32
+  %cst_1 = arith.constant 1.000000e+00 : f32
+  %0:2 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} outs(%arg0, %arg1 : tensor<8xf32>, tensor<8xf32>) {
+  ^bb0(%out: f32, %out_2: f32):
+    %1 = linalg.index 0 : index
+    %2 = arith.index_cast %1 : index to i64
+    %3 = arith.sitofp %2 : i64 to f32
+    %4 = arith.divf %3, %cst_0 : f32
+    linalg.yield %3, %4 : f32, f32
+  } -> (tensor<8xf32>, tensor<8xf32>)
+  linalg.generic {indexing_maps = [#map], iterator_types = ["parallel"]} ins(%0#1 : tensor<8xf32>) {
+  ^bb0(%in: f32):
+    %2 = arith.cmpf one, %in, %cst_1 : f32
+    cf.assert %2, "Side effect op"
+    linalg.yield
+  }
+  func.return %arg1 : tensor<8xf32>
+}
+
+// CHECK-LABEL: func @handle_unused_operands
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9]+]]: tensor<8xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9]+]]: tensor<8xf32>
+//       CHECK:   %[[EMPTY:.+]] = tensor.empty() : tensor<8xf32>
+//       CHECK:   %[[FUSED_OP:.+]] = linalg.generic
+//  CHECK-SAME:       outs(%[[EMPTY]] :
+//   CHECK-NOT:   linalg.generic

``````````

</details>


https://github.com/llvm/llvm-project/pull/117667