[mlir] Add readonly and nocapture arg attributes to LLVM dialect

[ftynse]

Please do the verification in the dialect rather than in the translation.

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

972–987

I see there is code around doing the same, but this should really be checked in the dialect's verifier for argument attributes rather than here. Some of the code around largely predates that kind of verifier.

[ezhulenev]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

972–987

Should it be legal to attach these attributes to non-llvm and non-pointer types if they'll become pointers once everything will be converted to LLVM? E.g. to memref (assuming it's going to use bare ptr conversion)? Or dialect verification should be identical to this translation?

[vzakhari]

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

972–987

To the same note, I tried using isCompatibleType in https://reviews.llvm.org/D137049#change-zc5yndDKCEkj to narrow the verification for result attributes only to the compatible types. But it is subtle, e.g. verification may allow attributes for a type that will convert to LLVM type which does not allow an attribute. So I guess we may need to run types verification during translation to LLVM IR as well as part of LLVM dialect verification. Or we may drop the former and rely on LLVM verification, which will report an issue for the LLVM instruction that we try to create in a wrong way.

[vzakhari]

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

350

This should not be needed for unit attributes.

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

let extraClassDeclaration = [{
     static StringRef getParallelAccessAttrName() { return "parallel_access"; }
     static StringRef getLoopOptionsAttrName() { return "options"; }
     static StringRef getAccessGroupsAttrName() { return "access_groups"; }
     static StringRef getStructAttrsAttrName() { return "llvm.struct_attrs"; }
     static StringRef getByValAttrName() { return "llvm.byval"; }
     static StringRef getByRefAttrName() { return "llvm.byref"; }
     static StringRef getStructRetAttrName() { return "llvm.sret"; }
     static StringRef getInAllocaAttrName() { return "llvm.inalloca"; }
+    static StringRef getNoCaptureAttrName() { return "llvm.nocapture"; }
+    static StringRef getReadOnlyAttrName() { return "llvm.readonly"; }
     static StringRef getNoUndefAttrName() { return "llvm.noundef"; }
     static StringRef getSExtAttrName() { return "llvm.signext"; }
     static StringRef getZExtAttrName() { return "llvm.zeroext"; }
 
     /// Verifies if the attribute is a well-formed value for "llvm.struct_attrs"
     static LogicalResult verifyStructAttr(
         Operation *op, Attribute attr, Type annotatedType);
 

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

if (ArrayAttr argAttrDicts = funcOp.getAllArgAttrs()) {
           } else if (attr.getName().getValue() ==
                      LLVM::LLVMDialect::getStructRetAttrName()) {
             convertedAttrs.push_back(rewriter.getNamedAttr(
                 LLVM::LLVMDialect::getStructRetAttrName(), convert(attr)));
           } else if (attr.getName().getValue() ==
                      LLVM::LLVMDialect::getInAllocaAttrName()) {
             convertedAttrs.push_back(rewriter.getNamedAttr(
                 LLVM::LLVMDialect::getInAllocaAttrName(), convert(attr)));
+          } else if (attr.getName().getValue() ==
+                     LLVM::LLVMDialect::getNoCaptureAttrName()) {
+            convertedAttrs.push_back(
+                rewriter.getNamedAttr(LLVM::LLVMDialect::getNoCaptureAttrName(),
+                                      rewriter.getUnitAttr()));
+          } else if (attr.getName().getValue() ==
+                     LLVM::LLVMDialect::getReadOnlyAttrName()) {
+            convertedAttrs.push_back(
+                rewriter.getNamedAttr(LLVM::LLVMDialect::getReadOnlyAttrName(),
+                                      rewriter.getUnitAttr()));
           } else {
             convertedAttrs.push_back(attr);
           }
         }
         auto mapping = result.getInputMapping(i);
         assert(mapping && "unexpected deletion of function argument");
         for (size_t j = 0; j < mapping->size; ++j)
           newArgAttrs[mapping->inputNo + j] =

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

if (auto *type = func->getParamStructRetType(it.index())) {
           TypeAttr::get(mlirType)));
     }
     if (auto *type = func->getParamInAllocaType(it.index())) {
       Type mlirType = convertType(type);
       argAttrs.push_back(NamedAttribute(
           builder.getStringAttr(LLVMDialect::getInAllocaAttrName()),
           TypeAttr::get(mlirType)));
     }
+    if (func->getArg(it.index())->hasNoCaptureAttr()) {
+      argAttrs.push_back(NamedAttribute(
+          builder.getStringAttr(LLVMDialect::getNoCaptureAttrName()),
+          UnitAttr::get(context)));
+    }
+    if (func->getArg(it.index())->onlyReadsMemory()) {
+      argAttrs.push_back(NamedAttribute(
+          builder.getStringAttr(LLVMDialect::getReadOnlyAttrName()),
+          UnitAttr::get(context)));
+    }
 
     funcOp.setArgAttrs(it.index(), argAttrs);
   }
 
   if (FlatSymbolRefAttr personality = getPersonalityAsAttr(func))
     funcOp.setPersonalityAttr(personality);
   else if (func->hasPersonalityFn())
     emitWarning(UnknownLoc::get(context),

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

for (auto [mlirArgTy, llvmArg] :
         if (!argTy.isOpaque() && argTy.getElementType() != attr.getValue())
           return function.emitError(
               "llvm.inalloca attribute attached to LLVM pointer "
               "argument of a different type");
         llvmArg.addAttrs(llvm::AttrBuilder(llvmArg.getContext())
                              .addInAllocaAttr(convertType(attr.getValue())));
       }
 
+      if (auto attr = function.getArgAttrOfType<UnitAttr>(
+              argIdx, LLVMDialect::getNoCaptureAttrName())) {
+        auto argTy = mlirArgTy.dyn_cast<LLVM::LLVMPointerType>();
+        if (!argTy)
+          return function.emitError(
+              "llvm.nocapture attribute attached to LLVM non-pointer argument");
+        llvmArg.addAttr(llvm::Attribute::NoCapture);
+      }
+
+      if (auto attr = function.getArgAttrOfType<UnitAttr>(
+              argIdx, LLVMDialect::getReadOnlyAttrName())) {
+        auto argTy = mlirArgTy.dyn_cast<LLVM::LLVMPointerType>();
+        if (!argTy)
+          return function.emitError(
+              "llvm.readonly attribute attached to LLVM non-pointer argument");
+        llvmArg.addAttr(llvm::Attribute::ReadOnly);
+      }
+
       if (auto attr =
               function.getArgAttrOfType<UnitAttr>(argIdx, "llvm.nest")) {
         if (!mlirArgTy.isa<LLVM::LLVMPointerType>())
           return function.emitError(
               "llvm.nest attribute attached to LLVM non-pointer argument");
         llvmArg.addAttrs(llvm::AttrBuilder(llvmArg.getContext())
                              .addAttribute(llvm::Attribute::Nest));
       }

mlir/test/Dialect/LLVMIR/func.mlir

module {
   // CHECK: llvm.func @nestattr_decl(!llvm.ptr<i32> {llvm.nest})
   llvm.func @nestattr_decl(!llvm.ptr<i32> {llvm.nest})
   // CHECK: llvm.func @noundefattr_decl(i32 {llvm.noundef})
   llvm.func @noundefattr_decl(i32 {llvm.noundef})
   // CHECK: llvm.func @llvm_align_decl(!llvm.ptr<f32> {llvm.align = 4 : i64})
   llvm.func @llvm_align_decl(!llvm.ptr<f32> {llvm.align = 4})
   // CHECK: llvm.func @inallocaattr_decl(!llvm.ptr<i32> {llvm.inalloca = i32})
   llvm.func @inallocaattr_decl(!llvm.ptr<i32> {llvm.inalloca = i32})
+  // CHECK: llvm.func @nocaptureattr_decl(!llvm.ptr<i32> {llvm.nocapture})
+  llvm.func @nocaptureattr_decl(!llvm.ptr<i32> {llvm.nocapture})
+  // CHECK: llvm.func @readonlyattr_decl(!llvm.ptr<i32> {llvm.readonly})
+  llvm.func @readonlyattr_decl(!llvm.ptr<i32> {llvm.readonly})
 
 
   // CHECK: llvm.func @variadic(...)
   llvm.func @variadic(...)
 
   // CHECK: llvm.func @variadic_args(i32, i32, ...)
   llvm.func @variadic_args(i32, i32, ...)
 

mlir/test/Target/LLVMIR/Import/func-attrs.ll

 ; RUN: mlir-translate -import-llvm %s | FileCheck %s
 
-; CHECK: llvm.func @foo(%arg0: !llvm.ptr {llvm.byval = i64}, %arg1: !llvm.ptr {llvm.byref = i64}, %arg2: !llvm.ptr {llvm.sret = i64}, %arg3: !llvm.ptr {llvm.inalloca = i64})
-define void @foo(ptr byval(i64) %arg0, ptr byref(i64) %arg1, ptr sret(i64) %arg2, ptr inalloca(i64) %arg3) {
+; CHECK: llvm.func @foo(%arg0: !llvm.ptr {llvm.byval = i64}, %arg1: !llvm.ptr {llvm.byref = i64}, %arg2: !llvm.ptr {llvm.sret = i64}, %arg3: !llvm.ptr {llvm.inalloca = i64}, %arg4: !llvm.ptr {llvm.nocapture}, %arg5: !llvm.ptr {llvm.readonly})
+define void @foo(ptr byval(i64) %arg0, ptr byref(i64) %arg1, ptr sret(i64) %arg2, ptr inalloca(i64) %arg3, ptr nocapture %arg4, ptr readonly %arg5) {
   ret void
 }

mlir/test/Target/LLVMIR/llvmir-invalid.mlir

 // expected-error @below{{llvm.byref attribute attached to LLVM pointer argument of a different type}}
 llvm.func @invalid_sret(%arg0 : !llvm.ptr<f32> {llvm.byref = i32}) -> !llvm.ptr<f32> {
   llvm.return %arg0 : !llvm.ptr<f32>
 }
 
 // -----
 
 // expected-error @below{{llvm.inalloca attribute attached to LLVM non-pointer argument}}
-llvm.func @invalid_byval(%arg0 : f32 {llvm.inalloca = f32}) -> f32 {
+llvm.func @invalid_inalloca(%arg0 : f32 {llvm.inalloca = f32}) -> f32 {
   llvm.return %arg0 : f32
 }
 
 // -----
 
 // expected-error @below{{llvm.inalloca attribute attached to LLVM pointer argument of a different type}}
 llvm.func @invalid_sret(%arg0 : !llvm.ptr<f32> {llvm.inalloca = i32}) -> !llvm.ptr<f32> {
   llvm.return %arg0 : !llvm.ptr<f32>
 }
 
 // -----
 
+// expected-error @below{{llvm.nocapture attribute attached to LLVM non-pointer argument}}
+llvm.func @invalid_nocapture(%arg0 : f32 {llvm.nocapture}) -> f32 {
+  llvm.return %arg0 : f32
+}
+
+// -----
+
+// expected-error @below{{llvm.readonly attribute attached to LLVM non-pointer argument}}
+llvm.func @invalid_readonly(%arg0 : f32 {llvm.readonly}) -> f32 {
+  llvm.return %arg0 : f32
+}
+
+// -----
+
 // expected-error @below{{llvm.align attribute attached to LLVM non-pointer argument}}
 llvm.func @invalid_align(%arg0 : f32 {llvm.align = 4}) -> f32 {
   llvm.return %arg0 : f32
 }
 
 // -----
 
 // expected-error @below{{llvm.signext attribute attached to LLVM non-integer argument}}

mlir/test/Target/LLVMIR/llvmir.mlir

 // CHECK-LABEL: define void @inallocaattr(ptr inalloca(i32) %
 llvm.func @inallocaattr(%arg0: !llvm.ptr<i32> {llvm.inalloca = i32}) {
   llvm.return
 }
 
 // CHECK-LABEL: declare void @inallocaattr_decl(ptr inalloca(i32))
 llvm.func @inallocaattr_decl(!llvm.ptr<i32> {llvm.inalloca = i32})
 
+// CHECK-LABEL: define void @nocaptureattr(ptr nocapture %
+llvm.func @nocaptureattr(%arg0: !llvm.ptr<i32> {llvm.nocapture}) {
+  llvm.return
+}
+
+// CHECK-LABEL: declare void @nocaptureattr_decl(ptr nocapture)
+llvm.func @nocaptureattr_decl(%arg0: !llvm.ptr<i32> {llvm.nocapture})
+
+// CHECK-LABEL: define void @readonlyattr(ptr readonly %
+llvm.func @readonlyattr(%arg0: !llvm.ptr<i32> {llvm.readonly}) {
+  llvm.return
+}
+
+// CHECK-LABEL: declare void @readonlyattr_decl(ptr readonly)
+llvm.func @readonlyattr_decl(%arg0: !llvm.ptr<i32> {llvm.readonly})
+
 // CHECK-LABEL: define void @signextattr(i1 signext %
 llvm.func @signextattr(%arg0: i1 {llvm.signext}) {
   llvm.return
 }
 
 // CHECK-LABEL: declare void @signextattr_decl(i1 signext)
 llvm.func @signextattr_decl(i1 {llvm.signext})