Revert "[LinalgExt] Fix FoldWithProducerReshapeByExpansion for >1 dyn dim" (iree-org#21947)

IanWood1 · web-flow · commit 515f29290ce1 · 2025-09-11T16:34:00.000-07:00
Reverts iree-org#21894 Breaks iree-org#21889 Fixes iree-org#21941
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ReshapeFusion.cpp
@@ -11,62 +11,17 @@
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtInterfaces.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
 #include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
-#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/PatternMatch.h"
-#include "mlir/Transforms/RegionUtils.h"
 
 #include <cstdint>
 #include <optional>
 
 namespace mlir::iree_compiler::IREE::LinalgExt {
-namespace {
-
-/// Represents the size of a dimension of some ShapedType value in the IR. This
-/// is used instead of OpFoldResult when modifying the IR is illegal. This can
-/// still be constructed from an OpFoldResult in cases where the value can be
-/// obtained without IR modification.
-class DimSize {
-public:
-  DimSize(TypedValue<ShapedType> val, int64_t dim)
-      : ofr(nullptr), val(val), dim(dim) {}
-  DimSize(OpFoldResult ofr) : ofr(ofr), val(nullptr), dim(-1) {}
-
-  bool isStatic() const {
-    if (ofr) {
-      return getConstantIntValue(ofr).has_value();
-    }
-    return val.getType().isStaticDim(dim);
-  }
-
-  // Get an OpFoldResult by possibly inserting IR.
-  OpFoldResult materialize(OpBuilder &b) const {
-    if (ofr) {
-      return ofr;
-    }
-    return getDim(b, val.getLoc(), val, dim);
-  }
-
-private:
-  OpFoldResult ofr;
-  TypedValue<ShapedType> val;
-  int64_t dim;
-};
-} // namespace
-
-static SmallVector<DimSize> getDimSizes(Value v) {
-  auto shapedVal = cast<TypedValue<ShapedType>>(v);
-  int64_t rank = shapedVal.getType().getRank();
-  SmallVector<DimSize> sizes;
-  for (int i = 0; i < rank; ++i) {
-    sizes.emplace_back(shapedVal, i);
-  }
-  return sizes;
-}
 
 static bool
 isIdentityReassoc(const SmallVector<ReassociationIndices> &indices) {
@@ -78,7 +33,7 @@ isIdentityReassoc(const SmallVector<ReassociationIndices> &indices) {
 };
 
 static SmallVector<ReassociationIndices>
-computeReassocFromShapeMap(ArrayRef<SmallVector<DimSize>> shapeMap) {
+computeReassocFromShapeMap(ArrayRef<SmallVector<int64_t>> shapeMap) {
   SmallVector<ReassociationIndices> reassoc;
   int64_t dimCount = 0;
   for (auto &shape : shapeMap) {
@@ -90,13 +45,14 @@ computeReassocFromShapeMap(ArrayRef<SmallVector<DimSize>> shapeMap) {
 }
 
 namespace {
+
 /// Helper class that supports fusing reshapes with operands when not all of the
 /// shape dims map to the iteration space.
 struct ReshapeOperandInfo {
   static constexpr int64_t kNoMapping = -1;
 
   // Original shape of this operand.
-  SmallVector<DimSize> originalShape;
+  ArrayRef<int64_t> originalShape;
 
   // Similar to the results of the operand's `AffineMap` except `kNoMapping` if
   // that dim doesn't map to the iteration space. For example, the indexed
@@ -116,7 +72,7 @@ class ExpansionInfo {
                         SmallVector<int64_t> loopRanges,
                         OpOperand *fusableOpOperand,
                         ArrayRef<ReassociationIndices> operandReassoc,
-                        ArrayRef<DimSize> expandedShape);
+                        ArrayRef<int64_t> expandedShape);
 
   std::optional<Value> getOrCreateExpanded(Location loc, OpOperand *operand,
                                            RewriterBase &rewriter) {
@@ -125,17 +81,13 @@ class ExpansionInfo {
     if (isIdentityReassoc(reassoc)) {
       return operand->get();
     }
-    SmallVector<OpFoldResult> outputShape;
+    SmallVector<int64_t> flattenedArray;
     for (auto &shape : shapeMap) {
-      llvm::append_range(
-          outputShape, llvm::map_range(shape, [&rewriter](const DimSize &size) {
-            return size.materialize(rewriter);
-          }));
+      flattenedArray.append(shape.begin(), shape.end());
     }
-    auto [staticShape, dynamicShape] = decomposeMixedValues(outputShape);
-    (void)dynamicShape;
     auto oldType = cast<ShapedType>(operand->get().getType());
-    auto newType = RankedTensorType::get(staticShape, oldType.getElementType());
+    auto newType =
+        RankedTensorType::get(flattenedArray, oldType.getElementType());
     if (failed(reshapeLikeShapesAreCompatible(
             [&](const Twine &msg) {
               return rewriter.notifyMatchFailure(loc, msg);
@@ -145,18 +97,18 @@ class ExpansionInfo {
       return {};
     }
     return tensor::ExpandShapeOp::create(rewriter, loc, newType, operand->get(),
-                                         reassoc, outputShape);
+                                         reassoc);
   };
 
   /// Get the shape map for the operand.
-  SmallVector<SmallVector<DimSize>> getShapeMap(OpOperand *operand) const {
+  SmallVector<SmallVector<int64_t>> getShapeMap(OpOperand *operand) const {
     auto info = reshapeInfos[operand->getOperandNumber()];
-    SmallVector<SmallVector<DimSize>> shapeMap;
+    SmallVector<SmallVector<int64_t>> shapeMap;
     for (auto [operandIdx, loopIdx] :
          llvm::enumerate(info.operandToIterationSpace)) {
       if (loopIdx == ReshapeOperandInfo::kNoMapping) {
         shapeMap.push_back(
-            SmallVector<DimSize>{info.originalShape[operandIdx]});
+            SmallVector<int64_t>{info.originalShape[operandIdx]});
       } else {
         shapeMap.push_back(loopShapeMap[loopIdx]);
       }
@@ -174,12 +126,17 @@ class ExpansionInfo {
   ReassociationIndicesRef getExpandedLoops(unsigned i) const {
     return loopReassoc[i];
   }
+  ArrayRef<int64_t> getExpandedShapeOfLoop(unsigned i) const {
+    return loopShapeMap[i];
+  }
 
 private:
+  /// Extent of the iteration space in the original operation.
+  SmallVector<int64_t> loopRanges;
   SmallVector<ReassociationIndices> loopReassoc;
   /// Mapping from extent of loops in the original operation, to the extent of
   /// loops in the expanded operation.
-  SmallVector<SmallVector<DimSize>> loopShapeMap;
+  SmallVector<SmallVector<int64_t>> loopShapeMap;
   unsigned expandedOpNumDims;
   /// Info about the reassociation and original shape for each operand.
   SmallVector<ReshapeOperandInfo> reshapeInfos;
@@ -239,7 +196,7 @@ class CollapsingInfo {
 LogicalResult ExpansionInfo::compute(
     SmallVector<ReshapeOperandInfo> infos, SmallVector<int64_t> loopRanges,
     OpOperand *fusableOpOperand, ArrayRef<ReassociationIndices> operandReassoc,
-    ArrayRef<DimSize> expandedShape) {
+    ArrayRef<int64_t> expandedShape) {
   if (operandReassoc.empty())
     return failure();
 
@@ -249,8 +206,7 @@ LogicalResult ExpansionInfo::compute(
     for (auto [operandDim, iterDim] :
          llvm::enumerate(info.operandToIterationSpace)) {
       if (iterDim != ReshapeOperandInfo::kNoMapping &&
-          ShapedType::isStatic(loopRanges[iterDim]) !=
-              info.originalShape[operandDim].isStatic()) {
+          loopRanges[iterDim] != info.originalShape[operandDim]) {
         return failure();
       }
     }
@@ -273,22 +229,12 @@ LogicalResult ExpansionInfo::compute(
     }
   }
 
-  // Fill in the remaining elements.
-  for (const ReshapeOperandInfo &info : infos) {
-    for (auto [operandIdx, loopIdx] :
-         llvm::enumerate(info.operandToIterationSpace)) {
-      if (loopIdx == ReshapeOperandInfo::kNoMapping ||
-          !this->loopShapeMap[loopIdx].empty()) {
-        continue;
-      }
-
-      this->loopShapeMap[loopIdx] =
-          SmallVector<DimSize>{info.originalShape[operandIdx]};
-    }
-  }
-
+  // Fill in the remaining elements with `loopRanges`
   this->expandedOpNumDims = 0;
-  for (const auto &shapeMap : this->loopShapeMap) {
+  for (const auto &[loopIdx, shapeMap] : llvm::enumerate(this->loopShapeMap)) {
+    if (shapeMap.empty()) {
+      this->loopShapeMap[loopIdx] = SmallVector<int64_t>{loopRanges[loopIdx]};
+    }
     this->expandedOpNumDims += shapeMap.size();
   }
 
@@ -298,6 +244,7 @@ LogicalResult ExpansionInfo::compute(
   }
   this->loopReassoc = computeReassocFromShapeMap(this->loopShapeMap);
   this->reshapeInfos = std::move(infos);
+  this->loopRanges = std::move(loopRanges);
   return success();
 }
 
@@ -360,7 +307,7 @@ getReshapeInfo(LinalgExt::AttentionOp attentionOp) {
           return operandInfo;
         }
 
-        operandInfo.originalShape = getDimSizes(opOperand.get());
+        operandInfo.originalShape = operandType.getShape();
         for (auto result :
              attentionOp.getMatchingIndexingMap(&opOperand).getResults()) {
           operandInfo.operandToIterationSpace.push_back(
@@ -378,13 +325,13 @@ getReshapeInfo(LinalgExt::ScatterOp scatterOp) {
   auto updateRank = scatterOp.getUpdateType().getRank();
 
   ReshapeOperandInfo updateInfo;
-  updateInfo.originalShape = getDimSizes(scatterOp.getUpdates());
+  updateInfo.originalShape = scatterOp.getUpdateType().getShape();
   llvm::append_range(updateInfo.operandToIterationSpace,
                      llvm::seq<int64_t>(0, updateRank));
   infos.push_back(std::move(updateInfo));
 
   ReshapeOperandInfo indicesInfo;
-  indicesInfo.originalShape = getDimSizes(scatterOp.getIndices());
+  indicesInfo.originalShape = scatterOp.getIndicesType().getShape();
   llvm::append_range(indicesInfo.operandToIterationSpace,
                      llvm::seq<int64_t>(0, scatterOp.getBatchRank()));
   if (scatterOp.getBatchRank() != scatterOp.getIndicesType().getRank())
@@ -393,7 +340,7 @@ getReshapeInfo(LinalgExt::ScatterOp scatterOp) {
   infos.push_back(std::move(indicesInfo));
 
   ReshapeOperandInfo originalInfo;
-  originalInfo.originalShape = getDimSizes(scatterOp.getOriginal());
+  originalInfo.originalShape = scatterOp.getOriginalType().getShape();
   originalInfo.operandToIterationSpace.append(scatterOp.getIndexDepth(),
                                               ReshapeOperandInfo::kNoMapping);
   llvm::append_range(originalInfo.operandToIterationSpace,
@@ -409,15 +356,15 @@ getReshapeInfo(LinalgExt::GatherOp gatherOp) {
   auto outputRank = gatherOp.getOutputType().getRank();
 
   ReshapeOperandInfo sourceInfo;
-  sourceInfo.originalShape = getDimSizes(gatherOp.getSource());
+  sourceInfo.originalShape = gatherOp.getSourceType().getShape();
   sourceInfo.operandToIterationSpace.append(gatherOp.getIndexDepth(),
                                             ReshapeOperandInfo::kNoMapping);
   llvm::append_range(sourceInfo.operandToIterationSpace,
                      llvm::seq(outputRank - rankOfContiguousSlice, outputRank));
   infos.push_back(std::move(sourceInfo));
 
   ReshapeOperandInfo indicesInfo;
-  indicesInfo.originalShape = getDimSizes(gatherOp.getIndices());
+  indicesInfo.originalShape = gatherOp.getIndicesType().getShape();
   llvm::append_range(indicesInfo.operandToIterationSpace,
                      llvm::seq<int64_t>(0, gatherOp.getBatchRank()));
   if (gatherOp.getBatchRank() != gatherOp.getIndicesType().getRank())
@@ -426,7 +373,7 @@ getReshapeInfo(LinalgExt::GatherOp gatherOp) {
   infos.push_back(std::move(indicesInfo));
 
   ReshapeOperandInfo outputInfo;
-  outputInfo.originalShape = getDimSizes(gatherOp.getOutput());
+  outputInfo.originalShape = gatherOp.getOutputType().getShape();
   llvm::append_range(outputInfo.operandToIterationSpace,
                      llvm::seq<int64_t>(0, outputRank));
   infos.push_back(std::move(outputInfo));
@@ -460,26 +407,15 @@ fuseWithReshapeByExpansion(OpTy op, Operation *reshapeOp,
   auto expandingReshapeOp = dyn_cast<tensor::ExpandShapeOp>(*reshapeOp);
   auto collapsingReshapeOp = dyn_cast<tensor::CollapseShapeOp>(*reshapeOp);
   bool isExpanding = (expandingReshapeOp != nullptr);
-  Value expandedVal = isExpanding ? expandingReshapeOp.getResult()
-                                  : collapsingReshapeOp.getSrc();
-  SmallVector<DimSize> expandedSize;
-  if (isExpanding) {
-    // The SSA dims must dominate `op` in order to use them to create new
-    // expand_shape ops.
-    if (failed(moveValueDefinitions(rewriter,
-                                    expandingReshapeOp.getOutputShape(), op))) {
-      return std::nullopt;
-    }
-    llvm::append_range(expandedSize, expandingReshapeOp.getMixedOutputShape());
-  } else {
-    expandedSize = getDimSizes(expandedVal);
-  }
+  RankedTensorType expandedType = isExpanding
+                                      ? expandingReshapeOp.getResultType()
+                                      : collapsingReshapeOp.getSrcType();
   ExpansionInfo info;
   if (failed(info.compute(
           getReshapeInfo(op), op.getStaticLoopRanges(), fusableOpOperand,
           isExpanding ? expandingReshapeOp.getReassociationIndices()
                       : collapsingReshapeOp.getReassociationIndices(),
-          expandedSize))) {
+          expandedType.getShape()))) {
     return std::nullopt;
   }
 
diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/reshape_fusion.mlir b/compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/reshape_fusion.mlir
@@ -196,13 +196,19 @@ util.func public @attention_dynamic(%arg0: tensor<?x?x?xf16>, %arg1: tensor<?x?x
 //  CHECK-DAG:   %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
 //  CHECK-DAG:   %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
 //  CHECK-DAG:   %[[D2:.+]] = tensor.dim %[[ARG0]], %[[C2]]
+//  CHECK-DAG:   %[[D4:.+]] = tensor.dim %[[ARG2]], %[[C2]]
 //  CHECK-DAG:   %[[SPLIT0:.+]] = arith.divsi %[[D0]]
+//  CHECK-DAG:   %[[EMPTY:.+]] = tensor.empty(%[[SPLIT0]], %[[D1]], %[[D4]]) : tensor<2x?x?x?xf16>
 //  CHECK-DAG:   %[[QUERY:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT0]], %[[D1]], %[[D2]]]
+//  CHECK-DAG:   %[[D5:.+]] = tensor.dim %[[ARG1]], %[[C0]]
 //  CHECK-DAG:   %[[D6:.+]] = tensor.dim %[[ARG1]], %[[C1]]
-//  CHECK-DAG:   %[[KEY:.+]] = tensor.expand_shape %[[ARG1]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT0]], %[[D6]], %[[D2]]]
-//  CHECK-DAG:   %[[D9:.+]] = tensor.dim %[[ARG2]], %[[C2]]
-//  CHECK-DAG:   %[[CACHE:.+]] = tensor.expand_shape %[[ARG2]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT0]], %[[D6]], %[[D9]]]
-//  CHECK-DAG:   %[[EMPTY:.+]] = tensor.empty(%[[SPLIT0]], %[[D1]], %[[D9]]) : tensor<2x?x?x?xf16>
+//  CHECK-DAG:   %[[D7:.+]] = tensor.dim %[[ARG1]], %[[C2]]
+//  CHECK-DAG:   %[[SPLIT1:.+]] = arith.divsi %[[D5]], %[[C2]]
+//  CHECK-DAG:   %[[KEY:.+]] = tensor.expand_shape %[[ARG1]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT1]], %[[D6]], %[[D7]]]
+//  CHECK-DAG:   %[[D8:.+]] = tensor.dim %[[ARG2]], %[[C0]]
+//  CHECK-DAG:   %[[D9:.+]] = tensor.dim %[[ARG2]], %[[C1]]
+//  CHECK-DAG:   %[[SPLIT2:.+]] = arith.divsi %[[D8]], %[[C2]]
+//  CHECK-DAG:   %[[CACHE:.+]] = tensor.expand_shape %[[ARG2]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT2]], %[[D9]], %[[D4]]]
 //      CHECK:   %[[ATTENTION:.+]] = iree_linalg_ext.attention
 //  CHECK-SAME:       indexing_maps =
 //  CHECK-SAME:       affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
@@ -250,13 +256,29 @@ util.func public @attention_dynamic_masked(%arg0: tensor<?x?x?xf16>, %arg1: tens
 // CHECK-SAME:     %[[ARG3:.+]]: f16
 // CHECK-SAME:     %[[ARG4:[a-zA-Z0-9]+]]: tensor<?x?x?xf16>)
 //  CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
+//  CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
 //  CHECK-DAG:   %[[C2:.+]] = arith.constant 2 : index
-//  CHECK-DAG:   %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]]
-//  CHECK-DAG:   %[[SPLIT:.+]] = arith.divsi %[[DIM]], %[[C2]]
-//  CHECK-DAG:   %[[QUERY:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT]],
-//  CHECK-DAG:   %[[KEY:.+]] = tensor.expand_shape %[[ARG1]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT]],
-//  CHECK-DAG:   %[[CACHE:.+]] = tensor.expand_shape %[[ARG2]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT]],
-//  CHECK-DAG:   %[[MASK:.+]] = tensor.expand_shape %[[ARG4]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT]],
+//  CHECK-DAG:   %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
+//  CHECK-DAG:   %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
+//  CHECK-DAG:   %[[D2:.+]] = tensor.dim %[[ARG0]], %[[C2]]
+//  CHECK-DAG:   %[[D4:.+]] = tensor.dim %[[ARG2]], %[[C2]]
+//  CHECK-DAG:   %[[SPLIT0:.+]] = arith.divsi %[[D0]]
+//  CHECK-DAG:   %[[EMPTY:.+]] = tensor.empty(%[[SPLIT0]], %[[D1]], %[[D4]]) : tensor<2x?x?x?xf16>
+//  CHECK-DAG:   %[[QUERY:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT0]], %[[D1]], %[[D2]]]
+//  CHECK-DAG:   %[[D5:.+]] = tensor.dim %[[ARG1]], %[[C0]]
+//  CHECK-DAG:   %[[D6:.+]] = tensor.dim %[[ARG1]], %[[C1]]
+//  CHECK-DAG:   %[[D7:.+]] = tensor.dim %[[ARG1]], %[[C2]]
+//  CHECK-DAG:   %[[SPLIT1:.+]] = arith.divsi %[[D5]], %[[C2]]
+//  CHECK-DAG:   %[[KEY:.+]] = tensor.expand_shape %[[ARG1]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT1]], %[[D6]], %[[D7]]]
+//  CHECK-DAG:   %[[D8:.+]] = tensor.dim %[[ARG2]], %[[C0]]
+//  CHECK-DAG:   %[[D9:.+]] = tensor.dim %[[ARG2]], %[[C1]]
+//  CHECK-DAG:   %[[SPLIT2:.+]] = arith.divsi %[[D8]], %[[C2]]
+//  CHECK-DAG:   %[[CACHE:.+]] = tensor.expand_shape %[[ARG2]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT2]], %[[D9]], %[[D4]]]
+//  CHECK-DAG:   %[[D10:.+]] = tensor.dim %[[ARG4]], %[[C0]]
+//  CHECK-DAG:   %[[D11:.+]] = tensor.dim %[[ARG4]], %[[C1]]
+//  CHECK-DAG:   %[[D12:.+]] = tensor.dim %[[ARG4]], %[[C2]]
+//  CHECK-DAG:   %[[SPLIT3:.+]] = arith.divsi %[[D10]], %[[C2]]
+//  CHECK-DAG:   %[[MASK:.+]] = tensor.expand_shape %[[ARG4]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, %[[SPLIT3]], %[[D11]], %[[D12]]]
 //      CHECK:   %[[ATTENTION:.+]] = iree_linalg_ext.attention
 //  CHECK-SAME:       indexing_maps =
 //  CHECK-SAME:       affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
@@ -266,6 +288,7 @@ util.func public @attention_dynamic_masked(%arg0: tensor<?x?x?xf16>, %arg1: tens
 //  CHECK-SAME:       affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d4)>
 //  CHECK-SAME:       affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d5)>
 // CHECK-SAME:       ins(%[[QUERY]], %[[KEY]], %[[CACHE]], %[[ARG3]], %[[MASK]] :
+// CHECK-SAME:       outs(%[[EMPTY]] :
 //      CHECK:   util.return %[[ATTENTION]]
 
 // -----
@@ -687,23 +710,6 @@ util.func public @scatter_collapse_noop(%arg0: tensor<10xf16>, %arg1: tensor<10x
 
 // -----
 
-util.func public @scatter_collapse_multiple_dynamic(%arg0 : tensor<?x?x4x32x32xf16>, %arg1 : tensor<?xi64>, %arg2 : tensor<?x4x32x32xf16>) -> (tensor<?x4x32x32xf16>){
-  %collapsed = tensor.collapse_shape %arg0 [[0, 1], [2], [3], [4]] : tensor<?x?x4x32x32xf16> into tensor<?x4x32x32xf16>
-  %0 = iree_linalg_ext.scatter dimension_map = [0] unique_indices(true) ins(%collapsed, %arg1 : tensor<?x4x32x32xf16>, tensor<?xi64>) outs(%arg2 : tensor<?x4x32x32xf16>) {
-  ^bb0(%arg7: f16, %arg8: f16):
-    iree_linalg_ext.yield %arg7 : f16
-  } -> tensor<?x4x32x32xf16>
-  util.return %0 : tensor<?x4x32x32xf16>
-}
-// CHECK-LABEL: util.func public @scatter_collapse_multiple_dynamic
-//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]:
-//  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]:
-//  CHECK-SAME:     %[[ARG2:[a-zA-Z0-9]+]]:
-//       CHECK:   %[[SCATTER:.+]] = iree_linalg_ext.scatter
-//  CHECK-SAME:       ins({{.*}} : tensor<?x?x4x32x32xf16>, tensor<?x?xi64>
-
-// -----
-
 util.func public @gather_expand(%arg0: tensor<100x128xf16>, %arg1: tensor<10xi32>, %arg2: tensor<10x128xf16>) -> tensor<2x5x4x32xf16> {
   %c0 = arith.constant 0 : index
   %0 = iree_linalg_ext.gather dimension_map = [0] ins(%arg0, %arg1 : tensor<100x128xf16>, tensor<10xi32>) outs(%arg2 : tensor<10x128xf16>) -> tensor<10x128xf16>
