From 92a15dd7482ff4e1fae7a07f888564e5b1d53eee Mon Sep 17 00:00:00 2001 From: Krzysztof Drewniak Date: Wed, 27 Nov 2024 12:31:51 -0600 Subject: [PATCH] [mlir][LLVM] Plumb range attributes on parameters and results through (#117801) We've had the ability to define LLVM's `range` attribute through #llvm.constant_range for some time, and have used this for some GPU intrinsics. This commit allows using `llvm.range` as a parameter or result attribute on function declarations and definitions. --- mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td | 1 + mlir/lib/Target/LLVMIR/AttrKindDetail.h | 1 + mlir/lib/Target/LLVMIR/ModuleImport.cpp | 6 +++++- mlir/lib/Target/LLVMIR/ModuleTranslation.cpp | 7 ++++++- mlir/test/Target/LLVMIR/Import/function-attributes.ll | 10 +++++++++- mlir/test/Target/LLVMIR/llvmir.mlir | 10 ++++++++++ 6 files changed, 32 insertions(+), 3 deletions(-) diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td index ea82f7f7b8e1..b5526bda9f2d 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td @@ -52,6 +52,7 @@ def LLVM_Dialect : Dialect { static StringRef getNoFreeAttrName() { return "llvm.nofree"; } static StringRef getNonNullAttrName() { return "llvm.nonnull"; } static StringRef getPreallocatedAttrName() { return "llvm.preallocated"; } + static StringRef getRangeAttrName() { return "llvm.range"; } static StringRef getReadonlyAttrName() { return "llvm.readonly"; } static StringRef getReturnedAttrName() { return "llvm.returned"; } static StringRef getSExtAttrName() { return "llvm.signext"; } diff --git a/mlir/lib/Target/LLVMIR/AttrKindDetail.h b/mlir/lib/Target/LLVMIR/AttrKindDetail.h index ddc6d46b90bb..aa2e1d04e1b4 100644 --- a/mlir/lib/Target/LLVMIR/AttrKindDetail.h +++ b/mlir/lib/Target/LLVMIR/AttrKindDetail.h @@ -45,6 +45,7 @@ getAttrKindToNameMapping() { {llvm::Attribute::AttrKind::NonNull, LLVMDialect::getNonNullAttrName()}, {llvm::Attribute::AttrKind::Preallocated, LLVMDialect::getPreallocatedAttrName()}, + {llvm::Attribute::AttrKind::Range, LLVMDialect::getRangeAttrName()}, {llvm::Attribute::AttrKind::ReadOnly, LLVMDialect::getReadonlyAttrName()}, {llvm::Attribute::AttrKind::ReadNone, LLVMDialect::getReadnoneAttrName()}, {llvm::Attribute::AttrKind::Returned, LLVMDialect::getReturnedAttrName()}, diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp index 0d416a5857fa..b0d5e635248d 100644 --- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp @@ -2022,7 +2022,11 @@ ModuleImport::convertParameterAttribute(llvm::AttributeSet llvmParamAttrs, mlirAttr = builder.getI64IntegerAttr(llvmAttr.getValueAsInt()); else if (llvmAttr.isEnumAttribute()) mlirAttr = builder.getUnitAttr(); - else + else if (llvmAttr.isConstantRangeAttribute()) { + const llvm::ConstantRange &value = llvmAttr.getValueAsConstantRange(); + mlirAttr = builder.getAttr(value.getLower(), + value.getUpper()); + } else llvm_unreachable("unexpected parameter attribute kind"); paramAttrs.push_back(builder.getNamedAttr(mlirName, mlirAttr)); } diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp index 9e58d2a29199..ad62ae0cef57 100644 --- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp @@ -1581,7 +1581,12 @@ ModuleTranslation::convertParameterAttrs(LLVMFuncOp func, int argIdx, .Case([&](auto intAttr) { attrBuilder.addRawIntAttr(llvmKind, intAttr.getInt()); }) - .Case([&](auto) { attrBuilder.addAttribute(llvmKind); }); + .Case([&](auto) { attrBuilder.addAttribute(llvmKind); }) + .Case([&](auto rangeAttr) { + attrBuilder.addConstantRangeAttr( + llvmKind, llvm::ConstantRange(rangeAttr.getLower(), + rangeAttr.getUpper())); + }); } else if (namedAttr.getNameDialect()) { if (failed(iface.convertParameterAttr(func, argIdx, namedAttr, *this))) return failure(); diff --git a/mlir/test/Target/LLVMIR/Import/function-attributes.ll b/mlir/test/Target/LLVMIR/Import/function-attributes.ll index 912f448657ba..079aa6f90bf1 100644 --- a/mlir/test/Target/LLVMIR/Import/function-attributes.ll +++ b/mlir/test/Target/LLVMIR/Import/function-attributes.ll @@ -45,6 +45,7 @@ attributes #0 = { readnone } ; CHECK-SAME: !llvm.ptr {llvm.returned} ; CHECK-SAME: !llvm.ptr {llvm.alignstack = 32 : i64} ; CHECK-SAME: !llvm.ptr {llvm.writeonly} +; CHECK-SAME: i64 {llvm.range = #llvm.constant_range} define ptr @func_arg_attrs( ptr byval(i64) %arg0, ptr byref(i64) %arg1, @@ -63,7 +64,8 @@ define ptr @func_arg_attrs( ptr preallocated(double) %arg16, ptr returned %arg17, ptr alignstack(32) %arg18, - ptr writeonly %arg19) { + ptr writeonly %arg19, + i64 range(i64 0, 4097) %arg20) { ret ptr %arg17 } @@ -141,6 +143,12 @@ declare inreg ptr @func_res_attr_inreg() ; // ----- +; CHECK-LABEL: @func_res_attr_range +; CHECK-SAME: (i64 {llvm.range = #llvm.constant_range}) +declare range(i64 0, 4097) i64 @func_res_attr_range() + +; // ----- + ; CHECK-LABEL: @entry_count ; CHECK-SAME: attributes {function_entry_count = 4242 : i64} define void @entry_count() !prof !1 { diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir index f75c2e471c94..51086bc655af 100644 --- a/mlir/test/Target/LLVMIR/llvmir.mlir +++ b/mlir/test/Target/LLVMIR/llvmir.mlir @@ -2328,6 +2328,16 @@ llvm.func @vararg_function(%arg0: i32, ...) -> i32 { // ----- +// CHECK: declare void @range_arg_function(i64 range(i64 0, 4097)) +llvm.func @range_arg_function(%arg0: i64 {llvm.range = #llvm.constant_range}) + +// ----- + +// CHECK: declare range(i64 0, 4097) i64 @range_res_function() +llvm.func @range_res_function() -> (i64 {llvm.range = #llvm.constant_range}) + +// ----- + // CHECK: declare void @readonly_function([[PTR:.+]] readonly) llvm.func @readonly_function(%arg0: !llvm.ptr {llvm.readonly})