[LinalgExt] Implement unit dim folding pattern for map_scatter (#21563)

Implements unit dim folding on the source of
`iree_linalg_ext.map_scatter`. This is needed for the
LLVMGPUVectorDistribute pipeline to avoid unit dim slices creating extra
copies. The unit dim folding only happens on the source of the
map_scatter, because we currently only see the op during codegen, where
it is the last op in the dispatch. We generally shouldn't see unit dims
on the output of the map_scatter, since unit dims in this phase will
come from tiling, which only affects the input.

---------

Signed-off-by: Max Dawkins <[email protected]>

Author: Max191

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

@@ -575,6 +575,55 @@ Value rankExpandValue(RewriterBase &rewriter, Location loc, Value destVal,
   }
 }
 
+struct DropMapScatterUnitDims final : public OpRewritePattern<MapScatterOp> {
+  using OpRewritePattern<MapScatterOp>::OpRewritePattern;
+  DropMapScatterUnitDims(MLIRContext *context,
+                         linalg::ControlDropUnitDims options,
+                         PatternBenefit benefit = 1)
+      : OpRewritePattern<MapScatterOp>(context, benefit),
+        options(std::move(options)) {}
+
+  LogicalResult matchAndRewrite(MapScatterOp mapScatterOp,
+                                PatternRewriter &rewriter) const override {
+    auto inputType = dyn_cast<RankedTensorType>(mapScatterOp.getInputType());
+    if (!inputType) {
+      return failure();
+    }
+    Location loc = mapScatterOp.getLoc();
+    FailureOr<Value> newInput = rankReduceOperand(
+        rewriter, loc, /*startDim=*/0, /*numDims=*/mapScatterOp.getInputRank(),
+        mapScatterOp.getInput(), mapScatterOp.getInputType(), options);
+    if (failed(newInput)) {
+      return failure();
+    }
+
+    auto newInputType = cast<ShapedType>(newInput->getType());
+    auto unitFoldingBuilder = [&](ArrayRef<BlockArgument> nonUnitIndices) {
+      Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+      int nonUnitArgIdx = 0;
+      return llvm::map_to_vector(
+          llvm::seq<int64_t>(inputType.getRank()), [&](int64_t dim) -> Value {
+            return inputType.getDimSize(dim) == 1
+                       ? zero
+                       : cast<Value>(nonUnitIndices[nonUnitArgIdx++]);
+          });
+    };
+    // The map_scatter op is generally only used in Codegen, where it is the
+    // last op in the dispatch, so for now, we don't bother collapsing the
+    // result shape and inserting an expansion after the op.
+    rewriter.modifyOpInPlace(mapScatterOp, [&]() {
+      mapScatterOp.getInputMutable().assign(newInput.value());
+      mapScatterOp.insertTransformationAtStart(
+          rewriter, unitFoldingBuilder,
+          /*numSourceIndices=*/newInputType.getRank());
+    });
+    return success();
+  }
+
+private:
+  linalg::ControlDropUnitDims options;
+};
+
 struct DropGatherUnitDims final : public OpRewritePattern<GatherOp> {
   DropGatherUnitDims(MLIRContext *context, linalg::ControlDropUnitDims options,
                      PatternBenefit benefit = 1)
@@ -1059,8 +1108,8 @@ struct DropAttentionUnitDims final
 void populateFoldUnitExtentDimsPatterns(
     RewritePatternSet &patterns, const linalg::ControlDropUnitDims &options) {
   patterns.add<DropScatterUnitIndexDepth>(patterns.getContext());
-  patterns.add<DropGatherUnitDims, DropScatterUnitDims, DropAttentionUnitDims>(
-      patterns.getContext(), options);
+  patterns.add<DropGatherUnitDims, DropScatterUnitDims, DropAttentionUnitDims,
+               DropMapScatterUnitDims>(patterns.getContext(), options);
 }
 
 } // namespace mlir::iree_compiler::IREE::LinalgExt

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

@@ -132,3 +132,63 @@ util.func public @scatter_no_change_output(%slice: tensor<1x2xf16>, %indices: te
 //       SLICE:   iree_linalg_ext.scatter
 //  SLICE-SAME:     ins(%[[UPDATE_SLICE]], %[[INDICES_SLICE]]
 //  SLICE-SAME:     outs(%[[ORIGINAL]]
+
+// -----
+
+util.func public @map_scatter(%input: tensor<1x2x1x1xf16>) -> tensor<2x2x2x2xf16> {
+  %empty = tensor.empty() : tensor<2x2x2x2xf16>
+  %0 = iree_linalg_ext.map_scatter %input into %empty {
+    ^bb0(%idx0: index, %idx1: index, %idx2: index, %idx3: index):
+      %mask = arith.constant true
+      iree_linalg_ext.yield %idx0, %idx1, %idx2, %idx3, %mask : index, index, index, index, i1
+  } : tensor<1x2x1x1xf16> into tensor<2x2x2x2xf16> -> tensor<2x2x2x2xf16>
+  util.return %0 : tensor<2x2x2x2xf16>
+}
+// RESHAPE-LABEL: util.func public @map_scatter
+//  RESHAPE-SAME:   %[[INPUT:[a-zA-Z0-9]+]]: tensor<1x2x1x1xf16>
+//   RESHAPE-DAG:   %[[DEST:.+]] = tensor.empty() : tensor<2x2x2x2xf16>
+//   RESHAPE-DAG:   %[[C0:.+]] = arith.constant 0 : index
+//   RESHAPE-DAG:   %[[INPUT_COLLAPSE:.+]] = tensor.collapse_shape %[[INPUT]]
+//  RESHAPE-SAME:     tensor<1x2x1x1xf16> into tensor<2xf16>
+//       RESHAPE:   iree_linalg_ext.map_scatter %[[INPUT_COLLAPSE]] into %[[DEST]]
+//       RESHAPE:     ^bb0(%[[IDX:.+]]: index):
+//       RESHAPE:       iree_linalg_ext.yield  %[[C0]], %[[IDX]], %[[C0]], %[[C0]]
+
+// SLICE-LABEL: util.func public @map_scatter
+//  SLICE-SAME:   %[[INPUT:[a-zA-Z0-9]+]]: tensor<1x2x1x1xf16>
+//   SLICE-DAG:   %[[DEST:.+]] = tensor.empty() : tensor<2x2x2x2xf16>
+//   SLICE-DAG:   %[[INPUT_SLICE:.+]] = tensor.extract_slice %[[INPUT]]
+//  SLICE-SAME:     tensor<1x2x1x1xf16> to tensor<2xf16>
+//       SLICE:   iree_linalg_ext.map_scatter %[[INPUT_SLICE]] into %[[DEST]]
+//       SLICE:     ^bb0(%[[IDX:.+]]: index):
+//       SLICE:       %[[C0:.+]] = arith.constant 0 : index
+//       SLICE:       iree_linalg_ext.yield  %[[C0]], %[[IDX]], %[[C0]], %[[C0]]
+
+// -----
+
+util.func public @map_scatter_all_unit(%input: tensor<1x1xf16>) -> tensor<2x2xf16> {
+  %empty = tensor.empty() : tensor<2x2xf16>
+  %0 = iree_linalg_ext.map_scatter %input into %empty {
+    ^bb0(%idx0: index, %idx1: index):
+      %mask = arith.constant true
+      iree_linalg_ext.yield %idx0, %idx1, %mask : index, index, i1
+  } : tensor<1x1xf16> into tensor<2x2xf16> -> tensor<2x2xf16>
+  util.return %0 : tensor<2x2xf16>
+}
+// RESHAPE-LABEL: util.func public @map_scatter_all_unit
+//  RESHAPE-SAME:   %[[INPUT:[a-zA-Z0-9]+]]: tensor<1x1xf16>
+//   RESHAPE-DAG:   %[[DEST:.+]] = tensor.empty() : tensor<2x2xf16>
+//   RESHAPE-DAG:   %[[C0:.+]] = arith.constant 0 : index
+//   RESHAPE-DAG:   %[[INPUT_COLLAPSE:.+]] = tensor.collapse_shape %[[INPUT]]
+//  RESHAPE-SAME:     tensor<1x1xf16> into tensor<1xf16>
+//       RESHAPE:   iree_linalg_ext.map_scatter %[[INPUT_COLLAPSE]] into %[[DEST]]
+//       RESHAPE:       iree_linalg_ext.yield  %[[C0]], %[[C0]]
+
+// SLICE-LABEL: util.func public @map_scatter_all_unit
+//  SLICE-SAME:   %[[INPUT:[a-zA-Z0-9]+]]: tensor<1x1xf16>
+//   SLICE-DAG:   %[[DEST:.+]] = tensor.empty() : tensor<2x2xf16>
+//   SLICE-DAG:   %[[INPUT_SLICE:.+]] = tensor.extract_slice %[[INPUT]]
+//  SLICE-SAME:     tensor<1x1xf16> to tensor<1xf16>
+//       SLICE:   iree_linalg_ext.map_scatter %[[INPUT_SLICE]] into %[[DEST]]
+//       SLICE:       %[[C0:.+]] = arith.constant 0 : index
+//       SLICE:       iree_linalg_ext.yield  %[[C0]], %[[C0]]