[mlir][LLVM] Plumb range attributes on parameters and results through

[krzysz00]

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.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Krzysztof Drewniak (krzysz00)

Changes

We've had the ability to define LLVM's range atrbitute through
#llvm.constant_range for some time, and have used this for some GPU
intrinsics. This commit allows using llvm.range as as a parameter or result attribute on function declarations and definitions.

---

Full diff: https://github.com/llvm/llvm-project/pull/117801.diff

6 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td (+1)
• (modified) mlir/lib/Target/LLVMIR/AttrKindDetail.h (+1)
• (modified) mlir/lib/Target/LLVMIR/ModuleImport.cpp (+5-1)
• (modified) mlir/lib/Target/LLVMIR/ModuleTranslation.cpp (+6-1)
• (modified) mlir/test/Target/LLVMIR/Import/function-attributes.ll (+9-1)
• (modified) mlir/test/Target/LLVMIR/llvmir.mlir (+10)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td
index ea82f7f7b8e124..b5526bda9f2de1 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 ddc6d46b90bb2b..aa2e1d04e1b4af 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 0d416a5857facb..b0d5e635248d3f 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<LLVM::ConstantRangeAttr>(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 9e58d2a29199e6..ad62ae0cef57be 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<IntegerAttr>([&](auto intAttr) {
             attrBuilder.addRawIntAttr(llvmKind, intAttr.getInt());
           })
-          .Case<UnitAttr>([&](auto) { attrBuilder.addAttribute(llvmKind); });
+          .Case<UnitAttr>([&](auto) { attrBuilder.addAttribute(llvmKind); })
+          .Case<LLVM::ConstantRangeAttr>([&](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 912f448657baa7..079aa6f90bf11d 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<i64, 0, 4097>}
 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<i64, 0, 4097>})
+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 f75c2e471c94bc..51086bc655af24 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<i64, 0, 4097>})
+
+// -----
+
+// CHECK: declare range(i64 0, 4097) i64 @range_res_function()
+llvm.func @range_res_function() -> (i64 {llvm.range = #llvm.constant_range<i64, 0, 4097>})
+
+// -----
+
 // CHECK: declare void @readonly_function([[PTR:.+]] readonly)
 llvm.func @readonly_function(%arg0: !llvm.ptr {llvm.readonly})
 

[gysit]

LGTM

Thanks for adding!

FYI there are some minor typos in the commit message:

• atrbitute -> attribute
• as as -> as

[Dinistro]

LGTM, thanks for the extension.

[krzysz00]

Thanks for the typo spotting!