Add structured atomic lowering

include/triton-shared/AnalysisStructured/PtrAnalysis.h

@@ -358,6 +358,11 @@ class PtrAnalysis {
 
   LogicalResult rewriteStoreOp(triton::StoreOp op, bool useUnsafeMask = false);
 
+  LogicalResult rewriteAtomicRMWOp(triton::AtomicRMWOp op,
+                                   bool useUnsafeMask = false);
+
+  LogicalResult rewriteAtomicCASOp(triton::AtomicCASOp op);
+
   LogicalResult rewriteOp(Operation *op, bool useUnsafeMask = false);
 };
 

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

@@ -11,6 +11,7 @@
 include "mlir/IR/OpBase.td"
 include "triton/Dialect/Triton/IR/TritonTypes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "triton/Dialect/Triton/IR/TritonAttrDefs.td"
 
 def Triton_Structured_Dialect : Dialect {
   let name = "tts";
@@ -326,6 +327,8 @@ def TTS_LoadOp : TTS_Op<"load", [
     }
   }];
 
+  let hasFolder = 1;
+
   // TODO
   //let hasCustomAssemblyFormat = 1;
   //let hasVerifier = 1;
@@ -357,9 +360,95 @@ def TTS_StoreOp : TTS_Op<"store", [
     }
   }];
 
+  let hasFolder = 1;
+
   // TODO
   //let hasCustomAssemblyFormat = 1;
   //let hasVerifier = 1;
 }
 
+def TTS_AtomicRMWOp : TTS_Op<"atomic_rmw", [
+  MemoryEffects<[MemWrite, MemRead]>
+]> {
+  let summary = "atomic read-modify-write";
+
+  let description = [{
+    Load the old value at $ptr, apply $atomic_rmw_op with $val, store the
+    result back to $ptr, and return the old value.
+  }];
+
+  let arguments = (ins
+    TT_AtomicRMWAttr:$atomic_rmw_op,
+    TT_PtrLike:$ptr,
+    TT_Tensor:$val,
+    Variadic<Index>:$mask_dims,
+    DenseI64ArrayAttr:$static_mask_dims,
+    TT_MemSemanticAttr:$sem,
+    TT_MemSyncScopeAttr:$scope
+  );
+
+  let results = (outs TT_Tensor:$result);
+
+  let builders = [
+    OpBuilder<(ins "triton::RMWOp":$atomic_rmw_op, "Value":$ptr, "Value":$value, "ArrayRef<OpFoldResult>":$dims, "triton::MemSemantic":$sem, "triton::MemSyncScope":$scope)>,
+  ];
+
+  let extraClassDeclaration = [{
+    /// Return a vector of all the static or dynamic fields
+    SmallVector<OpFoldResult> getMixedMaskDims() {
+      Builder b(getContext());
+      return ::mlir::getMixedValues(getStaticMaskDims(), getMaskDims(), b);
+    }
+
+    bool hasMask() {
+      return !getMixedMaskDims().empty();
+    }
+  }];
+
+  let hasFolder = 1;
+
+  // Explicitly list $atomic_rmw_op, $sem, and $scope rather than relying on
+  // attr-dict so they're printed as strings rather than opaque integers.
+  let assemblyFormat = [{
+    $atomic_rmw_op `,` $sem `,` $scope `,` $ptr `,` $val (`,` $mask_dims^)? attr-dict `:`
+    functional-type(operands, $result)
+  }];
+
+  let hasVerifier = 1;
+}
+
+def TTS_AtomicCASOp : TTS_Op<"atomic_cas", [
+  MemoryEffects<[MemWrite, MemRead]>
+]> {
+  let summary = "atomic compare-and-swap";
+
+  let description = [{
+    Compare $cmp with the old value at $ptr. Store $val when they match,
+    otherwise keep the old value, and return the old value.
+  }];
+
+  let arguments = (ins
+    TT_PtrLike:$ptr,
+    TT_Tensor:$cmp,
+    TT_Tensor:$val,
+    TT_MemSemanticAttr:$sem,
+    TT_MemSyncScopeAttr:$scope
+  );
+
+  let results = (outs TT_Tensor:$result);
+
+  let builders = [
+    OpBuilder<(ins "triton::MemSemantic":$sem, "triton::MemSyncScope":$scope, "Value":$ptr, "Value":$cmp, "Value":$val)>,
+  ];
+
+  // Explicitly list $sem and $scope rather than relying on attr-dict so
+  // they're printed as strings rather than opaque integers.
+  let assemblyFormat = [{
+    $sem `,` $scope `,` $ptr `,` $cmp `,` $val attr-dict `:`
+    functional-type(operands, $result)
+  }];
+
+  let hasVerifier = 1;
+}
+
 #endif // TRITON_STRUCTURED_DIALECT

lib/AnalysisStructured/PtrAnalysis.cpp

@@ -1748,6 +1748,97 @@ LogicalResult PtrAnalysis::rewriteStoreOp(triton::StoreOp op,
   return success();
 }
 
+LogicalResult PtrAnalysis::rewriteAtomicRMWOp(triton::AtomicRMWOp op,
+                                              bool useUnsafeMask) {
+  auto ptr = ptrMap.lookupOrNull(op.getPtr());
+  auto val = op.getVal();
+  auto mask = op.getMask();
+  auto loc = op.getLoc();
+
+  auto atomicRmwOp = op.getAtomicRmwOp();
+  auto sem = op.getSem();
+  auto scope = op.getScope();
+
+  if (!ptr) {
+    LLVM_DEBUG(op->emitRemark(
+        "PtrAnalysis: pointer is not replace with tts.make_tptr so "
+        "atomicRMWOp cannot be rewritten"));
+    return failure();
+  }
+
+  auto ptrType = dyn_cast<triton::PointerType>(ptr.getType());
+  if (ptrType && !isa<ShapedType>(ptrType.getPointeeType())) {
+    LLVM_DEBUG(op->emitRemark(
+        "PtrAnalysis: scalar atomicRMWOp will not be rewritten"));
+    return failure();
+  }
+
+  ArrayRef<OpFoldResult> dims;
+  mlir::triton::MaskState mstate(useUnsafeMask);
+
+  OpBuilder builder(op);
+
+  // Analyze the mask operand to determine at runtime the size of the data
+  // are moving.
+  if (mask) {
+    if (mstate.parse(mask, loc, builder).failed()) {
+      LLVM_DEBUG(op->emitRemark("MaskAnalysis failed"));
+      return failure();
+    }
+    dims = mstate.dims;
+  }
+
+  auto newOp = tts::AtomicRMWOp::create(builder, loc, atomicRmwOp, ptr, val,
+                                        dims, sem, scope);
+
+  LLVM_DEBUG({
+    llvm::dbgs() << "creating tts::atomic_rmw:\n";
+    newOp->dump();
+  });
+
+  op.replaceAllUsesWith(newOp.getResult());
+  op->erase();
+  return success();
+}
+
+LogicalResult PtrAnalysis::rewriteAtomicCASOp(triton::AtomicCASOp op) {
+  auto ptr = ptrMap.lookupOrNull(op.getPtr());
+  auto cmp = op.getCmp();
+  auto val = op.getVal();
+  auto loc = op.getLoc();
+
+  auto sem = op.getSem();
+  auto scope = op.getScope();
+
+  if (!ptr) {
+    LLVM_DEBUG(op->emitRemark(
+        "PtrAnalysis: pointer is not replace with tts.make_tptr so "
+        "atomicCASOp cannot be rewritten"));
+    return failure();
+  }
+
+  auto ptrType = dyn_cast<triton::PointerType>(ptr.getType());
+  if (ptrType && !isa<ShapedType>(ptrType.getPointeeType())) {
+    LLVM_DEBUG(op->emitRemark(
+        "PtrAnalysis: scalar atomicCASOp will not be rewritten"));
+    return failure();
+  }
+
+  OpBuilder builder(op);
+
+  auto newOp =
+      tts::AtomicCASOp::create(builder, loc, sem, scope, ptr, cmp, val);
+
+  LLVM_DEBUG({
+    llvm::dbgs() << "creating tts::atomic_cas:\n";
+    newOp->dump();
+  });
+
+  op.replaceAllUsesWith(newOp.getResult());
+  op->erase();
+  return success();
+}
+
 LogicalResult PtrAnalysis::rewriteOp(Operation *rootOp, bool useUnsafeMask) {
   LLVM_DEBUG({
     llvm::dbgs() << "rewriting rootOp\n";
@@ -1759,6 +1850,22 @@ LogicalResult PtrAnalysis::rewriteOp(Operation *rootOp, bool useUnsafeMask) {
       return WalkResult::advance();
     }
     return TypeSwitch<Operation *, WalkResult>(op)
+        .Case<triton::AtomicRMWOp>([&](auto atomicRmwOp) {
+          if (rewriteAtomicRMWOp(atomicRmwOp, useUnsafeMask).failed()) {
+            LLVM_DEBUG(atomicRmwOp->emitRemark(
+                "PtrAnalysis: Failed to rewrite AtomicRMWOp"));
+            return WalkResult::advance();
+          }
+          return WalkResult::skip();
+        })
+        .Case<triton::AtomicCASOp>([&](auto atomicCasOp) {
+          if (rewriteAtomicCASOp(atomicCasOp).failed()) {
+            LLVM_DEBUG(atomicCasOp->emitRemark(
+                "PtrAnalysis: Failed to rewrite AtomicCASOp"));
+            return WalkResult::advance();
+          }
+          return WalkResult::skip();
+        })
         .Case<triton::AddPtrOp>([&](auto addptr) {
           if (rewriteAddptrOp(addptr).failed()) {
             LLVM_DEBUG(