[mlir][GPU] Handle LLVM pointer attributes on memref arguments.

[stefankoncarevic]

Summary

Handle pointer attributes (noalias, nonnull, readonly, writeonly,
dereferencable, dereferencable_or_null). "noalias" attribute is
ignore for non-bare pointer.

[krzysz00]

Approving to move this along

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp

return rewriter.notifyMatchFailure(gpuFuncOp, [&](Diagnostic &diag) {
       diag << "failed to convert function signature type for: "
            << gpuFuncOp.getFunctionType();
     });
   }
 
   // Create the new function operation. Only copy those attributes that are
   // not specific to function modeling.
   SmallVector<NamedAttribute, 4> attributes;
+  ArrayAttr argAttrs;
   for (const auto &attr : gpuFuncOp->getAttrs()) {
     if (attr.getName() == SymbolTable::getSymbolAttrName() ||
         attr.getName() == gpuFuncOp.getFunctionTypeAttrName() ||
         attr.getName() ==
             gpu::GPUFuncOp::getNumWorkgroupAttributionsAttrName() ||
         attr.getName() == gpuFuncOp.getWorkgroupAttribAttrsAttrName() ||
         attr.getName() == gpuFuncOp.getPrivateAttribAttrsAttrName())
       continue;
+    if (attr.getName() == gpuFuncOp.getArgAttrsAttrName()) {
+      argAttrs = gpuFuncOp.getArgAttrsAttr();
+      continue;
+    }
     attributes.push_back(attr);
   }
   // Add a dialect specific kernel attribute in addition to GPU kernel
   // attribute. The former is necessary for further translation while the
   // latter is expected by gpu.launch_func.
   if (gpuFuncOp.isKernel())
     attributes.emplace_back(kernelAttributeName, rewriter.getUnitAttr());
   auto llvmFuncOp = rewriter.create<LLVM::LLVMFuncOp>(
for (const auto &en : llvm::enumerate(gpuFuncOp.getArgumentTypes())) {
           loc, getTypeConverter()->convertType(memrefTy));
       rewriter.replaceUsesOfBlockArgument(newArg, placeholder);
       Value desc = MemRefDescriptor::fromStaticShape(
           rewriter, loc, *getTypeConverter(), memrefTy, newArg);
       rewriter.replaceOp(placeholder, {desc});
     }
   }
 
+  // Get memref type from function arguments and set the noalias to
+  // pointer arguments.
+  for (const auto &en : llvm::enumerate(gpuFuncOp.getArgumentTypes())) {
+    auto memrefTy = en.value().dyn_cast<MemRefType>();
+    NamedAttrList argAttr = argAttrs
+                                ? argAttrs[en.index()].cast<DictionaryAttr>()
+                                : NamedAttrList();
+
+    auto copyPointerAttribute = [&](StringRef attrName) {
+      Attribute attr = argAttr.erase(attrName);
+
+      // This is a proxy for the bare pointer calling convention.
+      if (!attr)
+        return;
+      auto remapping = signatureConversion.getInputMapping(en.index());
+      if (remapping->size > 1 &&
+          attrName == LLVM::LLVMDialect::getNoAliasAttrName()) {
+        emitWarning(llvmFuncOp.getLoc(),
+                    "Cannot copy noalias with non-bare pointers.\n");
+        return;
+      }
+      for (size_t i = 0, e = remapping->size; i < e; ++i) {
+        if (llvmFuncOp.getArgument(remapping->inputNo + i)
+                .getType()
+                .isa<LLVM::LLVMPointerType>()) {
+          llvmFuncOp.setArgAttr(remapping->inputNo + i, attrName, attr);
+        }
+      }
+    };
+
+    if (argAttr.empty())
+      continue;
+
+    if (memrefTy) {
+      copyPointerAttribute(LLVM::LLVMDialect::getNoAliasAttrName());
+      copyPointerAttribute(LLVM::LLVMDialect::getReadonlyAttrName());
+      copyPointerAttribute(LLVM::LLVMDialect::getWriteOnlyAttrName());
+      copyPointerAttribute(LLVM::LLVMDialect::getNonNullAttrName());
+      copyPointerAttribute(LLVM::LLVMDialect::getDereferenceableAttrName());
+      copyPointerAttribute(
+          LLVM::LLVMDialect::getDereferenceableOrNullAttrName());
+    }
+  }
   rewriter.eraseOp(gpuFuncOp);
   return success();
 }
 
 static SmallString<16> getUniqueFormatGlobalName(gpu::GPUModuleOp moduleOp) {
   const char formatStringPrefix[] = "printfFormat_";
   // Get a unique global name.
   unsigned stringNumber = 0;

mlir/test/Conversion/GPUCommon/memref-arg-attrs.mlir

+// RUN: mlir-opt %s -split-input-file -convert-gpu-to-rocdl='use-opaque-pointers=1 use-bare-ptr-memref-call-conv=0' | FileCheck %s --check-prefixes=CHECK,ROCDL
+// RUN: mlir-opt %s -split-input-file -convert-gpu-to-nvvm='use-opaque-pointers=1 use-bare-ptr-memref-call-conv=0' | FileCheck %s --check-prefixes=CHECK,NVVM
+
+gpu.module @kernel {
+  gpu.func @test_func_readonly(%arg0 : memref<f32> {llvm.readonly} ) {
+    gpu.return
+  }
+}
+
+// CHECK-LABEL:  llvm.func @test_func_readonly
+// ROCDL-SAME:  !llvm.ptr {llvm.readonly}
+//  NVVM-SAME:  !llvm.ptr {llvm.readonly}
+
+
+// -----
+
+gpu.module @kernel {
+  gpu.func @test_func_writeonly(%arg0 : memref<f32> {llvm.writeonly} ) {
+    gpu.return
+  }
+}
+
+// CHECK-LABEL:  llvm.func @test_func_writeonly
+// ROCDL-SAME:  !llvm.ptr {llvm.writeonly}
+//  NVVM-SAME:  !llvm.ptr {llvm.writeonly}
+
+
+// -----
+
+gpu.module @kernel {
+  gpu.func @test_func_nonnull(%arg0 : memref<f32> {llvm.nonnull} ) {
+    gpu.return
+  }
+}
+
+// CHECK-LABEL:  llvm.func @test_func_nonnull
+// ROCDL-SAME:  !llvm.ptr {llvm.nonnull}
+//  NVVM-SAME:  !llvm.ptr {llvm.nonnull}
+
+
+// -----
+
+gpu.module @kernel {
+  gpu.func @test_func_dereferenceable(%arg0 : memref<f32> {llvm.dereferenceable = 4 : i64} ) {
+    gpu.return
+  }
+}
+
+// CHECK-LABEL:  llvm.func @test_func_dereferenceable
+// ROCDL-SAME:  !llvm.ptr {llvm.dereferenceable = 4 : i64}
+//  NVVM-SAME:  !llvm.ptr {llvm.dereferenceable = 4 : i64}
+
+
+// -----
+
+gpu.module @kernel {
+  gpu.func @test_func_dereferenceable_or_null(%arg0 : memref<f32> {llvm.dereferenceable_or_null = 4 : i64} ) {
+    gpu.return
+  }
+}
+
+// CHECK-LABEL:  llvm.func @test_func_dereferenceable_or_null
+// ROCDL-SAME:  !llvm.ptr {llvm.dereferenceable_or_null = 4 : i64}
+//  NVVM-SAME:  !llvm.ptr {llvm.dereferenceable_or_null = 4 : i64}

mlir/test/Conversion/GPUCommon/memref-arg-noalias-attrs.mlir

+// RUN: mlir-opt %s -split-input-file -convert-gpu-to-rocdl='use-opaque-pointers=1 use-bare-ptr-memref-call-conv=1' | FileCheck %s --check-prefixes=CHECK,ROCDL
+// RUN: mlir-opt %s -split-input-file -convert-gpu-to-nvvm='use-opaque-pointers=1 use-bare-ptr-memref-call-conv=1' | FileCheck %s --check-prefixes=CHECK,NVVM
+
+gpu.module @kernel {
+  gpu.func @func_with_noalias_attr(%arg0 : memref<f32> {llvm.noalias} ) {
+    gpu.return
+  }
+}
+
+// CHECK-LABEL:  llvm.func @func_with_noalias_attr
+// ROCDL-SAME:  !llvm.ptr {llvm.noalias}
+//  NVVM-SAME:  !llvm.ptr {llvm.noalias}
+
+
+// -----
+
+gpu.module @kernel {
+  gpu.func @func_without_any_attr(%arg0 : memref<f32> ) {
+    gpu.return
+  }
+}
+
+// CHECK-LABEL:  llvm.func @func_without_any_attr
+// ROCDL-SAME:  !llvm.ptr
+//  NVVM-SAME:  !llvm.ptr

mlir/test/Conversion/GPUCommon/memref-arg-noalias-warning.mlir

+// RUN: mlir-opt %s -split-input-file -convert-gpu-to-rocdl='use-opaque-pointers=1 use-bare-ptr-memref-call-conv=0' -verify-diagnostics
+
+gpu.module @kernel {
+// expected-warning @+1 {{Cannot copy noalias with non-bare pointers.}}
+  gpu.func @func_warning_for_not_bare_pointer(%arg0 : memref<f32> {llvm.noalias} ) {
+    gpu.return
+  }
+}