[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.

Author: krzysz00

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"; }

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()},

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));
   }

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();

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 {

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})