[Common] Update CombineSourceLayoutTransform pass to fold broadcast

-- This commit adds folding of linalg.generic broadcast op into
   MapGather op.
-- It also improves the folding algo to make sure that an indentity
   MapGather (and consequentially a folding trigger) is inserted only
   if the relayout chain originating from LoadFromBufferOp is "complex".
-- A "complex" relayout chain in this regard can be :-
   i. A chain of length >= 2.
  ii. Or, a chain of length == 1 but having a supported linalg relayout op.
     This is done to prevent creating MapGatherOp for simple primitives
     like tensor.extract_slice since it unneccesarily ends up creating an
     empty tensor for MapGather's destination -> which in turn leads to
     creation of big memref.alloca ops later in the pipeline causing stack
     size limit issue.

Signed-off-by: Abhishek Varma <abhvarma@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: Abhishek-Varma

.pre-commit-config.yaml

@@ -41,10 +41,11 @@ repos:
         exclude_types: ["jupyter"]
 
   - repo: https://github.com/igorshubovych/markdownlint-cli
-    rev: v0.47.0
+    rev: v0.39.0
     hooks:
       - id: markdownlint
         name: Run markdownlint on .md files
+        language_version: "system"
         args: ["--config", "docs/.markdownlint.yml"]
         files: "docs/website/.*.md"
         exclude: "mlir-dialects/!(index).md"

compiler/src/iree/compiler/Codegen/Common/CombineLayoutTransformation.cpp

@@ -13,6 +13,7 @@
 #include "llvm/Support/DebugLog.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/Dialect/UB/IR/UBOps.h"
@@ -532,9 +533,46 @@ static MapStoreOp insertIdentityMapStore(RewriterBase &rewriter,
 }
 
 bool isSupportedSingleInputRelayoutOp(Operation *op) {
-  return isa<tensor::ExpandShapeOp, tensor::CollapseShapeOp,
-             tensor::ExtractSliceOp, tensor::PadOp, linalg::CopyOp,
-             linalg::TransposeOp>(op);
+  if (isa<tensor::ExpandShapeOp, tensor::CollapseShapeOp,
+          tensor::ExtractSliceOp, tensor::PadOp, linalg::CopyOp,
+          linalg::TransposeOp>(op)) {
+    return true;
+  }
+  auto genericOp = dyn_cast<linalg::GenericOp>(op);
+  return genericOp && linalg::isaBroadcastOpInterface(genericOp).has_value();
+}
+
+/// Returns true if the relayout op starts a "complex" chain.
+/// A "complex" chain is :-
+/// - a chain of relayout ops with length >= 2,
+/// - or a chain of length 1 with one of the supported linalg relayout ops.
+static bool isComplexRelayoutChain(Operation *relayoutOp) {
+  assert(isSupportedSingleInputRelayoutOp(relayoutOp) &&
+         "expected a supported relayout op");
+  Value result = relayoutOp->getResult(0);
+  bool hasRelayoutUser = llvm::any_of(result.getUsers(), [](Operation *user) {
+    return isSupportedSingleInputRelayoutOp(user);
+  });
+  // Chain length >= 2 -> complex.
+  if (hasRelayoutUser) {
+    return true;
+  }
+  // Chain length 1: complex only if the op is a linalg op.
+  return isa<linalg::LinalgOp>(relayoutOp);
+}
+
+/// Collects direct relayout op users of `loadResult` that start a complex
+/// relayout chain.
+static SmallPtrSet<Operation *, 4>
+getComplexChainRelayoutUsers(Value loadResult) {
+  SmallPtrSet<Operation *, 4> complexUsers;
+  for (Operation *user : loadResult.getUsers()) {
+    if (isSupportedSingleInputRelayoutOp(user) &&
+        user->getOperand(0) == loadResult && isComplexRelayoutChain(user)) {
+      complexUsers.insert(user);
+    }
+  }
+  return complexUsers;
 }
 
 // This is only desirable in the dispatch scope but not in the workgroup scope.
@@ -983,7 +1021,31 @@ foldExtractSliceIntoMapLoad(RewriterBase &rewriter,
                                      indexTransformBuilder);
 }
 
-/// Fold a consumer `padOp` into a producer `mapLoadOp`.
+/// Fold a consumer broadcast `linalg.generic` into a producer `map_load`.
+static FailureOr<MapLoadOp> foldBroadcastGenericIntoMapLoad(
+    RewriterBase &rewriter, linalg::GenericOp genericOp, MapLoadOp mapLoadOp) {
+  assert(genericOp.getDpsInputs()[0] == mapLoadOp.getResult(0) &&
+         "expected map_load to be the producer of genericOp input");
+  if (!linalg::isaBroadcastOpInterface(genericOp).has_value()) {
+    return rewriter.notifyMatchFailure(genericOp,
+                                       "generic op is not a broadcast");
+  }
+
+  AffineMap inputMap = genericOp.getIndexingMapsArray()[0];
+  return foldConsumerIntoMapLoadImpl(
+      rewriter, genericOp, mapLoadOp,
+      [inputMap](ArrayRef<BlockArgument> indices) -> SmallVector<Value> {
+        SmallVector<Value> sourceIndices;
+        sourceIndices.reserve(inputMap.getNumResults());
+        for (AffineExpr expr : inputMap.getResults()) {
+          unsigned pos = cast<AffineDimExpr>(expr).getPosition();
+          sourceIndices.push_back(indices[pos]);
+        }
+        return sourceIndices;
+      });
+}
+
+/// Fold a consumer `padOp` into a producer `mapGatherOp`.
 /// Index transformation: source_idx = new_idx - low_pad
 /// Fill value is set to the pad value.
 static FailureOr<MapLoadOp> foldPadIntoMapLoad(RewriterBase &rewriter,
@@ -1050,6 +1112,9 @@ FailureOr<MapLoadOp> foldIntoMapLoad(RewriterBase &rewriter, Operation *op,
       .Case<tensor::PadOp>([&](tensor::PadOp padOp) {
         return foldPadIntoMapLoad(rewriter, padOp, mapLoadOp);
       })
+      .Case<linalg::GenericOp>([&](linalg::GenericOp genericOp) {
+        return foldBroadcastGenericIntoMapLoad(rewriter, genericOp, mapLoadOp);
+      })
       .Default([](Operation *) { return failure(); });
 }
 
@@ -1078,12 +1143,18 @@ struct FoldConsumerRelayoutIntoMapLoadPattern
   }
 };
 
-// Insert identity map_load op after the root and replace uses.
-static MapLoadOp insertIdentityMapLoad(RewriterBase &rewriter, OpResult root) {
+// Insert identity map_gather op after the root and replace only uses whose
+// owner is in `complexChainUsers` (i.e. uses that are part of a complex
+// relayout chain). Other uses keep using the load/root directly.
+static MapLoadOp
+insertIdentityMapLoad(RewriterBase &rewriter, OpResult root,
+                      const SmallPtrSetImpl<Operation *> &complexChainUsers) {
   Location loc = root.getLoc();
   SetVector<OpOperand *> originalUses;
   for (OpOperand &use : root.getUses()) {
-    originalUses.insert(&use);
+    if (complexChainUsers.contains(use.getOwner())) {
+      originalUses.insert(&use);
+    }
   }
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPointAfterValue(root);
@@ -1109,12 +1180,11 @@ struct InsertMapLoadOpPattern
 
   LogicalResult matchAndRewrite(IREE::Codegen::LoadFromBufferOp loadOp,
                                 PatternRewriter &rewriter) const override {
-    // Check if the load has at least one relayout op user.
-    bool hasRelayoutUser =
-        llvm::any_of(loadOp->getUsers(), [](Operation *user) {
-          return isSupportedSingleInputRelayoutOp(user);
-        });
-    if (!hasRelayoutUser) {
+    Value loadResult = loadOp.getResult();
+    SmallPtrSet<Operation *, 4> complexChainUsers =
+        getComplexChainRelayoutUsers(loadResult);
+    // Only introduce map_gather when there is at least one complex chain.
+    if (complexChainUsers.empty()) {
       return failure();
     }
     // Check that the load doesn't already have a map_load user (avoid
@@ -1125,7 +1195,8 @@ struct InsertMapLoadOpPattern
     if (hasMapLoadUser) {
       return failure();
     }
-    (void)insertIdentityMapLoad(rewriter, loadOp->getResult(0));
+    (void)insertIdentityMapLoad(rewriter, cast<OpResult>(loadResult),
+                                complexChainUsers);
     return success();
   }
 };

compiler/src/iree/compiler/Codegen/Common/test/combine_source_layout_transformation.mlir

@@ -28,16 +28,9 @@ func.func @fold_expand_shape(%buffer : memref<8x16xf32>) -> tensor<2x4x16xf32> {
 }
 // CHECK-LABEL: @fold_expand_shape
 //  CHECK-SAME:   %[[BUFFER:[a-zA-Z0-9_]+]]
-//       CHECK:   %[[SOURCE:.+]] = iree_codegen.load_from_buffer %[[BUFFER]]
-//       CHECK:   %[[DEST:.+]] = tensor.empty() : tensor<2x4x16xf32>
-//   CHECK-NOT:   tensor.expand_shape
-//       CHECK:   %[[MAP_GATHER:.+]] = iree_linalg_ext.map_load
-//  CHECK-SAME:     %[[SOURCE]] into %[[DEST]] {
-//  CHECK-NEXT:   ^bb0(%[[IDX0:.+]]: index, %[[IDX1:.+]]: index, %[[IDX2:.+]]: index):
-//       CHECK:     %[[LINEARIZE:.+]] = affine.linearize_index
-//  CHECK-SAME:       [%[[IDX0]], %[[IDX1]]] by (2, 4)
-//       CHECK:     iree_linalg_ext.yield %[[LINEARIZE]], %[[IDX2]],
-//       CHECK:   } : tensor<8x16xf32> into tensor<2x4x16xf32> -> tensor<2x4x16xf32>
+//       CHECK:   iree_codegen.load_from_buffer %[[BUFFER]]
+//       CHECK:   tensor.expand_shape
+//   CHECK-NOT:   iree_linalg_ext.map_load
 
 // -----
 
@@ -48,15 +41,9 @@ func.func @fold_collapse_shape(%buffer : memref<2x4x16xf32>) -> tensor<8x16xf32>
 }
 // CHECK-LABEL: @fold_collapse_shape
 //  CHECK-SAME:   %[[BUFFER:[a-zA-Z0-9_]+]]
-//       CHECK:   %[[SOURCE:.+]] = iree_codegen.load_from_buffer %[[BUFFER]]
-//       CHECK:   %[[DEST:.+]] = tensor.empty() : tensor<8x16xf32>
-//   CHECK-NOT:   tensor.collapse_shape
-//       CHECK:   %[[MAP_GATHER:.+]] = iree_linalg_ext.map_load
-//  CHECK-SAME:     %[[SOURCE]] into %[[DEST]] {
-//  CHECK-NEXT:   ^bb0(%[[IDX0:.+]]: index, %[[IDX1:.+]]: index):
-//       CHECK:     %[[DELINEARIZE:.+]]:2 = affine.delinearize_index %[[IDX0]] into (2, 4)
-//       CHECK:     iree_linalg_ext.yield %[[DELINEARIZE]]#0, %[[DELINEARIZE]]#1, %[[IDX1]],
-//       CHECK:   } : tensor<2x4x16xf32> into tensor<8x16xf32> -> tensor<8x16xf32>
+//       CHECK:   iree_codegen.load_from_buffer %[[BUFFER]]
+//       CHECK:   tensor.collapse_shape
+//   CHECK-NOT:   iree_linalg_ext.map_load
 
 // -----
 
@@ -67,16 +54,9 @@ func.func @fold_extract_slice(%buffer : memref<64xf32>) -> tensor<16xf32> {
 }
 // CHECK-LABEL: @fold_extract_slice
 //  CHECK-SAME:   %[[BUFFER:[a-zA-Z0-9_]+]]
-//   CHECK-DAG:   %[[C8:.+]] = arith.constant 8 : index
 //       CHECK:   %[[SOURCE:.+]] = iree_codegen.load_from_buffer %[[BUFFER]]
-//       CHECK:   %[[DEST:.+]] = tensor.empty() : tensor<16xf32>
-//   CHECK-NOT:   tensor.extract_slice
-//       CHECK:   %[[MAP_GATHER:.+]] = iree_linalg_ext.map_load
-//  CHECK-SAME:     %[[SOURCE]] into %[[DEST]] {
-//  CHECK-NEXT:   ^bb0(%[[IDX0:.+]]: index):
-//       CHECK:     %[[NEW_IDX:.+]] = arith.addi %[[IDX0]], %[[C8]] overflow<nsw>
-//       CHECK:     iree_linalg_ext.yield %[[NEW_IDX]],
-//       CHECK:   } : tensor<64xf32> into tensor<16xf32> -> tensor<16xf32>
+//       CHECK:   tensor.extract_slice %[[SOURCE]][8] [16] [1]
+//   CHECK-NOT:   iree_linalg_ext.map_load
 
 // -----
 
@@ -104,7 +84,6 @@ func.func @fold_copy_transpose(%buffer : memref<4x16xf32>) -> tensor<16x4xf32> {
 
 // -----
 
-// Low padding is [0, 0, 0], so indices are passed through unchanged due to subi with 0.
 func.func @fold_pad_with_zero_low_padding_offsets(%buffer : memref<1x50x64xf32>) -> tensor<1x64x64xf32> {
   %cst = arith.constant 0.000000e+00 : f32
   %source = iree_codegen.load_from_buffer %buffer : memref<1x50x64xf32> -> tensor<1x50x64xf32>
@@ -116,15 +95,9 @@ func.func @fold_pad_with_zero_low_padding_offsets(%buffer : memref<1x50x64xf32>)
 }
 // CHECK-LABEL: @fold_pad_with_zero_low_padding_offsets
 //  CHECK-SAME:   %[[BUFFER:[a-zA-Z0-9_]+]]
-//   CHECK-DAG:   %[[CST:.+]] = arith.constant 0.000000e+00 : f32
-//       CHECK:   %[[SOURCE:.+]] = iree_codegen.load_from_buffer %[[BUFFER]]
-//       CHECK:   %[[DEST:.+]] = tensor.empty() : tensor<1x64x64xf32>
-//   CHECK-NOT:   tensor.pad
-//       CHECK:   %[[MAP_GATHER:.+]] = iree_linalg_ext.map_load
-//  CHECK-SAME:     %[[SOURCE]] into %[[DEST]] {
-//  CHECK-NEXT:   ^bb0(%[[IDX0:.+]]: index, %[[IDX1:.+]]: index, %[[IDX2:.+]]: index):
-//       CHECK:     iree_linalg_ext.yield %[[IDX0]], %[[IDX1]], %[[IDX2]], %[[CST]] :
-//       CHECK:   } : tensor<1x50x64xf32> into tensor<1x64x64xf32> -> tensor<1x64x64xf32>
+//       CHECK:   iree_codegen.load_from_buffer %[[BUFFER]]
+//       CHECK:   tensor.pad
+//   CHECK-NOT:   iree_linalg_ext.map_load
 
 // -----
 
@@ -139,19 +112,9 @@ func.func @fold_pad_with_non_zero_low_padding_offsets(%buffer : memref<8x16xf32>
 }
 // CHECK-LABEL: @fold_pad_with_non_zero_low_padding_offsets
 //  CHECK-SAME:   %[[BUFFER:[a-zA-Z0-9_]+]]
-//   CHECK-DAG:   %[[CST:.+]] = arith.constant 1.000000e+00 : f32
-//   CHECK-DAG:   %[[C1:.+]] = arith.constant 1 : index
-//   CHECK-DAG:   %[[C2:.+]] = arith.constant 2 : index
-//       CHECK:   %[[SOURCE:.+]] = iree_codegen.load_from_buffer %[[BUFFER]]
-//       CHECK:   %[[DEST:.+]] = tensor.empty() : tensor<10x20xf32>
-//   CHECK-NOT:   tensor.pad
-//       CHECK:   %[[MAP_GATHER:.+]] = iree_linalg_ext.map_load
-//  CHECK-SAME:     %[[SOURCE]] into %[[DEST]] {
-//  CHECK-NEXT:   ^bb0(%[[IDX0:.+]]: index, %[[IDX1:.+]]: index):
-//       CHECK:     %[[NEW_IDX0:.+]] = arith.subi %[[IDX0]], %[[C1]] overflow<nsw> : index
-//       CHECK:     %[[NEW_IDX1:.+]] = arith.subi %[[IDX1]], %[[C2]] overflow<nsw> : index
-//       CHECK:     iree_linalg_ext.yield %[[NEW_IDX0]], %[[NEW_IDX1]], %[[CST]] :
-//       CHECK:   } : tensor<8x16xf32> into tensor<10x20xf32> -> tensor<10x20xf32>
+//       CHECK:   iree_codegen.load_from_buffer %[[BUFFER]]
+//       CHECK:   tensor.pad
+//   CHECK-NOT:   iree_linalg_ext.map_load
 
 // -----
 
@@ -182,3 +145,54 @@ func.func @nested_pads_different_values(%buffer : memref<8x16xf32>) -> tensor<14
 // Second pad is NOT folded because the map_load already has a padding value.
 //       CHECK:   tensor.pad
 //       CHECK:     tensor.yield %[[CST1]] : f32
+
+// -----
+
+func.func @fold_broadcast_generic(%buffer : memref<2x3xf32>) -> tensor<2x3x4x5xf32> {
+  %source = iree_codegen.load_from_buffer %buffer : memref<2x3xf32> -> tensor<2x3xf32>
+  %init = tensor.empty() : tensor<2x3x4x5xf32>
+  %broadcast = linalg.generic {
+    indexing_maps = [
+      affine_map<(d0, d1, d2, d3) -> (d0, d1)>,
+      affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+    ],
+    iterator_types = ["parallel", "parallel", "parallel", "parallel"]
+  } ins(%source : tensor<2x3xf32>) outs(%init : tensor<2x3x4x5xf32>) {
+  ^bb0(%in: f32, %out: f32):
+    linalg.yield %in : f32
+  } -> tensor<2x3x4x5xf32>
+  return %broadcast : tensor<2x3x4x5xf32>
+}
+// CHECK-LABEL: @fold_broadcast_generic
+//  CHECK-SAME:   %[[BUFFER:[a-zA-Z0-9_]+]]
+//       CHECK:   %[[SOURCE:.+]] = iree_codegen.load_from_buffer %[[BUFFER]]
+//       CHECK:   %[[DEST:.+]] = tensor.empty() : tensor<2x3x4x5xf32>
+//   CHECK-NOT:   linalg.generic
+//       CHECK:   iree_linalg_ext.map_load
+//  CHECK-SAME:     %[[SOURCE]] into %[[DEST]] {
+//  CHECK-NEXT:   ^bb0(%[[IDX0:.+]]: index, %[[IDX1:.+]]: index, %[[IDX2:.+]]: index, %[[IDX3:.+]]: index):
+// Broadcast: output (d0,d1,d2,d3) reads from source at (d0,d1)
+//       CHECK:     iree_linalg_ext.yield %[[IDX0]], %[[IDX1]],
+//       CHECK:   } : tensor<2x3xf32> into tensor<2x3x4x5xf32> -> tensor<2x3x4x5xf32>
+
+// -----
+
+func.func @complex_relayout_chain(%buffer : memref<8x16xf32>) -> tensor<16x8xf32> {
+  %source = iree_codegen.load_from_buffer %buffer : memref<8x16xf32> -> tensor<8x16xf32>
+  %expanded = tensor.expand_shape %source [[0, 1], [2]] output_shape [2, 4, 16] : tensor<8x16xf32> into tensor<2x4x16xf32>
+  %collapsed = tensor.collapse_shape %expanded [[0, 1], [2]] : tensor<2x4x16xf32> into tensor<8x16xf32>
+  %init = tensor.empty() : tensor<16x8xf32>
+  %transposed = linalg.transpose ins(%collapsed : tensor<8x16xf32>) outs(%init : tensor<16x8xf32>) permutation = [1, 0]
+  return %transposed : tensor<16x8xf32>
+}
+// CHECK-LABEL: @complex_relayout_chain
+//  CHECK-SAME:   %[[BUFFER:[a-zA-Z0-9_]+]]
+//       CHECK:   iree_codegen.load_from_buffer %[[BUFFER]]
+//       CHECK:   tensor.empty() : tensor<16x8xf32>
+//   CHECK-NOT:   tensor.expand_shape
+//   CHECK-NOT:   tensor.collapse_shape
+//   CHECK-NOT:   linalg.transpose
+//       CHECK:   iree_linalg_ext.map_load {{.*}} into {{.*}} {
+//  CHECK-NEXT:   ^bb0(%[[IDX0:.+]]: index, %[[IDX1:.+]]: index):
+//  CHECK-NEXT:     iree_linalg_ext.yield %[[IDX1]], %[[IDX0]], {{.*}} : index, index, f32
+//       CHECK:   } : tensor<8x16xf32> into tensor<16x8xf32> -> tensor<16x8xf32>