diff --git a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp --- a/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp +++ b/mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp @@ -65,6 +65,7 @@ // Create the new function operation. Only copy those attributes that are // not specific to function modeling. SmallVector attributes; + ArrayAttr argAttrs; for (const auto &attr : gpuFuncOp->getAttrs()) { if (attr.getName() == SymbolTable::getSymbolAttrName() || attr.getName() == gpuFuncOp.getFunctionTypeAttrName() || @@ -73,6 +74,10 @@ 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 @@ -188,6 +193,49 @@ } } + // 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(); + NamedAttrList argAttr = argAttrs + ? argAttrs[en.index()].cast() + : 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()) { + 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(); } diff --git a/mlir/test/Conversion/GPUCommon/memref-arg-attrs.mlir b/mlir/test/Conversion/GPUCommon/memref-arg-attrs.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Conversion/GPUCommon/memref-arg-attrs.mlir @@ -0,0 +1,64 @@ +// 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 {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 {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 {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 {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 {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} diff --git a/mlir/test/Conversion/GPUCommon/memref-arg-noalias-attrs.mlir b/mlir/test/Conversion/GPUCommon/memref-arg-noalias-attrs.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Conversion/GPUCommon/memref-arg-noalias-attrs.mlir @@ -0,0 +1,25 @@ +// 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 {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 ) { + gpu.return + } +} + +// CHECK-LABEL: llvm.func @func_without_any_attr +// ROCDL-SAME: !llvm.ptr +// NVVM-SAME: !llvm.ptr diff --git a/mlir/test/Conversion/GPUCommon/memref-arg-noalias-warning.mlir b/mlir/test/Conversion/GPUCommon/memref-arg-noalias-warning.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Conversion/GPUCommon/memref-arg-noalias-warning.mlir @@ -0,0 +1,8 @@ +// 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 {llvm.noalias} ) { + gpu.return + } +}