[mlir][llvm] Add atomic support to the StoreOp.

[gysit]

Summary

This revision adds atomic support to the StoreOp. It chooses
to print the atomic keywords together with the syncscope and
ordering arguments. The revision also implements verifiers to
ensure the constraints that apply to atomic store operations
are checked.

Depends on D144112

[Dinistro]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

918–931

Couldn't this be shared with the load verifier through a helper function? This could be done by passing in a set of parameters or by templating the function, I guess.

[gysit]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

918–931

I tried to do that. Note that the set of unsupported atomic orderings is different and that the error messages are also custom. Additionally, we can access all the arguments directly. Overall, it seemed nicer to have separate functions in this case.

[gysit]

Factor out some of the verification logic that LoadOp and StoreOp share.

[Dinistro]

LGTM module minor comment.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

852

NIT: If I'm not mistaken. ArrayRefs can also be constructed from initialiser lists, correct me if I'm wrong. Assuming this is indeed the case, it would be cleaner to replace SmallVector with ArrayRef here.

919

As above.

[gysit]

Use initializer lists.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 
 def LLVM_StoreOp : LLVM_Op<"store">, MemoryOpBase {
   let arguments = (ins LLVM_LoadableType:$value,
                    Arg<LLVM_PointerTo<LLVM_LoadableType>,"",[MemWrite]>:$addr,
                    OptionalAttr<SymbolRefArrayAttr>:$access_groups,
                    OptionalAttr<SymbolRefArrayAttr>:$alias_scopes,
                    OptionalAttr<SymbolRefArrayAttr>:$noalias_scopes,
                    OptionalAttr<SymbolRefArrayAttr>:$tbaa,
-                   OptionalAttr<I64Attr>:$alignment, UnitAttr:$volatile_,
-                   UnitAttr:$nontemporal);
+                   OptionalAttr<I64Attr>:$alignment,
+                   UnitAttr:$volatile_,
+                   UnitAttr:$nontemporal,
+                   DefaultValuedAttr<AtomicOrdering,
+                                     "AtomicOrdering::not_atomic">:$ordering,
+                   OptionalAttr<StrAttr>:$syncscope);
   string llvmInstName = "Store";
+  let description = [{
+    The `store` operation is used to write to memory. A store may be marked as
+    atomic, volatile, and/or nontemporal, and takes a number of optional
+    attributes that specify aliasing information.
+
+    An atomic store only supports a limited set of pointer, integer, and
+    floating point types, and requires an explicit alignment.
+
+    Examples:
+    ```mlir
+    // A volatile store of a float variable.
+    llvm.store volatile %val, %ptr : f32, !llvm.ptr
+
+    // A nontemporal store of a float variable.
+    llvm.store %val, %ptr {nontemporal} : f32, !llvm.ptr
+
+    // An atomic store of an integer variable.
+    llvm.store %val, %ptr atomic monotonic {alignment = 8 : i64}
+        : i64, !llvm.ptr
+    ```
+
+    See the following link for more details:
+    https://llvm.org/docs/LangRef.html#store-instruction
+  }];
   let assemblyFormat = [{
-    (`volatile` $volatile_^)? $value `,` $addr attr-dict `:`
-    custom<StoreType>(type($value), type($addr))
+    (`volatile` $volatile_^)? $value `,` $addr
+    (`atomic` (`syncscope` `(` $syncscope^ `)`)? $ordering^)?
+    attr-dict `:` custom<StoreType>(type($value), type($addr))
   }];
   string llvmBuilder = [{
     auto *inst = builder.CreateStore($value, $addr, $volatile_);
-  }] # setAlignmentCode
+  }] # setOrderingCode
+     # setSyncScopeCode
+     # setAlignmentCode
      # setNonTemporalMetadataCode
      # setAccessGroupsMetadataCode
      # setAliasScopeMetadataCode
      # setTBAAMetadataCode;
   string mlirBuilder = [{
     auto *storeInst = cast<llvm::StoreInst>(inst);
     unsigned alignment = storeInst->getAlign().value();
     $_op = $_builder.create<LLVM::StoreOp>($_location, $value, $addr,
         alignment, storeInst->isVolatile(),
-        storeInst->hasMetadata(llvm::LLVMContext::MD_nontemporal));
+        storeInst->hasMetadata(llvm::LLVMContext::MD_nontemporal),
+        convertAtomicOrderingFromLLVM(storeInst->getOrdering()),
+        getLLVMSyncScope(storeInst));
   }];
   let builders = [
     OpBuilder<(ins "Value":$value, "Value":$addr,
       CArg<"unsigned", "0">:$alignment, CArg<"bool", "false">:$isVolatile,
-      CArg<"bool", "false">:$isNonTemporal)>
+      CArg<"bool", "false">:$isNonTemporal,
+      CArg<"AtomicOrdering", "AtomicOrdering::not_atomic">:$ordering,
+      CArg<"StringRef", "StringRef()">:$syncscope)>
   ];
   let hasVerifier = 1;
 }
 
 // Casts.
 class LLVM_CastOp<string mnemonic, string instName, Type type,
                   Type resultType, list<Trait> traits = []> :
     LLVM_Op<mnemonic, !listconcat([Pure], traits)>,

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

return *bitWidth == 8 || *bitWidth == 16 || *bitWidth == 32 ||
          *bitWidth == 64;
 }
 
 LogicalResult LoadOp::verify() {
   if (getOrdering() != AtomicOrdering::not_atomic) {
     if (!isTypeCompatibleWithAtomicOp(getResult().getType(),
                                       /*isPointerTypeAllowed=*/true))
       return emitOpError("unsupported type ")
-             << getResult().getType() << " for atomic access";
+             << getResult().getType() << " for atomic load";
     if (getOrdering() == AtomicOrdering::release ||
         getOrdering() == AtomicOrdering::acq_rel)
       return emitOpError("unsupported ordering '")
              << stringifyAtomicOrdering(getOrdering()) << "'";
     if (!getAlignment())
-      return emitOpError("expected alignment for atomic access");
+      return emitOpError("expected alignment for atomic load");
   } else if (getSyncscope()) {
-    return emitOpError("expected syncscope to be null for non-atomic access");
+    return emitOpError("expected syncscope to be null for non-atomic load");
   }
   return verifyMemOpMetadata(*this);
 }
 
 void LoadOp::build(OpBuilder &builder, OperationState &state, Value addr,
                    unsigned alignment, bool isVolatile, bool isNonTemporal) {
   auto type = addr.getType().cast<LLVMPointerType>().getElementType();
   assert(type && "must provide explicit element type to the constructor "
                  "when the pointer type is opaque");
   build(builder, state, type, addr, alignment, isVolatile, isNonTemporal);
 }
 
 void LoadOp::build(OpBuilder &builder, OperationState &state, Type type,
                    Value addr, unsigned alignment, bool isVolatile,
                    bool isNonTemporal, AtomicOrdering ordering,
                    StringRef syncscope) {
   build(builder, state, type, addr,
         /*access_groups=*/nullptr,
static void printLoadType(OpAsmPrinter &printer, Operation *op, Type type,
   if (pointerType.isOpaque())
     printer << " -> " << elementType;
 }
 
 //===----------------------------------------------------------------------===//
 // StoreOp
 //===----------------------------------------------------------------------===//
 
-LogicalResult StoreOp::verify() { return verifyMemOpMetadata(*this); }
+LogicalResult StoreOp::verify() {
+  if (getOrdering() != AtomicOrdering::not_atomic) {
+    if (!isTypeCompatibleWithAtomicOp(getValue().getType(),
+                                      /*isPointerTypeAllowed=*/true))
+      return emitOpError("unsupported type ")
+             << getValue().getType() << " for atomic store";
+    if (getOrdering() == AtomicOrdering::acquire ||
+        getOrdering() == AtomicOrdering::acq_rel)
+      return emitOpError("unsupported ordering '")
+             << stringifyAtomicOrdering(getOrdering()) << "'";
+    if (!getAlignment())
+      return emitOpError("expected alignment for atomic store");
+  } else if (getSyncscope()) {
+    return emitOpError("expected syncscope to be null for non-atomic store");
+  }
+  return verifyMemOpMetadata(*this);
+}
 
 void StoreOp::build(OpBuilder &builder, OperationState &state, Value value,
                     Value addr, unsigned alignment, bool isVolatile,
-                    bool isNonTemporal) {
+                    bool isNonTemporal, AtomicOrdering ordering,
+                    StringRef syncscope) {
   build(builder, state, value, addr, /*access_groups=*/nullptr,
         /*alias_scopes=*/nullptr, /*noalias_scopes=*/nullptr, /*tbaa=*/nullptr,
         alignment ? builder.getI64IntegerAttr(alignment) : nullptr, isVolatile,
-        isNonTemporal);
+        isNonTemporal, ordering,
+        syncscope.empty() ? nullptr : builder.getStringAttr(syncscope));
 }
 
 /// Parses the StoreOp type either using the typed or opaque pointer format.
 // TODO: Drop once the typed pointer assembly format is not needed anymore.
 static ParseResult parseStoreType(OpAsmParser &parser, Type &elementType,
                                   Type &type) {
   SMLoc trailingTypeLoc;
   if (parser.getCurrentLocation(&trailingTypeLoc) ||

mlir/test/Dialect/LLVMIR/invalid.mlir

 func.func @load_non_ptr_type(%foo : f32) {
   // expected-error@+1 {{expected LLVM pointer type}}
   llvm.load %foo : f32
 }
 
 // -----
 
 func.func @load_syncscope(%ptr : !llvm.ptr) {
-  // expected-error@below {{expected syncscope to be null for non-atomic access}}
+  // expected-error@below {{expected syncscope to be null for non-atomic load}}
   %1 = "llvm.load"(%ptr) {syncscope = "singlethread"} : (!llvm.ptr) -> (f32)
 }
 
 // -----
 
 func.func @load_unsupported_ordering(%ptr : !llvm.ptr) {
   // expected-error@below {{unsupported ordering 'release'}}
   %1 = llvm.load %ptr atomic release {alignment = 4 : i64} : !llvm.ptr -> f32
 }
 
 // -----
 
 func.func @load_unsupported_type(%ptr : !llvm.ptr) {
-  // expected-error@below {{unsupported type 'f80' for atomic access}}
+  // expected-error@below {{unsupported type 'f80' for atomic load}}
   %1 = llvm.load %ptr atomic monotonic {alignment = 16 : i64} : !llvm.ptr -> f80
 }
 
 // -----
 
 func.func @load_unsupported_type(%ptr : !llvm.ptr) {
-  // expected-error@below {{unsupported type 'i1' for atomic access}}
+  // expected-error@below {{unsupported type 'i1' for atomic load}}
   %1 = llvm.load %ptr atomic monotonic {alignment = 16 : i64} : !llvm.ptr -> i1
 }
 
 // -----
 
 func.func @load_unaligned_atomic(%ptr : !llvm.ptr) {
-  // expected-error@below {{expected alignment for atomic access}}
+  // expected-error@below {{expected alignment for atomic load}}
   %1 = llvm.load %ptr atomic monotonic : !llvm.ptr -> f32
 }
 
 // -----
 
 func.func @store_non_llvm_type(%foo : memref<f32>, %bar : f32) {
   // expected-error@+1 {{expected LLVM pointer type}}
   llvm.store %bar, %foo : memref<f32>

 
 func.func @store_malformed_elem_type(%foo: !llvm.ptr, %bar: f32) {
   // expected-error@+1 {{expected non-function type}}
   llvm.store %bar, %foo : !llvm.ptr, "f32"
 }
 
 // -----
 
+func.func @store_syncscope(%val : f32, %ptr : !llvm.ptr) {
+  // expected-error@below {{expected syncscope to be null for non-atomic store}}
+  "llvm.store"(%val, %ptr) {syncscope = "singlethread"} : (f32, !llvm.ptr) -> ()
+}
+
+// -----
+
+func.func @store_unsupported_ordering(%val : f32, %ptr : !llvm.ptr) {
+  // expected-error@below {{unsupported ordering 'acquire'}}
+  llvm.store %val, %ptr atomic acquire {alignment = 4 : i64} : f32, !llvm.ptr
+}
+
+// -----
+
+func.func @store_unsupported_type(%val : f80, %ptr : !llvm.ptr) {
+  // expected-error@below {{unsupported type 'f80' for atomic store}}
+  llvm.store %val, %ptr atomic monotonic {alignment = 16 : i64} : f80, !llvm.ptr
+}
+
+// -----
+
+func.func @store_unsupported_type(%val : i1, %ptr : !llvm.ptr) {
+  // expected-error@below {{unsupported type 'i1' for atomic store}}
+  llvm.store %val, %ptr atomic monotonic {alignment = 16 : i64} : i1, !llvm.ptr
+}
+
+// -----
+
+func.func @store_unaligned_atomic(%val : f32, %ptr : !llvm.ptr) {
+  // expected-error@below {{expected alignment for atomic store}}
+  llvm.store %val, %ptr atomic monotonic : f32, !llvm.ptr
+}
+
+// -----
+
 func.func @invalid_call() {
   // expected-error@+1 {{'llvm.call' op must have either a `callee` attribute or at least an operand}}
   "llvm.call"() : () -> ()
   llvm.return
 }
 
 // -----
 

mlir/test/Dialect/LLVMIR/roundtrip.mlir

 func.func @atomic_load(%ptr : !llvm.ptr) {
   // CHECK: llvm.load %{{.*}} atomic monotonic {alignment = 4 : i64} : !llvm.ptr -> f32
   %0 = llvm.load %ptr atomic monotonic {alignment = 4 : i64} : !llvm.ptr -> f32
   // CHECK: llvm.load volatile %{{.*}} atomic syncscope("singlethread") monotonic {alignment = 16 : i64} : !llvm.ptr -> f32
   %1 = llvm.load volatile %ptr atomic syncscope("singlethread") monotonic {alignment = 16 : i64} : !llvm.ptr -> f32
   llvm.return
 }
 
+// CHECK-LABEL: @atomic_store
+func.func @atomic_store(%val : f32, %ptr : !llvm.ptr) {
+  // CHECK: llvm.store %{{.*}}, %{{.*}} atomic monotonic {alignment = 4 : i64} : f32, !llvm.ptr
+  llvm.store %val, %ptr atomic monotonic {alignment = 4 : i64} : f32, !llvm.ptr
+  // CHECK: llvm.store volatile %{{.*}}, %{{.*}} atomic syncscope("singlethread") monotonic {alignment = 16 : i64} : f32, !llvm.ptr
+  llvm.store volatile %val, %ptr atomic syncscope("singlethread") monotonic {alignment = 16 : i64} : f32, !llvm.ptr
+  llvm.return
+}
+
 // CHECK-LABEL: @atomicrmw
 func.func @atomicrmw(%ptr : !llvm.ptr, %val : f32) {
   // CHECK: llvm.atomicrmw fadd %{{.*}}, %{{.*}} monotonic : !llvm.ptr, f32
   %0 = llvm.atomicrmw fadd %ptr, %val monotonic : !llvm.ptr, f32
   // CHECK: llvm.atomicrmw volatile fsub %{{.*}}, %{{.*}} syncscope("singlethread") monotonic {alignment = 16 : i64} : !llvm.ptr, f32
   %1 = llvm.atomicrmw volatile fsub %ptr, %val syncscope("singlethread") monotonic {alignment = 16 : i64} : !llvm.ptr, f32
   llvm.return
 }

mlir/test/Target/LLVMIR/Import/instructions.ll

define void @load_store(ptr %ptr) {
   store volatile double %2, ptr %ptr, align 16, !nontemporal !0
   ret void
 }
 
 !0 = !{i32 1}
 
 ; // -----
 
-; CHECK-LABEL: @atomic_load
+; CHECK-LABEL: @atomic_load_store
 ; CHECK-SAME:  %[[PTR:[a-zA-Z0-9]+]]
-define void @atomic_load(ptr %ptr) {
+define void @atomic_load_store(ptr %ptr) {
   ; CHECK:  %[[V1:[0-9]+]] = llvm.load %[[PTR]] atomic acquire {alignment = 8 : i64} : !llvm.ptr -> f64
   ; CHECK:  %[[V2:[0-9]+]] = llvm.load volatile %[[PTR]] atomic syncscope("singlethreaded") acquire {alignment = 16 : i64} : !llvm.ptr -> f64
   %1 = load atomic double, ptr %ptr acquire, align 8
   %2 = load atomic volatile double, ptr %ptr syncscope("singlethreaded") acquire, align 16
+
+  ; CHECK:  llvm.store %[[V1]], %[[PTR]] atomic release {alignment = 8 : i64} : f64, !llvm.ptr
+  ; CHECK:  llvm.store volatile %[[V2]], %[[PTR]] atomic syncscope("singlethreaded") release {alignment = 16 : i64} : f64, !llvm.ptr
+  store atomic double %1, ptr %ptr release, align 8
+  store atomic volatile double %2, ptr %ptr syncscope("singlethreaded") release, align 16
   ret void
 }
 
 ; // -----
 
 ; CHECK-LABEL: @atomic_rmw
 ; CHECK-SAME:  %[[PTR1:[a-zA-Z0-9]+]]
 ; CHECK-SAME:  %[[VAL1:[a-zA-Z0-9]+]]

mlir/test/Target/LLVMIR/llvmir.mlir

llvm.func @nontemporal_store_and_load() {
   %1 = llvm.load %0 {nontemporal} : !llvm.ptr<i32>
   llvm.return
 }
 
 // CHECK: ![[NODE]] = !{i32 1}
 
 // -----
 
-llvm.func @atomic_load(%ptr : !llvm.ptr) {
+llvm.func @atomic_store_and_load(%ptr : !llvm.ptr) {
   // CHECK: load atomic
-  // CHECK-SAME:  monotonic, align 4
-  %1 = llvm.load %ptr atomic monotonic {alignment = 4 : i64} : !llvm.ptr -> f32
+  // CHECK-SAME:  acquire, align 4
+  %1 = llvm.load %ptr atomic acquire {alignment = 4 : i64} : !llvm.ptr -> f32
   // CHECK: load atomic
-  // CHECK-SAME:  syncscope("singlethread") monotonic, align 4
-  %2 = llvm.load %ptr atomic syncscope("singlethread") monotonic {alignment = 4 : i64} : !llvm.ptr -> f32
+  // CHECK-SAME:  syncscope("singlethread") acquire, align 4
+  %2 = llvm.load %ptr atomic syncscope("singlethread") acquire {alignment = 4 : i64} : !llvm.ptr -> f32
+
+  // CHECK: store atomic
+  // CHECK-SAME:  release, align 4
+  llvm.store %1, %ptr atomic release {alignment = 4 : i64} : f32, !llvm.ptr
+  // CHECK: store atomic
+  // CHECK-SAME:  syncscope("singlethread") release, align 4
+  llvm.store %2, %ptr atomic syncscope("singlethread") release {alignment = 4 : i64} : f32, !llvm.ptr
   llvm.return
 }
 
 // -----
 
 // Check that the translation does not crash in absence of a data layout.
 module {
   // CHECK: declare void @module_default_layout