[MLIR][LLVMDialect] Added volatile and nontemporal attributes to load and store

[georgemitenkov]

Summary

This patch introduces 2 new optional attributes to llvm.load
and llvm.store ops: isVolatile and isNonTemporal. These attributes
are translated into proper LLVM as a volatile marker and a metadata node
respectively. They are also helpful with SPIR-V to LLVM dialect conversion
since they are the mappings for Volatile and NonTemporal Memory Operands.

[georgemitenkov]

There is also a bit of duplication when setting the attributes. I kept it like this for now, as I am not sure what is the best practise for handling duplication in builders in tablegen?

[ftynse]

There is also a bit of duplication when setting the attributes. I kept it like this for now, as I am not sure what is the best practise for handling duplication in builders in tablegen?

Builders are just C++ code. You can call helper functions, including other builders. There are many places where we call the auto-generated builder from the custom builder after casting the types and adding the missing default arguments.

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

285

I'd use UnitAttr instead. Optional<Bool> gives you a tri-state attribute, but volatile is actually and non-temporal are actually bi-state. The presence of the attribute is meaningful by itself, so is its absence.

286

Nit: the convention seems to lean towards using property-like names (e.g. volatile) as opposed to function-like names (e.g. isVolatile) for attributes. I understand this generates less meaningful function names, but let's be consistent.

302

This looks long enough to go to the .cpp file instead.

[georgemitenkov]

Used _volatile and nontemporal UnitAttr instead of optional bool attribute.
Builder code is moved to .cpp file, and prining is also prettier now.

[ftynse]

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

34 ↗

(On Diff #280407)

Let's use a trailing underscore, names with leading underscores are often reserved, and also mess up with alphabetically-ordered lists.

203 ↗

(On Diff #280407)

Since we have a custom printer anyway, we should drop the underscore from the pretty syntax.

[georgemitenkov]

Addressed comments: put underscore in the end and removed it when printing.

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

code setAlignmentCode = [{
     if ($alignment.hasValue()) {
       auto align = $alignment.getValue().getZExtValue();
       if (align != 0)
         inst->setAlignment(llvm::Align(align));
     }
   }];
 }
 
+// Code definition that is used for nontemporal metadata creation.
+class MemoryOpWithAlignmentAndAttributes : MemoryOpWithAlignmentBase {
+  code setNonTemporalMetadataCode = [{
+    if ($isNonTemporal.hasValue() && $isNonTemporal.getValue()) {
+      llvm::Module *module = builder.GetInsertBlock()->getModule();
+      llvm::MDNode *metadata = llvm::MDNode::get(
+          inst->getContext(), llvm::ConstantAsMetadata::get(
+              builder.getInt32(1)));
+      inst->setMetadata(module->getMDKindID("nontemporal"), metadata);
+    }
+  }];
+}
+
 // Memory-related operations.
 def LLVM_AllocaOp :
     MemoryOpWithAlignmentBase,
     LLVM_OneResultOp<"alloca">,
     Arguments<(ins LLVM_Type:$arraySize, OptionalAttr<I64Attr>:$alignment)> {
   string llvmBuilder = [{
     auto *inst = builder.CreateAlloca(
       $_resultType->getPointerElementType(), $arraySize);

 def LLVM_GEPOp : LLVM_OneResultOp<"getelementptr", [NoSideEffect]>,
                  Arguments<(ins LLVM_Type:$base, Variadic<LLVM_Type>:$indices)>,
                  LLVM_Builder<"$res = builder.CreateGEP($base, $indices);"> {
   let assemblyFormat = [{
     $base `[` $indices `]` attr-dict `:` functional-type(operands, results)
   }];
 }
 def LLVM_LoadOp :
-    MemoryOpWithAlignmentBase,
+    MemoryOpWithAlignmentAndAttributes,
     LLVM_OneResultOp<"load">,
-    Arguments<(ins LLVM_Type:$addr, OptionalAttr<I64Attr>:$alignment)> {
+    Arguments<(ins LLVM_Type:$addr,
+                   OptionalAttr<BoolAttr>:$isVolatile,
+                   OptionalAttr<BoolAttr>:$isNonTemporal,
+                   OptionalAttr<I64Attr>:$alignment)> {
   string llvmBuilder = [{
-    auto *inst = builder.CreateLoad($addr);
-  }] # setAlignmentCode # [{
+    auto *inst =
+        $isVolatile.hasValue()
+            ? builder.CreateLoad($addr, $isVolatile.getValue())
+            : builder.CreateLoad($addr);
+  }] # setAlignmentCode # setNonTemporalMetadataCode # [{
     $res = inst;
   }];
   let builders = [OpBuilder<
-    "OpBuilder &b, OperationState &result, Value addr, unsigned alignment = 0",
+    "OpBuilder &b, OperationState &result, Value addr, unsigned alignment = 0, "
+    "bool isVolatile = false, bool isNonTemporal = false",
     [{
       auto type = addr.getType().cast<LLVM::LLVMType>().getPointerElementTy();
-      build(b, result, type, addr, alignment);
+      auto volatileAttr = isVolatile ? b.getBoolAttr(isVolatile) : BoolAttr();
+      auto nonTemporalAttr =
+          isNonTemporal
+              ? b.getBoolAttr(isNonTemporal)
+              : BoolAttr();
+      auto alignAttr =
+          alignment == 0
+              ? IntegerAttr()
+              : b.getI64IntegerAttr(alignment);
+      build(b, result, type, addr, volatileAttr, nonTemporalAttr, alignAttr);
     }]>,
     OpBuilder<
     "OpBuilder &b, OperationState &result, Type t, Value addr, "
-    "unsigned alignment = 0",
+    "unsigned alignment = 0, bool isVolatile = false, "
+    "bool isNonTemporal = false",
     [{
-      if (alignment == 0)
-        return build(b, result, t, addr, IntegerAttr());
-      build(b, result, t, addr, b.getI64IntegerAttr(alignment));
+      auto volatileAttr = isVolatile ? b.getBoolAttr(isVolatile) : BoolAttr();
+      auto nonTemporalAttr =
+          isNonTemporal
+              ? b.getBoolAttr(isNonTemporal)
+              : BoolAttr();
+      auto alignAttr =
+          alignment == 0
+              ? IntegerAttr()
+              : b.getI64IntegerAttr(alignment);
+      build(b, result, t, addr, volatileAttr, nonTemporalAttr, alignAttr);
     }]>];
   let parser = [{ return parseLoadOp(parser, result); }];
   let printer = [{ printLoadOp(p, *this); }];
   let verifier = alignmentVerifierCode;
 }
 def LLVM_StoreOp :
-    MemoryOpWithAlignmentBase,
+    MemoryOpWithAlignmentAndAttributes,
     LLVM_ZeroResultOp<"store">,
     Arguments<(ins LLVM_Type:$value,
                    LLVM_Type:$addr,
+                   OptionalAttr<BoolAttr>:$isVolatile,
+                   OptionalAttr<BoolAttr>:$isNonTemporal,
                    OptionalAttr<I64Attr>:$alignment)> {
   string llvmBuilder = [{
-    auto *inst = builder.CreateStore($value, $addr);
-  }] # setAlignmentCode;
+    auto *inst =
+        $isVolatile.hasValue()
+            ? builder.CreateStore($value, $addr, $isVolatile.getValue())
+            : builder.CreateStore($value, $addr);
+  }] # setAlignmentCode # setNonTemporalMetadataCode;
   let builders = [
     OpBuilder<
     "OpBuilder &b, OperationState &result, Value value, Value addr, "
+    "bool isVolatile = false, bool isNonTemporal = false, "
     "unsigned alignment = 0",
     [{
-      if (alignment == 0)
-        return build(b, result, ArrayRef<Type>{}, value, addr, IntegerAttr());
-      build(b, result, ArrayRef<Type>{}, value, addr, 
-            b.getI64IntegerAttr(alignment));
+      auto volatileAttr = isVolatile ? b.getBoolAttr(isVolatile) : BoolAttr();
+      auto nonTemporalAttr =
+          isNonTemporal
+              ? b.getBoolAttr(isNonTemporal)
+              : BoolAttr();
+      auto alignAttr =
+          alignment == 0
+              ? IntegerAttr()
+              : b.getI64IntegerAttr(alignment);
+      build(b, result, ArrayRef<Type>{}, value, addr, volatileAttr,
+            nonTemporalAttr, alignAttr);
     }]
   >];
   let parser = [{ return parseStoreOp(parser, result); }];
   let printer = [{ printStoreOp(p, *this); }];
   let verifier = alignmentVerifierCode;
 }
 
 // Casts.

mlir/test/Target/llvmir.mlir

 llvm.func @address_taken() {
   llvm.return
 }
 
 llvm.mlir.global internal constant @taker_of_address() : !llvm<"void()*"> {
   %0 = llvm.mlir.addressof @address_taken : !llvm<"void()*">
   llvm.return %0 : !llvm<"void()*">
 }
+
+// -----
+
+llvm.func @volatile_store_and_load() {
+  %val = llvm.mlir.constant(5 : i32) : !llvm.i32
+  %size = llvm.mlir.constant(1 : i64) : !llvm.i64
+  %0 = llvm.alloca %size x !llvm.i32 : (!llvm.i64) -> (!llvm<"i32*">)
+  // CHECK: store volatile i32 5, i32* %{{.*}}
+  llvm.store %val, %0 {isVolatile = true} : !llvm<"i32*">
+  // CHECK: %{{.*}} = load volatile i32, i32* %{{.*}}
+  %1 = llvm.load %0 {isVolatile = true} : !llvm<"i32*">
+  llvm.return
+}
+
+// -----
+
+// Check that nontemporal attribute is exported as metadata node.
+llvm.func @nontemoral_store_and_load() {
+  %val = llvm.mlir.constant(5 : i32) : !llvm.i32
+  %size = llvm.mlir.constant(1 : i64) : !llvm.i64
+  %0 = llvm.alloca %size x !llvm.i32 : (!llvm.i64) -> (!llvm<"i32*">)
+  // CHECK: !nontemporal ![[NODE:[0-9]+]]
+  llvm.store %val, %0 {isNonTemporal = true} : !llvm<"i32*">
+  // CHECK: !nontemporal ![[NODE]]
+  %1 = llvm.load %0 {isNonTemporal = true} : !llvm<"i32*">
+  llvm.return
+}
+
+// CHECK: ![[NODE]] = !{i32 1}