Keep masked online attention finalization fused

Signed-off-by: Keshav Vinayak Jha <keshavvinayakjha@gmail.com>

Author: keshavvinayak01

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

@@ -8,7 +8,6 @@
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
 #include "iree/compiler/Dialect/LinalgExt/Transforms/Passes.h"
 #include "iree/compiler/Dialect/LinalgExt/Utils/IndexingUtils.h"
-#include "iree/compiler/Dialect/LinalgExt/Utils/Utils.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -124,40 +123,32 @@ void convertToOnlineAttention(IREE::LinalgExt::AttentionOp attnOp,
 
   Value x = onlineAttn.getResult(0);
   Value sum = onlineAttn.getResult(2);
-  Value fullyMaskedRows;
-  if (mask) {
-    fullyMaskedRows = createFullyMaskedRowsFromMask(
-        rewriter, loc, *attnOp.getMaskMap(), sumMap, rowRedSize, *mask);
-  }
+  bool hasMask = static_cast<bool>(mask);
 
   // Finalize online attention. With a mask, fully-masked rows can have
-  // `sum == 0` and `x == 0`; zero those rows after normalization to produce 0
-  // instead of NaN.
+  // `sum == 0` and `x == 0`; guard that case in the finalization loop to
+  // produce 0 instead of NaN. Keep this fused with finalization because the
+  // online-attention lowering path expects a single finalization consumer.
 
   // Compress the indexing maps.
-  SmallVector<Value> inputs = {sum, x};
-  SmallVector<AffineMap> indexingMapsForFinalization = {sumMap, accMap};
-  if (fullyMaskedRows) {
-    inputs.push_back(fullyMaskedRows);
-    indexingMapsForFinalization.push_back(sumMap);
-  }
-  indexingMapsForFinalization.push_back(accMap);
   SmallVector<AffineMap> compressedMaps =
-      compressUnusedDims(indexingMapsForFinalization);
+      compressUnusedDims(SmallVector<AffineMap>{sumMap, accMap, accMap});
 
   SmallVector<utils::IteratorType> iteratorTypes(compressedMaps[0].getNumDims(),
                                                  utils::IteratorType::parallel);
 
   auto genericOp = linalg::GenericOp::create(
-      rewriter, loc, attnOp.getOutput().getType(), inputs, attnOp.getOutput(),
-      compressedMaps, iteratorTypes,
+      rewriter, loc, attnOp.getOutput().getType(), ValueRange{sum, x},
+      attnOp.getOutput(), compressedMaps, iteratorTypes,
       [&](OpBuilder &b, Location loc, ValueRange args) {
         Value result;
-        if (fullyMaskedRows) {
+        if (hasMask) {
           result = arith::DivFOp::create(b, loc, args[1], args[0]);
-          Value zero = arith::ConstantOp::create(
-              b, loc, b.getFloatAttr(args[1].getType(), 0.0));
-          result = arith::SelectOp::create(b, loc, args[2], zero, result);
+          Value zero =
+              arith::ConstantOp::create(b, loc, b.getFloatAttr(f32Type, 0.0));
+          Value isZero = arith::CmpFOp::create(
+              b, loc, arith::CmpFPredicate::OEQ, args[0], zero);
+          result = arith::SelectOp::create(b, loc, isZero, zero, result);
         } else {
           Value one = arith::ConstantOp::create(
               b, loc, b.getFloatAttr(args[0].getType(), 1.0));
@@ -167,7 +158,7 @@ void convertToOnlineAttention(IREE::LinalgExt::AttentionOp attnOp,
           result = arith::MulFOp::create(b, loc, reciprocal, args[1]);
         }
         // Cast result to the required type by attention output.
-        result = convertScalarToDtype(b, loc, result, args.back().getType(),
+        result = convertScalarToDtype(b, loc, result, args[2].getType(),
                                       /*isUnsignedCast=*/false);
         linalg::YieldOp::create(b, loc, result);
       });

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

@@ -66,17 +66,12 @@ func.func @masked_attention(%q: tensor<2x10x4096x128xf16>, %k: tensor<2x10x4096x
 
 // CHECK-LABEL: func.func @masked_attention
 // CHECK-SAME: %[[MASK:.+]]: tensor<2x10x4096x4096xi1>
-// Masked: compute the fully-masked row predicate from the mask and use it to
-// zero those rows after normalization.
+// Masked: keep the fully-masked row guard fused with finalization.
 // CHECK: %[[OUT:.+]]:3 = iree_linalg_ext.online_attention
-// CHECK: %[[FULLY_MASKED:.+]] = linalg.generic
-// CHECK-SAME: ins(%[[MASK]]
-// CHECK: arith.xori
-// CHECK: arith.andi
-// CHECK: linalg.yield
 // CHECK: linalg.generic
-// CHECK-SAME: ins(%[[OUT]]#2, %[[OUT]]#0, %[[FULLY_MASKED]]
+// CHECK-SAME: ins(%[[OUT]]#2, %[[OUT]]#0
 // CHECK: arith.divf
+// CHECK: arith.cmpf oeq
 // CHECK: arith.select
 // CHECK: arith.truncf
 // CHECK: linalg.yield

compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp

@@ -139,31 +139,6 @@ Value createFullyMaskedRowsFromScores(OpBuilder &builder, Location loc,
       });
 }
 
-Value createFullyMaskedRowsFromMask(OpBuilder &builder, Location loc,
-                                    AffineMap maskMap, AffineMap rowMap,
-                                    ArrayRef<OpFoldResult> rowSizes,
-                                    Value mask) {
-  return createFullyMaskedRows(
-      builder, loc, maskMap, rowMap, rowSizes, mask,
-      [](OpBuilder &b, Location loc, Value maskValue) -> Value {
-        Type maskType = maskValue.getType();
-        if (maskType.isInteger()) {
-          if (maskType.getIntOrFloatBitWidth() != 1) {
-            maskValue =
-                arith::TruncIOp::create(b, loc, b.getI1Type(), maskValue);
-          }
-          Value trueValue =
-              arith::ConstantOp::create(b, loc, b.getBoolAttr(/*value=*/true));
-          return arith::XOrIOp::create(b, loc, maskValue, trueValue);
-        }
-        Value negInf = arith::ConstantOp::create(
-            b, loc,
-            b.getFloatAttr(maskType, -std::numeric_limits<double>::infinity()));
-        return arith::CmpFOp::create(b, loc, arith::CmpFPredicate::OEQ,
-                                     maskValue, negInf);
-      });
-}
-
 Value zeroFullyMaskedRows(OpBuilder &builder, Location loc, AffineMap valueMap,
                           AffineMap rowMap, Value value,
                           Value fullyMaskedRows) {

compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h

@@ -50,14 +50,6 @@ Value createFullyMaskedRowsFromScores(OpBuilder &builder, Location loc,
                                       ArrayRef<OpFoldResult> rowSizes,
                                       Value scores);
 
-/// Compute a row predicate for safe masked-softmax finalization directly from
-/// an attention mask. Integer masks are normalized to i1 and use `false`/0 as
-/// masked; floating-point masks use `-inf` as masked.
-Value createFullyMaskedRowsFromMask(OpBuilder &builder, Location loc,
-                                    AffineMap maskMap, AffineMap rowMap,
-                                    ArrayRef<OpFoldResult> rowSizes,
-                                    Value mask);
-
 /// Zero every element in rows whose row predicate is true.
 Value zeroFullyMaskedRows(OpBuilder &builder, Location loc, AffineMap valueMap,
                           AffineMap rowMap, Value value, Value fullyMaskedRows);