[LinalgExt] Fix reshape fusion crash (#21472)

The upstream reshape builder requires that no change in static/dynamic
dims. This causes `getOrCreateExpanded` to fail but at that point the
pattern has already modified the IR so it's too late for the pattern to
return `failure()` (will cause infinite looping). Instead, exit early
when computing expansion info for ops with mixed static & dynamic dims.


Related issue for propagating static shape info
#21471
Closes #21439

Signed-off-by: Ian Wood <[email protected]>

Author: IanWood1

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp

@@ -200,6 +200,18 @@ LogicalResult ExpansionInfo::compute(
   if (operandReassoc.empty())
     return failure();
 
+  // Check that the operand dim size matches the iteration space dim size. This
+  // can fail when one is static and the other is dynamic.
+  for (const ReshapeOperandInfo &info : infos) {
+    for (auto [operandDim, iterDim] :
+         llvm::enumerate(info.operandToIterationSpace)) {
+      if (iterDim != ReshapeOperandInfo::kNoMapping &&
+          loopRanges[iterDim] != info.originalShape[operandDim]) {
+        return failure();
+      }
+    }
+  }
+
   int64_t operandNum = fusableOpOperand->getOperandNumber();
   ReshapeOperandInfo &fusionOperandInfo = infos[operandNum];
   this->loopShapeMap.clear();
@@ -413,8 +425,10 @@ fuseWithReshapeByExpansion(OpTy op, Operation *reshapeOp,
 
   IRMapping mapping;
   for (OpOperand &operand : op->getOpOperands()) {
-    mapping.map(operand.get(),
-                info.getOrCreateExpanded(loc, &operand, rewriter).value());
+    std::optional<Value> maybeNewOperand =
+        info.getOrCreateExpanded(loc, &operand, rewriter);
+    assert(maybeNewOperand.has_value());
+    mapping.map(operand.get(), maybeNewOperand.value());
   }
 
   assert(op.getNumDpsInits() == 1);

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/reshape_fusion.mlir

@@ -482,6 +482,32 @@ util.func public @sink_through_k2(%0 : tensor<128x16x2x2x128xf16>, %1 : tensor<1
 
 // -----
 
+#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d2)>
+#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d4)>
+#map3 = affine_map<(d0, d1, d2, d3, d4) -> ()>
+#map4 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d4)>
+
+util.func public @no_fuse_attention_mixed_static_dynamic(%arg0: tensor<?x4096x16xf16>, %arg1: tensor<20x1024x16xf16>, %arg2: tensor<20x1024x64xf16>, %arg3: f16) -> tensor<2x10x4096x64xf16> {
+  %0 = tensor.empty() : tensor<20x4096x64xf16>
+  %1 = iree_linalg_ext.attention {indexing_maps = [#map, #map1, #map2, #map3, #map4]} ins(%arg0, %arg1, %arg2, %arg3 : tensor<?x4096x16xf16>, tensor<20x1024x16xf16>, tensor<20x1024x64xf16>, f16) outs(%0 : tensor<20x4096x64xf16>) {
+    ^bb0(%score: f16):
+      iree_linalg_ext.yield %score: f16
+  } -> tensor<20x4096x64xf16>
+  %expanded = tensor.expand_shape %1 [[0, 1], [2], [3]] output_shape [2, 10, 4096, 64] : tensor<20x4096x64xf16> into tensor<2x10x4096x64xf16>
+  util.return %expanded : tensor<2x10x4096x64xf16>
+}
+
+//CHECK-LABEL: func public @no_fuse_attention_mixed_static_dynamic(
+// CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x4096x16xf16>
+// CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: tensor<20x1024x16xf16>
+// CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]: tensor<20x1024x64xf16>
+// CHECK-SAME:     %[[ARG3:.+]]: f16)
+//      CHECK:   %[[ATTENTION:.+]] = iree_linalg_ext.attention
+// CHECK-SAME:       ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]] :
+//      CHECK:   tensor.expand_shape %[[ATTENTION]]
+
+
 util.func @scatter_collapse_updates(%arg0: tensor<4x?x2x16x4x128xf16>, %arg1: tensor<?x1xi32>, %arg2: tensor<?x2x16x4x128xf16>) -> tensor<?x2x16x4x128xf16> {
   %collapsed = tensor.collapse_shape %arg0[[0, 1], [2], [3], [4], [5]] : tensor<4x?x2x16x4x128xf16> into tensor<?x2x16x4x128xf16>
   %1 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(true) ins(%collapsed, %arg1 : tensor<?x2x16x4x128xf16>, tensor<?x1xi32>) outs(%arg2 : tensor<?x2x16x4x128xf16>) {