Keep masked attention finalization reciprocal form

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

Author: keshavvinayak01

compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp

@@ -476,8 +476,14 @@ FailureOr<SmallVector<Value>> AttentionOp::decomposeOperation(OpBuilder &b) {
 
   Value fullyMaskedRows;
   if (mask != nullptr) {
-    fullyMaskedRows =
-        createFullyMaskedRowsFromScores(b, loc, sMap, maxMap, rowRedSize, s);
+    Type maskElementType = getElementTypeOrSelf(mask.value().getType());
+    if (isa<IntegerType>(maskElementType)) {
+      fullyMaskedRows = createFullyMaskedRowsFromMask(
+          b, loc, *getMaskMap(), maxMap, rowRedSize, mask.value());
+    } else {
+      fullyMaskedRows =
+          createFullyMaskedRowsFromScores(b, loc, sMap, maxMap, rowRedSize, s);
+    }
   }
 
   // max = rowMax(S)

compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/decompose_aggregate_op.mlir

@@ -185,9 +185,9 @@ func.func @attention_f16_masked(%query: tensor<192x1024x64xf16>,
 // CHECK: linalg.generic
 // CHECK:   arith.addf
 // CHECK:   linalg.yield
-// masked_rows = rowAll(isneginf(S))
+// masked_rows = rowAll(!mask)
 // CHECK: linalg.generic
-// CHECK:   arith.cmpf oeq
+// CHECK:   arith.xori
 // CHECK:   arith.andi
 // CHECK:   linalg.yield
 // max = rowMax(S)

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

@@ -127,35 +127,34 @@ void convertToOnlineAttention(IREE::LinalgExt::AttentionOp attnOp,
 
   // Finalize online attention. With a mask, fully-masked rows can have
   // `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.
+  // produce 0 instead of NaN. Keep the normal reciprocal-multiply shape so the
+  // common path stays as close as possible to unmasked finalization.
 
   // Compress the indexing maps.
-  SmallVector<AffineMap> compressedMaps =
-      compressUnusedDims(SmallVector<AffineMap>{sumMap, accMap, accMap});
+  SmallVector<Value> finalizeInputs = {sum, x};
+  SmallVector<AffineMap> finalizeMaps = {sumMap, accMap, accMap};
+
+  SmallVector<AffineMap> compressedMaps = compressUnusedDims(finalizeMaps);
 
   SmallVector<utils::IteratorType> iteratorTypes(compressedMaps[0].getNumDims(),
                                                  utils::IteratorType::parallel);
 
   auto genericOp = linalg::GenericOp::create(
-      rewriter, loc, attnOp.getOutput().getType(), ValueRange{sum, x},
+      rewriter, loc, attnOp.getOutput().getType(), finalizeInputs,
       attnOp.getOutput(), compressedMaps, iteratorTypes,
       [&](OpBuilder &b, Location loc, ValueRange args) {
-        Value result;
+        Value one = arith::ConstantOp::create(
+            b, loc, b.getFloatAttr(args[0].getType(), 1.0));
+        Value reciprocal = arith::DivFOp::create(b, loc, one, args[0]);
+        // Both sum and x are in fp32, as created earlier, so we only need to
+        // cast after the mul.
+        Value result = arith::MulFOp::create(b, loc, reciprocal, args[1]);
         if (hasMask) {
-          result = arith::DivFOp::create(b, loc, args[1], args[0]);
           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));
-          Value reciprocal = arith::DivFOp::create(b, loc, one, args[0]);
-          // Both sum and x are in fp32, as created earlier, so we only need to
-          // cast after the mul.
-          result = arith::MulFOp::create(b, loc, reciprocal, args[1]);
         }
         // Cast result to the required type by attention output.
         result = convertScalarToDtype(b, loc, result, args[2].getType(),

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

@@ -66,11 +66,13 @@ func.func @masked_attention(%q: tensor<2x10x4096x128xf16>, %k: tensor<2x10x4096x
 
 // CHECK-LABEL: func.func @masked_attention
 // CHECK-SAME: %[[MASK:.+]]: tensor<2x10x4096x4096xi1>
-// Masked: keep the fully-masked row guard fused with finalization.
+// Masked: keep the same reciprocal-multiply finalization shape and guard the
+// fully-masked row case.
 // CHECK: %[[OUT:.+]]:3 = iree_linalg_ext.online_attention
 // CHECK: linalg.generic
 // CHECK-SAME: ins(%[[OUT]]#2, %[[OUT]]#0
 // CHECK: arith.divf
+// CHECK: arith.mulf
 // CHECK: arith.cmpf oeq
 // CHECK: arith.select
 // CHECK: arith.truncf

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

@@ -139,6 +139,22 @@ 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) {
+        if (maskValue.getType().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 zeroFullyMaskedRows(OpBuilder &builder, Location loc, AffineMap valueMap,
                           AffineMap rowMap, Value value,
                           Value fullyMaskedRows) {

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

@@ -50,6 +50,13 @@ Value createFullyMaskedRowsFromScores(OpBuilder &builder, Location loc,
                                       ArrayRef<OpFoldResult> rowSizes,
                                       Value scores);
 
+/// Compute a row predicate for safe masked-softmax finalization by checking
+/// whether every integer mask element in a softmax row is false.
+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);