[MLIR][Linalg] Fix linalg crash during elementwise op fusion

[IanWood1]

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). This is because concatAffineMaps wasn't handling the maps being empty properly.

Similar to canOpOperandsBeDroppedImpl, I added an early return when the maps are of size zero. Additionally, concatAffineMaps's declaration comment says it returns an empty map when maps is empty but it has no way to get the MLIRContext needed to construct the empty affine map when the array is empty. So, I changed this to take the context.

NOTE: concatAffineMaps now takes an MLIRContext to be able to construct an empty map in the case where maps is empty.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Ian Wood (IanWood1)

Changes

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 --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

[qedawkins]

LGTM

[banach-space]

If this is fixing a crash, then there should be a test demonstrating that, no?

Argh, ignore me, sorry 🤦🏻