[DispatchCreation] Limit attention broadcast fusion (#24127)

Reject fusions that would broadcast away an entire attention dim group.



Closes #24051

Signed-off-by: Ian Wood <ianwood@u.northwestern.edu>

Author: IanWood1

compiler/src/iree/compiler/DispatchCreation/ElementwiseOpFusion.cpp

@@ -12,7 +12,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
+#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
 #include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "iree/compiler/Dialect/LinalgExt/Utils/IndexingUtils.h"
 #include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
 #include "iree/compiler/DispatchCreation/FusionUtils.h"
 #include "iree/compiler/DispatchCreation/Passes.h"
@@ -47,6 +49,56 @@ static SmallVector<T> applyProjectedPermutation(const SmallVectorImpl<T> &input,
   return result;
 }
 
+static bool wouldBroadcastAwayAttentionDimGroup(OpOperand *fusedOperand) {
+  auto attentionOp =
+      dyn_cast<IREE::LinalgExt::AttentionOp>(fusedOperand->getOwner());
+  if (!attentionOp) {
+    return false;
+  }
+
+  auto producer = dyn_cast_if_present<linalg::LinalgOp>(
+      fusedOperand->get().getDefiningOp());
+  if (!producer || !IREE::LinalgExt::isBroadcastingOp(producer)) {
+    return false;
+  }
+
+  FailureOr<IREE::LinalgExt::AttentionOpDetail> maybeOpInfo =
+      IREE::LinalgExt::AttentionOpDetail::get(
+          attentionOp.getQueryMap(), attentionOp.getKeyMap(),
+          attentionOp.getValueMap(), attentionOp.getOutputMap());
+  if (failed(maybeOpInfo)) {
+    return true;
+  }
+  IREE::LinalgExt::AttentionOpDetail opInfo = std::move(*maybeOpInfo);
+
+  AffineMap producerInputMap =
+      producer.getMatchingIndexingMap(producer.getDpsInputOperand(0));
+  AffineMap producerResultMap =
+      producer.getMatchingIndexingMap(producer.getDpsInitOperand(0));
+  AffineMap consumerInputMap = attentionOp.getMatchingIndexingMap(fusedOperand);
+  AffineMap fusedInputMap =
+      producerInputMap.compose(inversePermutation(producerResultMap))
+          .compose(consumerInputMap);
+
+  auto broadcastsAwayDimGroup = [&](ArrayRef<int64_t> dims) {
+    auto usesAnyDim = [](AffineMap map, ArrayRef<int64_t> dims) {
+      return llvm::any_of(map.getResults(), [&](AffineExpr expr) {
+        auto dimExpr = dyn_cast<AffineDimExpr>(expr);
+        return dimExpr && llvm::any_of(dims, [&](int64_t dim) {
+                 return dimExpr.getPosition() == dim;
+               });
+      });
+    };
+    return !dims.empty() && usesAnyDim(consumerInputMap, dims) &&
+           !usesAnyDim(fusedInputMap, dims);
+  };
+
+  SmallVector<ArrayRef<int64_t>> importantDimGroups = {
+      opInfo.getBatchDims(), opInfo.getMDims(), opInfo.getK1Dims(),
+      opInfo.getK2Dims(), opInfo.getNDims()};
+  return llvm::any_of(importantDimGroups, broadcastsAwayDimGroup);
+}
+
 //===----------------------------------------------------------------------===//
 // GatherFusionPattern
 //===----------------------------------------------------------------------===//
@@ -192,7 +244,13 @@ void ElementwiseOpFusionPass::runOnOperation() {
     Operation *producer = fusedOperand->get().getDefiningOp();
     Operation *consumer = fusedOperand->getOwner();
 
-    return IREE::Flow::isNonNullAndOutsideDispatch({producer, consumer});
+    if (!IREE::Flow::isNonNullAndOutsideDispatch({producer, consumer})) {
+      return false;
+    }
+    if (wouldBroadcastAwayAttentionDimGroup(fusedOperand)) {
+      return false;
+    }
+    return true;
   };
   RewritePatternSet linalgExtFusionPatterns(context);
   IREE::LinalgExt::populateFuseLinalgExtOpsWithTransposes(

compiler/src/iree/compiler/DispatchCreation/test/elementwise_op_fusion.mlir

@@ -284,6 +284,48 @@ util.func public @fuse_attention_with_broadcast(%arg0: tensor<4x8x128x?xf16>, %a
 //  CHECK-SAME:       ins(%[[ARG1]], %[[ARG2]], %[[ARG0]], %[[ARG3]], %[[ARG4]] :
 //       CHECK:   util.return %[[ATTENTION]]
 
+// -----
+
+util.func public @dont_fuse_attention_with_broadcasted_away_n_dim(
+    %q: tensor<4x32x64x16xf16>,
+    %k: tensor<4x32x64x16xf16>,
+    %v: tensor<4x32x64xf16>,
+    %scale: f16) -> tensor<4x32x64x128xf16> {
+  %empty_v = tensor.empty() : tensor<4x32x64x128xf16>
+  %v_broadcast = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>,
+                       affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
+      iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+      ins(%v : tensor<4x32x64xf16>) outs(%empty_v : tensor<4x32x64x128xf16>) {
+  ^bb0(%in: f16, %out: f16):
+    linalg.yield %in : f16
+  } -> tensor<4x32x64x128xf16>
+  %empty_out = tensor.empty() : tensor<4x32x64x128xf16>
+  %attention = iree_linalg_ext.attention {
+      indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>,
+                       affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d4, d3)>,
+                       affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d4, d5)>,
+                       affine_map<(d0, d1, d2, d3, d4, d5) -> ()>,
+                       affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d5)>]}
+      ins(%q, %k, %v_broadcast, %scale :
+          tensor<4x32x64x16xf16>, tensor<4x32x64x16xf16>,
+          tensor<4x32x64x128xf16>, f16)
+      outs(%empty_out : tensor<4x32x64x128xf16>) {
+  ^bb0(%score: f16):
+    iree_linalg_ext.yield %score : f16
+  } -> tensor<4x32x64x128xf16>
+  util.return %attention : tensor<4x32x64x128xf16>
+}
+// CHECK-LABEL: func public @dont_fuse_attention_with_broadcasted_away_n_dim
+//  CHECK-SAME:     %[[Q:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[K:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[V:[a-zA-Z0-9]+]]:
+//  CHECK-SAME:     %[[SCALE:[a-zA-Z0-9]+]]:
+//       CHECK:   %[[V_BCAST:.+]] = linalg.generic
+//       CHECK:   %[[ATTENTION:.+]] = iree_linalg_ext.attention
+//  CHECK-SAME:     ins(%[[Q]], %[[K]], %[[V_BCAST]], %[[SCALE]] :
+//       CHECK:   util.return %[[ATTENTION]]
+
 
 // -----