Add support for unstructured masks in MaskAnalysis (#333)

- Enhanced `MaskAnalysis` to handle unstructured masks. MaskState will
save unstructured mask for a dimension if that dimension failed
MaskAnalysis.
- Updated `TritonStructuredDialect` to include `gather_scatter_mask` in
`MakeGatherScatterTensorPtrOp`.
- Modified `PtrAnalysis` to apply unstructured masks during pointer
analysis.
- Extended `StructuredToMemref` conversion to support generic masks in
load/store operations.
- Added verification logic for `MakeGatherScatterTensorPtrOp` to ensure
compatibility with masks.

---------

Co-authored-by: Xiang Li <[email protected]>

Author: python3kgae

include/triton-shared/Analysis/MaskAnalysis.h

@@ -44,17 +44,36 @@ namespace triton {
 //
 // Example of creating 2D mask:
 //  mask = (rows[:, None] < M) & (cols[None, :] < N)
+//
+// Bool tensor mask could be saved into masks in case that dimension failed
+// MaskAnalysis. These is to allow case where only one dimension failed while
+// others passed. A MakeGatherScatterTensorPtrOp operation could be generated
+// for the failed dimension. Only 3 patterns are supported for this.
+// 1. offsets[:, None] < n where the offsets is 1d tensor.
+//    It will in pattern of expandDims -> broadcast -> cmp
+// 2. mask[:, None] where mask is 1d bool tensor.
+//    It will in pattern of cmp -> expandDims -> broadcast
+// 3. scalar_mask[:, None] where scalar mask is scalar bool.
+//    It will in pattern of splat -> expandDims -> broadcast
+// These 3 patterns are only about how a bool tensor was created from 1D or
+// scalar bool. How the 1D and scalar bool were created is not important for the
+// unstructured mask.
+// Only one tensor mask is allowed. If multiple dimensions have failed
+// MaskAnalysis, then MaskAnalysis will still fail on the current operation.
 struct MaskState {
   OpFoldResult start;
   OpFoldResult end;
   SmallVector<OpFoldResult> dims;
+  SmallVector<Value> masks;
   OpFoldResult scalar;
   const bool useUnsafeMask;
 
   void dump() const;
 
   MaskState(bool useUnsafeMask = false) : useUnsafeMask(useUnsafeMask) {}
 
+  SmallVector<std::pair<unsigned, Value>> getUnstructuredMasks();
+
   int64_t getRank() const { return dims.size(); }
 
   bool isEmpty() const { return getRank() == 0 && !scalar && !start && !end; }

include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td

@@ -159,6 +159,7 @@ def TTS_MakeGatherScatterTensorPtrOp
   // strides: The strides of the parent tensor, which means how much to increase the pointer
   //          by when moving by 1 element in a specific axis.
   // offsets: Offset of the block along each dimension from base.
+  // gather_scatter_mask: Optional bool mask for mask which failed MaskAnalysis.
   // result:  A tensor of pointers.
 
   let arguments = (ins TT_Ptr:$base,
@@ -168,19 +169,21 @@ def TTS_MakeGatherScatterTensorPtrOp
                        Variadic<Index>:$strides,
                        Variadic<Index>:$offsets,
                        DenseI64ArrayAttr:$static_strides,
-                       DenseI64ArrayAttr:$static_offsets);
+                       DenseI64ArrayAttr:$static_offsets,
+                       Optional<TT_BoolLike>:$gather_scatter_mask);
 
   let results = (outs TT_PtrLike:$result);
 
   let assemblyFormat = [{
     $base `to` `sizes` `` `:` $sizes
     `gather_scatter_dim` `` `:` $gather_scatter_dim
     `gather_scatter_offset` `` `:` $gather_scatter_offset
+    (`gather_scatter_mask` `` `:` $gather_scatter_mask^)?
     `` `,` `strides` `` `:`
     custom<DynamicIndexList>($strides, $static_strides)
     `` `,` `offsets` `` `:`
     custom<DynamicIndexList>($offsets, $static_offsets)
-    attr-dict `:`  type($gather_scatter_offset) type($base) `to` type($result)
+    attr-dict `:`  type($gather_scatter_offset) type($gather_scatter_mask) type($base) `to` type($result)
   }];
 
 
@@ -193,6 +196,15 @@ def TTS_MakeGatherScatterTensorPtrOp
       "ArrayRef<int64_t>":$sizes,
       "ArrayRef<OpFoldResult>":$strides,
       "ArrayRef<OpFoldResult>":$offsets)>,
+
+    OpBuilder<(ins
+      "Value":$base,
+      "Value":$gather_scatter_offset,
+      "Value":$gather_scatter_mask,
+      "int":$gather_scatter_dim,
+      "ArrayRef<int64_t>":$sizes,
+      "ArrayRef<OpFoldResult>":$strides,
+      "ArrayRef<OpFoldResult>":$offsets)>,
   ];
 
   let extraClassDeclaration = [{
@@ -213,9 +225,8 @@ def TTS_MakeGatherScatterTensorPtrOp
     }
   }];
 
-  // TODO
-  //let hasVerifier = 1;
-  //let hasCanonicalizer = 1;
+  let hasVerifier = 1;
+  let hasCanonicalizer = 0;
 }
 
 def TTS_GetStructuredStateOp : TTS_Op<"get_structured_state", [AttrSizedResultSegments, Pure]> {

lib/Analysis/MaskAnalysis.cpp

@@ -320,6 +320,11 @@ void MaskState::dump() const {
   llvm::dbgs() << "dims: ";
   for (auto dim : dims)
     llvm::dbgs() << "\t" << dim << "\n";
+  if (!masks.empty()) {
+    llvm::dbgs() << "masks: ";
+    for (auto mask : masks)
+      llvm::dbgs() << "\t" << mask << "\n";
+  }
   llvm::dbgs() << "\n";
 }
 
@@ -341,14 +346,96 @@ LogicalResult MaskState::parseAdd(arith::AddIOp addOp, const Location loc,
 LogicalResult MaskState::parseAnd(arith::AndIOp andOp, const Location loc,
                                   OpBuilder &builder) {
   assert(this->isEmpty());
-
+  bool isBoolOp = false;
+  unsigned rank = 1;
+  if (auto shapedType = dyn_cast<ShapedType>(andOp.getType())) {
+    isBoolOp = shapedType.getElementType().isInteger(1);
+    rank = shapedType.getRank();
+  }
   MaskState lhsState;
-  if (failed(lhsState.parse(andOp.getLhs(), loc, builder)))
+  LogicalResult lResult = lhsState.parse(andOp.getLhs(), loc, builder);
+  if (failed(lResult) && !isBoolOp) {
     return failure();
+  }
 
   MaskState rhsState;
-  if (failed(rhsState.parse(andOp.getRhs(), loc, builder)))
+  LogicalResult rResult = rhsState.parse(andOp.getRhs(), loc, builder);
+  if (failed(rResult) && !isBoolOp) {
     return failure();
+  }
+
+  if (isBoolOp) {
+    if (lhsState.masks.size() != rank) {
+      return failure();
+    }
+
+    if (lhsState.masks.size() != rhsState.masks.size()) {
+      return failure();
+    }
+
+    // merge the masks.
+    if (lhsState.masks.size() == rhsState.masks.size()) {
+      auto shapedType = cast<ShapedType>(andOp.getType());
+      assert(shapedType.hasStaticShape());
+      for (size_t i = 0; i < lhsState.masks.size(); i++) {
+        Value lhsV = lhsState.masks[i];
+        Value rhsV = rhsState.masks[i];
+        if (!lhsV && !rhsV) {
+          masks.push_back(nullptr);
+        } else {
+          uint32_t size = shapedType.getShape()[i];
+          auto structuredMaskToUnstructuredMask = [](MaskState state,
+                                                     unsigned dim,
+                                                     uint32_t size,
+                                                     OpBuilder &builder,
+                                                     Location loc) {
+            OpFoldResult ofr = state.isMask() ? state.dims[dim] : state.scalar;
+            if (auto intV = getIntAttr(ofr)) {
+              if (intV == size) {
+                // Full mask.
+                return Value();
+              }
+            }
+            auto targetTensorType =
+                RankedTensorType::get({size}, builder.getI32Type());
+            Value range =
+                builder
+                    .create<triton::MakeRangeOp>(loc, targetTensorType, 0, size)
+                    .getResult();
+            Value v = ofrToIndexValue(ofr, loc, builder);
+            v = builder
+                    .create<arith::IndexCastUIOp>(loc, builder.getI32Type(), v)
+                    .getResult();
+            v = builder.create<triton::SplatOp>(loc, targetTensorType, v)
+                    .getResult();
+            return builder
+                .create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult, range, v)
+                .getResult();
+          };
+          if (!lhsV) {
+            lhsV = structuredMaskToUnstructuredMask(lhsState, i, size, builder,
+                                                    loc);
+          } else if (!rhsV) {
+            rhsV = structuredMaskToUnstructuredMask(rhsState, i, size, builder,
+                                                    loc);
+          }
+          if (!lhsV) {
+            masks.push_back(rhsV);
+            continue;
+          } else if (!rhsV) {
+            masks.push_back(lhsV);
+            continue;
+          }
+          // And the mask.
+          masks.push_back(builder.create<arith::AndIOp>(loc, lhsV, rhsV));
+        }
+      }
+      // Only support one unstructured mask.
+      if (getUnstructuredMasks().size() > 1) {
+        return failure();
+      }
+    }
+  }
 
   if (!lhsState.isMask() || !rhsState.isMask()) {
     return this->minStateScalar(lhsState, rhsState, loc, builder);
@@ -365,7 +452,48 @@ LogicalResult MaskState::parseExtSI(arith::ExtSIOp op, const Location loc,
 LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc,
                                   OpBuilder &builder) {
   assert(this->isEmpty());
-
+  int cmpOpDim = -1;
+  if (auto shapedType = dyn_cast<ShapedType>(cmpOp.getType())) {
+    for (unsigned r = 0; r < shapedType.getRank(); r++) {
+      if (shapedType.getShape()[r] != 1) {
+        if (cmpOpDim != -1) {
+          // This will happen when the cmp has more than one dimension with size
+          // larger than 1.
+          // Like a < b while both a and b are tensors with shape 2x2.
+          cmpOpDim = -1;
+          break;
+        }
+        cmpOpDim = r;
+      }
+    }
+    masks.clear();
+    for (unsigned r = 0; r < shapedType.getRank(); r++) {
+      masks.push_back(nullptr);
+    }
+    // If cmpOpDim == -1, parseCmp must fail later.
+    // Here just setup unstructured masks when cmpOpDim != -1.
+    if (cmpOpDim != -1) {
+      // Save cmpOp as unstructured mask for failure case, will recover it to
+      // nullptr later if success.
+      Value unstructuredMask = cmpOp;
+      if (shapedType.getRank() > 1) {
+        // If cmpOp is not 1D, collapse it to 1D.
+        auto flatType = RankedTensorType::get({shapedType.getShape()[cmpOpDim]},
+                                              shapedType.getElementType());
+        auto maybeReassociationMap =
+            getReassociationIndicesForReshape(shapedType, flatType);
+        SmallVector<ReassociationIndices> reassociation =
+            *maybeReassociationMap;
+        // Set masks.
+        unstructuredMask = builder.create<tensor::CollapseShapeOp>(
+            loc, flatType, cmpOp, reassociation);
+      }
+      masks[cmpOpDim] = unstructuredMask;
+    }
+  } else {
+    cmpOpDim = 0;
+    masks.push_back(cmpOp);
+  }
   if (cmpOp.getPredicate() != arith::CmpIPredicate::slt &&
       cmpOp.getPredicate() != arith::CmpIPredicate::ult &&
       cmpOp.getPredicate() != arith::CmpIPredicate::sge) {
@@ -453,7 +581,10 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc,
     else
       this->dims.push_back(lhsState.dims[i]);
   }
-
+  if (cmpOpDim != -1) {
+    // Clear masks when success.
+    masks[cmpOpDim] = nullptr;
+  }
   return success();
 }
 
@@ -623,7 +754,15 @@ LogicalResult MaskState::parseSplat(triton::SplatOp splatOp, const Location loc,
 
   for (auto s : dstShape)
     this->dims.push_back(builder.getIndexAttr(s));
-
+  bool isBool = src.getType().isInteger(1);
+  if (isBool) {
+    // If src is a 1D boolean tensor and parse success.
+    // Create masks.
+    masks.clear();
+    for (unsigned i = 0; i < dstShape.size(); i++) {
+      masks.push_back(nullptr);
+    }
+  }
   return success();
 }
 
@@ -632,18 +771,76 @@ LogicalResult MaskState::parseExpandDims(triton::ExpandDimsOp expandDimsOp,
                                          OpBuilder &builder) {
   assert(this->isEmpty());
 
-  if (failed(this->parse(expandDimsOp.getSrc(), loc, builder)))
-    return failure();
-
   auto dstShape =
       cast<ShapedType>(expandDimsOp.getResult().getType()).getShape();
   auto axis = expandDimsOp.getAxis();
+  Value src = expandDimsOp.getSrc();
+  auto srcType = cast<ShapedType>(src.getType());
+  bool isBoolOp = srcType.getElementType().isInteger(1);
+  LogicalResult result = parse(src, loc, builder);
+  if (failed(result)) {
+    if (isBoolOp) {
+      if (srcType.getRank() > 1 && masks.size() != srcType.getRank()) {
+        return failure();
+      }
+    } else {
+      return failure();
+    }
+  }
+
+  if (isBoolOp) {
+    // Save mask for 1D boolean tensor
+    if (srcType.getRank() == 1) {
+      assert(dstShape.size() == 2);
+      masks.resize(dstShape.size());
+      masks[axis] = nullptr;
+      if (failed(result)) {
+        // Recover dims to allow other dim to be processed.
+        dims.clear();
+        dims.push_back(builder.getIndexAttr(srcType.getShape()[0]));
+        // Save src as unstructured mask.
+        masks[1 - axis] = src;
+      } else {
+        // save nullptr when parse success.
+        masks[1 - axis] = nullptr;
+      }
+    } else {
+      if (failed(result)) {
+        auto unstructuredMasks = getUnstructuredMasks();
+        if (unstructuredMasks.empty()) {
+          return failure();
+        }
+        if (unstructuredMasks.size() > 1) {
+          return failure();
+        }
+        auto [dim, mask] = unstructuredMasks.front();
+        // Recover dims for unstructured mask dim to allow other dim to be
+        // processed.
+        dims[dim] = builder.getIndexAttr(srcType.getShape()[dim]);
+      }
+      masks.insert(masks.begin() + axis, nullptr);
+    }
+  }
+
   assert(dstShape[axis] == 1 &&
          "expect changed dimension to be 1 in expand_dims");
   this->dims.insert(this->dims.begin() + axis, builder.getIndexAttr(1));
 
   return success();
 }
 
+// Return all non-nullptr masks along with their dimensions.
+SmallVector<std::pair<unsigned, Value>> MaskState::getUnstructuredMasks() {
+  SmallVector<std::pair<unsigned, Value>> result;
+
+  for (auto [i, m] : llvm::enumerate(masks)) {
+    if (m) {
+      result.push_back({i, m});
+    }
+  }
+
+  return result;
+}
+
 } // namespace triton
 } // namespace mlir