tmpusers/fabianmcg/ptr-atomics

5 files changed, 181 insertions, 8 deletions

diff --git a/mlir/include/mlir/Dialect/Ptr/IR/PtrEnums.td b/mlir/include/mlir/Dialect/Ptr/IR/PtrEnums.td
index c169f48..8cd343c 100644
--- a/mlir/include/mlir/Dialect/Ptr/IR/PtrEnums.td
+++ b/mlir/include/mlir/Dialect/Ptr/IR/PtrEnums.td
@@ -32,18 +32,25 @@ def AtomicBinOpFMax : I32EnumCase<"fmax", 13, "fmax">;
 def AtomicBinOpFMin : I32EnumCase<"fmin", 14, "fmin">;
 def AtomicBinOpUIncWrap : I32EnumCase<"uinc_wrap", 15, "uinc_wrap">;
 def AtomicBinOpUDecWrap : I32EnumCase<"udec_wrap", 16, "udec_wrap">;
+def AtomicBinOpUSubCond : I32EnumCase<"usub_cond", 17, "usub_cond">;
+def AtomicBinOpUSubSat : I32EnumCase<"usub_sat", 18, "usub_sat">;
+def AtomicBinOpFMaximum : I32EnumCase<"fmaximum", 19, "fmaximum">;
+def AtomicBinOpFMinimum : I32EnumCase<"fminimum", 20, "fminimum">;
 
 def AtomicBinOp : I32Enum<
     "AtomicBinOp",
-    "ptr.atomicrmw binary operations",
+    "ptr.atomic_rmw binary operations",
     [AtomicBinOpXchg, AtomicBinOpAdd, AtomicBinOpSub, AtomicBinOpAnd,
      AtomicBinOpNand, AtomicBinOpOr, AtomicBinOpXor, AtomicBinOpMax,
      AtomicBinOpMin, AtomicBinOpUMax, AtomicBinOpUMin, AtomicBinOpFAdd,
      AtomicBinOpFSub, AtomicBinOpFMax, AtomicBinOpFMin, AtomicBinOpUIncWrap,
-     AtomicBinOpUDecWrap]> {
+     AtomicBinOpUDecWrap, AtomicBinOpUSubCond, AtomicBinOpUSubSat,
+     AtomicBinOpFMaximum, AtomicBinOpFMinimum]> {
   let cppNamespace = "::mlir::ptr";
 }
 
+def AtomicBinOpProp : EnumProp<AtomicBinOp>;
+
 //===----------------------------------------------------------------------===//
 // Atomic ordering enum attribute.
 //===----------------------------------------------------------------------===//
@@ -58,7 +65,7 @@ def AtomicOrderingSeqCst    : I32EnumCase<"seq_cst", 6, "seq_cst">;
 
 def AtomicOrdering : I32Enum<
     "AtomicOrdering",
-    "Atomic ordering for LLVM's memory model",
+    "Atomic ordering for ptr's memory model",
     [AtomicOrderingNotAtomic, AtomicOrderingUnordered, AtomicOrderingMonotonic,
      AtomicOrderingAcquire, AtomicOrderingRelease, AtomicOrderingAcqRel,
      AtomicOrderingSeqCst
diff --git a/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td b/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td
index 468a300..d939b66 100644
--- a/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td
+++ b/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td
@@ -58,6 +58,106 @@ def Ptr_Ptr1DType :
   Ptr_ShapedValueType<[Ptr_PtrType], [HasAnyRankOfPred<[1]>]>;
 
 //===----------------------------------------------------------------------===//
+// AtomicRMWOp
+//===----------------------------------------------------------------------===//
+
+def Ptr_AtomicRMWOp : Pointer_Op<"atomic_rmw", [
+    AllTypesMatch<["value", "result"]>
+  ]> {
+  let summary = "Atomic read-modify-write operation";
+  let description = [{
+    The `atomic_rmw` operation performs an atomic read-modify-write operation
+    on a memory location specified by a pointer. The operation is defined by
+    the provided binary operator.
+
+    Example:
+
+    ```mlir
+    %result = ptr.atomic_rmw add %ptr, %value monotonic :
+      !ptr.ptr<#ptr.generic_space>, i32
+    ```
+
+    See the following link for more details on the meaning of `alignment`,
+    `volatile_`, `ordering`, and `syncscope`:
+    https://llvm.org/docs/LangRef.html#atomicrmw-instruction
+  }];
+  let arguments = (ins AtomicBinOpProp:$bin_op,
+                       Ptr_PtrType:$ptr,
+                       AnyType:$value,
+                       AtomicOrderingProp:$ordering,
+                       OptionalAttr<StrAttr>:$syncscope,
+                       AlignmentProp:$alignment,
+                       UnitProp:$volatile_);
+  let results = (outs AnyType:$result);
+  let assemblyFormat = [{
+    (`volatile` $volatile_^)? $bin_op $ptr `,` $value
+    (`syncscope` `(` $syncscope^ `)`)? $ordering
+    (`alignment` `=` $alignment^)?
+    attr-dict `:` qualified(type($ptr)) `->` type($value)
+  }];
+  let builders = [
+    OpBuilder<(ins "AtomicBinOp":$binOp, "Value":$ptr, "Value":$value,
+      "AtomicOrdering":$ordering, CArg<"StringRef", "StringRef()">:$syncscope,
+      CArg<"unsigned", "0">:$alignment, CArg<"bool", "false">:$isVolatile)>
+  ];
+  let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// AtomicCmpXchgOp
+//===----------------------------------------------------------------------===//
+
+def Ptr_AtomicCmpXchgOp : Pointer_Op<"cmpxchg", [
+    AllTypesMatch<["compare", "result", "value"]>,
+  ]> {
+  let summary = "Atomic compare-and-exchange operation";
+  let description = [{
+    The `cmpxchg` operation performs an atomic compare-and-exchange operation
+    on a memory location specified by a pointer. The op atomically compares the
+    value at the memory location with a given compare value, and if they are
+    equal, it tries to store a new value into the memory.
+    The operation returns the original value at the memory location and a
+    boolean status indicating whether the exchange was successful.
+
+    Example:
+
+    ```mlir
+    %result, %status = ptr.cmpxchg %ptr, %compare, %value monotonic monotonic :
+      !ptr.ptr<#ptr.generic_space>, i32
+    ```
+
+    See the following link for more details on the meaning of `alignment`,
+    `volatile_`, `weak`, `success_ordering`, `failure_ordering`, and
+    `syncscope`:
+    https://llvm.org/docs/LangRef.html#atomicrmw-instruction
+  }];
+  let arguments = (ins Ptr_PtrType:$ptr,
+                       AnyType:$compare,
+                       AnyType:$value,
+                       AtomicOrderingProp:$success_ordering,
+                       AtomicOrderingProp:$failure_ordering,
+                       OptionalAttr<StrAttr>:$syncscope,
+                       AlignmentProp:$alignment,
+                       UnitProp:$weak,
+                       UnitProp:$volatile_);
+  let results = (outs AnyType:$result, I1:$status);
+  let assemblyFormat = [{
+    (`weak` $weak^)? (`volatile` $volatile_^)? $ptr `,` $compare `,` $value
+    (`syncscope` `(` $syncscope^ `)`)? $success_ordering $failure_ordering
+    (`alignment` `=` $alignment^)? attr-dict `:` type($ptr) `,` type($value)
+  }];
+  let builders = [
+    OpBuilder<(ins "Value":$ptr, "Value":$compare, "Value":$value,
+      "AtomicOrdering":$successOrdering,
+      "AtomicOrdering":$failureOrdering,
+      CArg<"StringRef", "StringRef()">:$syncscope,
+      CArg<"unsigned", "0">:$alignment, CArg<"bool", "false">:$isWeak,
+      CArg<"bool", "false">:$isVolatile)>
+  ];
+  let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/include/mlir/TableGen/Operator.h b/mlir/include/mlir/TableGen/Operator.h
index 9e57037..16da9f4 100644
--- a/mlir/include/mlir/TableGen/Operator.h
+++ b/mlir/include/mlir/TableGen/Operator.h
@@ -325,12 +325,12 @@ public:
 
   /// Pair consisting kind of argument and index into operands or attributes.
   struct OperandOrAttribute {
-    enum class Kind { Operand, Attribute };
+    enum class Kind { Operand = 0x0, Attribute = 0x1, Property = 0x2 };
     OperandOrAttribute(Kind kind, int index) {
-      packed = (index << 1) | (kind == Kind::Attribute);
+      packed = (index << 2) | static_cast<int>(kind);
     }
-    int operandOrAttributeIndex() const { return (packed >> 1); }
-    Kind kind() { return (packed & 0x1) ? Kind::Attribute : Kind::Operand; }
+    int operandOrAttributeIndex() const { return (packed >> 2); }
+    Kind kind() const { return static_cast<Kind>(packed & 0x3); }
 
   private:
     int packed;
diff --git a/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp b/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
index f0209af..681d4b6 100644
--- a/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
+++ b/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
@@ -57,6 +57,69 @@ verifyAlignment(std::optional<int64_t> alignment,
 }
 
 //===----------------------------------------------------------------------===//
+// AtomicRMWOp
+//===----------------------------------------------------------------------===//
+
+void AtomicRMWOp::build(OpBuilder &builder, OperationState &state,
+                        AtomicBinOp binOp, Value ptr, Value val,
+                        AtomicOrdering ordering, StringRef syncscope,
+                        unsigned alignment, bool isVolatile) {
+  build(builder, state, val.getType(), binOp, ptr, val, ordering,
+        !syncscope.empty() ? builder.getStringAttr(syncscope) : nullptr,
+        alignment ? std::optional<int64_t>(alignment) : std::nullopt,
+        isVolatile);
+}
+
+LogicalResult AtomicRMWOp::verify() {
+  auto emitDiag = [&]() -> InFlightDiagnostic { return emitError(); };
+  MemorySpaceAttrInterface ms = getPtr().getType().getMemorySpace();
+  DataLayout dataLayout = DataLayout::closest(*this);
+  if (!ms.isValidAtomicOp(getBinOp(), getValue().getType(), getOrdering(),
+                          getAlignment(), &dataLayout, emitDiag))
+    return failure();
+  if (getOrdering() < AtomicOrdering::monotonic) {
+    return emitError() << "expected at least '"
+                       << stringifyAtomicOrdering(AtomicOrdering::monotonic)
+                       << "' ordering";
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// AtomicCmpXchgOp
+//===----------------------------------------------------------------------===//
+
+void AtomicCmpXchgOp::build(OpBuilder &builder, OperationState &state,
+                            Value ptr, Value cmp, Value val,
+                            AtomicOrdering successOrdering,
+                            AtomicOrdering failureOrdering, StringRef syncscope,
+                            unsigned alignment, bool isWeak, bool isVolatile) {
+  build(builder, state, ptr, cmp, val, successOrdering, failureOrdering,
+        !syncscope.empty() ? builder.getStringAttr(syncscope) : nullptr,
+        alignment ? std::optional<int64_t>(alignment) : std::nullopt, isWeak,
+        isVolatile);
+}
+
+LogicalResult AtomicCmpXchgOp::verify() {
+  auto emitDiag = [&]() -> InFlightDiagnostic { return emitError(); };
+  MemorySpaceAttrInterface ms = getPtr().getType().getMemorySpace();
+  DataLayout dataLayout = DataLayout::closest(*this);
+  if (!ms.isValidAtomicXchg(getResult().getType(), getSuccessOrdering(),
+                            getFailureOrdering(), getAlignment(), &dataLayout,
+                            emitDiag))
+    return failure();
+  if (failed(verifyAlignment(getAlignment(), emitDiag)))
+    return failure();
+  if (getSuccessOrdering() < AtomicOrdering::monotonic ||
+      getFailureOrdering() < AtomicOrdering::monotonic)
+    return emitError("ordering must be at least 'monotonic'");
+  if (getFailureOrdering() == AtomicOrdering::release ||
+      getFailureOrdering() == AtomicOrdering::acq_rel)
+    return emitError("failure ordering cannot be 'release' or 'acq_rel'");
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/lib/TableGen/Operator.cpp b/mlir/lib/TableGen/Operator.cpp
index da86b00..953d8bc 100644
--- a/mlir/lib/TableGen/Operator.cpp
+++ b/mlir/lib/TableGen/Operator.cpp
@@ -385,7 +385,8 @@ void Operator::populateTypeInferenceInfo(
   if (getTrait("::mlir::OpTrait::SameOperandsAndResultType")) {
     // Check for a non-variable length operand to use as the type anchor.
     auto *operandI = llvm::find_if(arguments, [](const Argument &arg) {
-      NamedTypeConstraint *operand = llvm::dyn_cast_if_present<NamedTypeConstraint *>(arg);
+      NamedTypeConstraint *operand =
+          llvm::dyn_cast_if_present<NamedTypeConstraint *>(arg);
       return operand && !operand->isVariableLength();
     });
     if (operandI == arguments.end())
@@ -672,6 +673,8 @@ void Operator::populateOpStructure() {
       arguments.emplace_back(&attributes[attrIndex++]);
     } else {
       assert(argDef->isSubClassOf(propertyClass));
+      attrOrOperandMapping.push_back(
+          {OperandOrAttribute::Kind::Property, propIndex});
       arguments.emplace_back(&properties[propIndex++]);
     }
   }