Add support for unstructured masks in MaskAnalysis

include/triton-shared/Analysis/MaskAnalysis.h

@@ -44,17 +44,33 @@ 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
+// 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>> getGenericMasks();
+
   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

@@ -316,6 +316,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";
 }
 
@@ -337,14 +342,51 @@ 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()) {
+      for (size_t i = 0; i < lhsState.masks.size(); i++) {
+        if (lhsState.masks[i] && rhsState.masks[i]) {
+          // And the mask.
+          masks.push_back(builder.create<arith::AndIOp>(loc, lhsState.masks[i],
+                                                        rhsState.masks[i]));
+        } else {
+          masks.push_back(lhsState.masks[i] ? lhsState.masks[i]
+                                            : rhsState.masks[i]);
+        }
+      }
+      // Only support one generic mask.
+      if (getGenericMasks().size() > 1) {
+        return failure();
+      }
+    }
+  }
 
   if (!lhsState.isMask() || !rhsState.isMask()) {
     return this->minStateScalar(lhsState, rhsState, loc, builder);
@@ -361,7 +403,45 @@ 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) {
+          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 generic masks when cmpOpDim != -1.
+    if (cmpOpDim != -1) {
+      // Save cmpOp as generic mask for failure case, will recover it to nullptr
+      // later if success.
+      Value genericMask = 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.
+        genericMask = builder.create<tensor::CollapseShapeOp>(
+            loc, flatType, cmpOp, reassociation);
+      }
+      masks[cmpOpDim] = genericMask;
+    }
+  } else {
+    cmpOpDim = 0;
+    masks.push_back(cmpOp);
+  }
   if (cmpOp.getPredicate() != arith::CmpIPredicate::slt &&
       cmpOp.getPredicate() != arith::CmpIPredicate::ult &&
       cmpOp.getPredicate() != arith::CmpIPredicate::sge) {
@@ -449,7 +529,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();
 }
 
@@ -619,7 +702,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();
 }
 
@@ -628,18 +719,75 @@ 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 generic mask.
+        masks[1 - axis] = src;
+      } else {
+        // save nullptr when parse success.
+        masks[1 - axis] = nullptr;
+      }
+    } else {
+      if (failed(result)) {
+        auto genericMasks = getGenericMasks();
+        if (genericMasks.empty()) {
+          return failure();
+        }
+        if (genericMasks.size() > 1) {
+          return failure();
+        }
+        auto [dim, mask] = genericMasks.front();
+        // Recover dims for generic 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::getGenericMasks() {
+  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