[MLIR][LLVM] Promote noinline/alwaysinline/optnone out of passthrough (…

…#95110) The `noinline`, `alwaysinline`, and `optnone` function attributes are already being used in MLIR code for the LLVM inlining interface and in some SPIR-V lowering, despite residing in the passthrough dictionary, which is intended as exactly that -- a pass through MLIR -- and not to model any actual semantics being handled in MLIR itself. Promote the `noinline`, `alwaysinline`, and `optnone` attributes out of the passthrough dictionary on `llvm.func` into first class unit attributes, updating the import and export accordingly. Add a verifier to `llvm.func` that checks that these attributes are not set in an incompatible way according to the LLVM specification. Update the LLVM dialect inlining interface to use the first class attributes to check whether inlining is possible.
llvm · Jun 12, 2024 · c012e48 · c012e48
1 parent 93d4fb0
commit c012e48
Show file tree

Hide file tree

Showing 13 changed files with 149 additions and 44 deletions.
diff --git a/mlir/docs/Dialects/LLVM.md b/mlir/docs/Dialects/LLVM.md
@@ -179,7 +179,7 @@ Example:
 
 ```mlir
 llvm.func @func() attributes {
-  passthrough = ["noinline",           // value-less attribute
+  passthrough = ["readonly",           // value-less attribute
                  ["alignstack", "4"],  // integer attribute with value
                  ["other", "attr"]]    // attribute unknown to LLVM
 } {

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -1452,7 +1452,10 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     OptionalAttr<BoolAttr>:$no_infs_fp_math,
     OptionalAttr<BoolAttr>:$no_nans_fp_math,
     OptionalAttr<BoolAttr>:$approx_func_fp_math,
-    OptionalAttr<BoolAttr>:$no_signed_zeros_fp_math
+    OptionalAttr<BoolAttr>:$no_signed_zeros_fp_math,
+    OptionalAttr<UnitAttr>:$no_inline,
+    OptionalAttr<UnitAttr>:$always_inline,
+    OptionalAttr<UnitAttr>:$optimize_none
   );
 
   let regions = (region AnyRegion:$body);
@@ -1490,6 +1493,15 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     /// Returns the callable region, which is the function body. If the function
     /// is external, returns null.
     Region *getCallableRegion();
+
+    /// Returns true if the `no_inline` attribute is set, false otherwise.
+    bool isNoInline() { return bool(getNoInlineAttr()); }
+
+    /// Returns true if the `always_inline` attribute is set, false otherwise.
+    bool isAlwaysInline() { return bool(getAlwaysInlineAttr()); }
+
+    /// Returns true if the `optimize_none` attribute is set, false otherwise.
+    bool isOptimizeNone() { return bool(getOptimizeNoneAttr()); }
   }];
 
   let hasCustomAssemblyFormat = 1;

diff --git a/mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp b/mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp
@@ -1425,15 +1425,18 @@ class FuncConversionPattern : public SPIRVToLLVMConversion<spirv::FuncOp> {
     // Convert SPIR-V Function Control to equivalent LLVM function attribute
     MLIRContext *context = funcOp.getContext();
     switch (funcOp.getFunctionControl()) {
+    case spirv::FunctionControl::Inline:
+      newFuncOp.setAlwaysInline(true);
+      break;
+    case spirv::FunctionControl::DontInline:
+      newFuncOp.setNoInline(true);
+      break;
+
 #define DISPATCH(functionControl, llvmAttr)                                    \
   case functionControl:                                                        \
     newFuncOp->setAttr("passthrough", ArrayAttr::get(context, {llvmAttr}));    \
     break;
 
-      DISPATCH(spirv::FunctionControl::Inline,
-               StringAttr::get(context, "alwaysinline"));
-      DISPATCH(spirv::FunctionControl::DontInline,
-               StringAttr::get(context, "noinline"));
       DISPATCH(spirv::FunctionControl::Pure,
                StringAttr::get(context, "readonly"));
       DISPATCH(spirv::FunctionControl::Const,

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -2492,6 +2492,13 @@ LogicalResult LLVMFuncOp::verify() {
     return success();
   }
 
+  // In LLVM IR, these attributes are composed by convention, not by design.
+  if (isNoInline() && isAlwaysInline())
+    return emitError("no_inline and always_inline attributes are incompatible");
+
+  if (isOptimizeNone() && !isNoInline())
+    return emitOpError("with optimize_none must also be no_inline");
+
   Type landingpadResultTy;
   StringRef diagnosticMessage;
   bool isLandingpadTypeConsistent =

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
@@ -690,8 +690,6 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
         // Cache set of StringAttrs for fast lookup in `isLegalToInline`.
         disallowedFunctionAttrs({
             StringAttr::get(dialect->getContext(), "noduplicate"),
-            StringAttr::get(dialect->getContext(), "noinline"),
-            StringAttr::get(dialect->getContext(), "optnone"),
             StringAttr::get(dialect->getContext(), "presplitcoroutine"),
             StringAttr::get(dialect->getContext(), "returns_twice"),
             StringAttr::get(dialect->getContext(), "strictfp"),
@@ -702,14 +700,20 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
     if (!wouldBeCloned)
       return false;
     if (!isa<LLVM::CallOp>(call)) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "Cannot inline: call is not an LLVM::CallOp\n");
+      LLVM_DEBUG(llvm::dbgs() << "Cannot inline: call is not an '"
+                              << LLVM::CallOp::getOperationName() << "' op\n");
       return false;
     }
     auto funcOp = dyn_cast<LLVM::LLVMFuncOp>(callable);
     if (!funcOp) {
       LLVM_DEBUG(llvm::dbgs()
-                 << "Cannot inline: callable is not an LLVM::LLVMFuncOp\n");
+                 << "Cannot inline: callable is not an '"
+                 << LLVM::LLVMFuncOp::getOperationName() << "' op\n");
+      return false;
+    }
+    if (funcOp.isNoInline()) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << "Cannot inline: function is marked no_inline\n");
       return false;
     }
     if (funcOp.isVarArg()) {

diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1664,23 +1664,26 @@ static void processMemoryEffects(llvm::Function *func, LLVMFuncOp funcOp) {
 
 // List of LLVM IR attributes that map to an explicit attribute on the MLIR
 // LLVMFuncOp.
-static constexpr std::array ExplicitAttributes{
-    StringLiteral("aarch64_pstate_sm_enabled"),
-    StringLiteral("aarch64_pstate_sm_body"),
-    StringLiteral("aarch64_pstate_sm_compatible"),
-    StringLiteral("aarch64_new_za"),
-    StringLiteral("aarch64_preserves_za"),
+static constexpr std::array kExplicitAttributes{
     StringLiteral("aarch64_in_za"),
-    StringLiteral("aarch64_out_za"),
     StringLiteral("aarch64_inout_za"),
-    StringLiteral("vscale_range"),
+    StringLiteral("aarch64_new_za"),
+    StringLiteral("aarch64_out_za"),
+    StringLiteral("aarch64_preserves_za"),
+    StringLiteral("aarch64_pstate_sm_body"),
+    StringLiteral("aarch64_pstate_sm_compatible"),
+    StringLiteral("aarch64_pstate_sm_enabled"),
+    StringLiteral("alwaysinline"),
+    StringLiteral("approx-func-fp-math"),
     StringLiteral("frame-pointer"),
-    StringLiteral("target-features"),
-    StringLiteral("unsafe-fp-math"),
     StringLiteral("no-infs-fp-math"),
     StringLiteral("no-nans-fp-math"),
-    StringLiteral("approx-func-fp-math"),
     StringLiteral("no-signed-zeros-fp-math"),
+    StringLiteral("noinline"),
+    StringLiteral("optnone"),
+    StringLiteral("target-features"),
+    StringLiteral("unsafe-fp-math"),
+    StringLiteral("vscale_range"),
 };
 
 static void processPassthroughAttrs(llvm::Function *func, LLVMFuncOp funcOp) {
@@ -1709,7 +1712,7 @@ static void processPassthroughAttrs(llvm::Function *func, LLVMFuncOp funcOp) {
     auto keyAttr = StringAttr::get(context, attrName);
 
     // Skip attributes that map to an explicit attribute on the LLVMFuncOp.
-    if (llvm::is_contained(ExplicitAttributes, attrName))
+    if (llvm::is_contained(kExplicitAttributes, attrName))
       continue;
 
     if (attr.isStringAttribute()) {
@@ -1745,6 +1748,13 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
   processMemoryEffects(func, funcOp);
   processPassthroughAttrs(func, funcOp);
 
+  if (func->hasFnAttribute(llvm::Attribute::NoInline))
+    funcOp.setNoInline(true);
+  if (func->hasFnAttribute(llvm::Attribute::AlwaysInline))
+    funcOp.setAlwaysInline(true);
+  if (func->hasFnAttribute(llvm::Attribute::OptimizeNone))
+    funcOp.setOptimizeNone(true);
+
   if (func->hasFnAttribute("aarch64_pstate_sm_enabled"))
     funcOp.setArmStreaming(true);
   else if (func->hasFnAttribute("aarch64_pstate_sm_body"))

diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -1390,10 +1390,10 @@ LogicalResult ModuleTranslation::convertDialectAttributes(
   return success();
 }
 
-/// Converts the function attributes from LLVMFuncOp and attaches them to the
-/// llvm::Function.
-static void convertFunctionAttributes(LLVMFuncOp func,
-                                      llvm::Function *llvmFunc) {
+/// Converts memory effect attributes from `func` and attaches them to
+/// `llvmFunc`.
+static void convertFunctionMemoryAttributes(LLVMFuncOp func,
+                                            llvm::Function *llvmFunc) {
   if (!func.getMemory())
     return;
 
@@ -1412,6 +1412,18 @@ static void convertFunctionAttributes(LLVMFuncOp func,
   llvmFunc->setMemoryEffects(newMemEffects);
 }
 
+/// Converts function attributes from `func` and attaches them to `llvmFunc`.
+static void convertFunctionAttributes(LLVMFuncOp func,
+                                      llvm::Function *llvmFunc) {
+  if (func.getNoInlineAttr())
+    llvmFunc->addFnAttr(llvm::Attribute::NoInline);
+  if (func.getAlwaysInlineAttr())
+    llvmFunc->addFnAttr(llvm::Attribute::AlwaysInline);
+  if (func.getOptimizeNoneAttr())
+    llvmFunc->addFnAttr(llvm::Attribute::OptimizeNone);
+  convertFunctionMemoryAttributes(func, llvmFunc);
+}
+
 FailureOr<llvm::AttrBuilder>
 ModuleTranslation::convertParameterAttrs(LLVMFuncOp func, int argIdx,
                                          DictionaryAttr paramAttrs) {

diff --git a/mlir/test/Conversion/SPIRVToLLVM/func-ops-to-llvm.mlir b/mlir/test/Conversion/SPIRVToLLVM/func-ops-to-llvm.mlir
@@ -29,12 +29,12 @@ spirv.func @none() "None" {
   spirv.Return
 }
 
-// CHECK-LABEL: llvm.func @inline() attributes {passthrough = ["alwaysinline"]}
+// CHECK-LABEL: llvm.func @inline() attributes {always_inline}
 spirv.func @inline() "Inline" {
   spirv.Return
 }
 
-// CHECK-LABEL: llvm.func @dont_inline() attributes {passthrough = ["noinline"]}
+// CHECK-LABEL: llvm.func @dont_inline() attributes {no_inline}
 spirv.func @dont_inline() "DontInline" {
   spirv.Return
 }

diff --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -90,12 +90,12 @@ llvm.func @caller() -> (i32) {
 
 // -----
 
-llvm.func @foo() -> (i32) attributes { passthrough = ["noinline"] } {
+llvm.func @foo() -> (i32) attributes { no_inline } {
   %0 = llvm.mlir.constant(0 : i32) : i32
   llvm.return %0 : i32
 }
 
-llvm.func @bar() -> (i32) attributes { passthrough = ["noinline"] } {
+llvm.func @bar() -> (i32) attributes { no_inline } {
   %0 = llvm.mlir.constant(1 : i32) : i32
   llvm.return %0 : i32
 }
@@ -161,11 +161,7 @@ llvm.func @caller() {
 
 // -----
 
-llvm.func @callee_noinline() attributes { passthrough = ["noinline"] } {
-  llvm.return
-}
-
-llvm.func @callee_optnone() attributes { passthrough = ["optnone"] } {
+llvm.func @callee_noinline() attributes { no_inline } {
   llvm.return
 }
 
@@ -187,15 +183,13 @@ llvm.func @callee_strictfp() attributes { passthrough = ["strictfp"] } {
 
 // CHECK-LABEL: llvm.func @caller
 // CHECK-NEXT: llvm.call @callee_noinline
-// CHECK-NEXT: llvm.call @callee_optnone
 // CHECK-NEXT: llvm.call @callee_noduplicate
 // CHECK-NEXT: llvm.call @callee_presplitcoroutine
 // CHECK-NEXT: llvm.call @callee_returns_twice
 // CHECK-NEXT: llvm.call @callee_strictfp
 // CHECK-NEXT: llvm.return
 llvm.func @caller() {
   llvm.call @callee_noinline() : () -> ()
-  llvm.call @callee_optnone() : () -> ()
   llvm.call @callee_noduplicate() : () -> ()
   llvm.call @callee_presplitcoroutine() : () -> ()
   llvm.call @callee_returns_twice() : () -> ()

diff --git a/mlir/test/Dialect/LLVMIR/invalid.mlir b/mlir/test/Dialect/LLVMIR/invalid.mlir
@@ -1472,3 +1472,17 @@ func.func @tma_load(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<3>, %barrier: !
   nvvm.cp.async.bulk.tensor.shared.cluster.global %dest, %tmaDescriptor,  %barrier, box[%crd0,%crd1,%crd2,%crd3,%crd0,%crd1,%crd2,%crd3]: !llvm.ptr<3>, !llvm.ptr  
   return
 }
+
+// -----
+
+// expected-error @below {{no_inline and always_inline attributes are incompatible}}
+llvm.func @alwaysinline_noinline() attributes { always_inline, no_inline } {
+  llvm.return
+}
+
+// -----
+
+// expected-error @below {{'llvm.func' op with optimize_none must also be no_inline}}
+llvm.func @optnone_requires_noinline() attributes { optimize_none } {
+  llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/Import/function-attributes.ll b/mlir/test/Target/LLVMIR/Import/function-attributes.ll
@@ -163,11 +163,10 @@ define void @func_memory() memory(readwrite, argmem: none) {
 ; CHECK-LABEL: @passthrough_combined
 ; CHECK-SAME: attributes {passthrough = [
 ; CHECK-DAG: ["alignstack", "16"]
-; CHECK-DAG: "noinline"
 ; CHECK-DAG: "probe-stack"
 ; CHECK-DAG: ["alloc-family", "malloc"]
 ; CHECK:   llvm.return
-define void @passthrough_combined() alignstack(16) noinline "probe-stack" "alloc-family"="malloc" {
+define void @passthrough_combined() alignstack(16) "probe-stack" "alloc-family"="malloc" {
   ret void
 }
 
@@ -344,3 +343,21 @@ declare void @func_attr_no_signed_zeros_fp_math_true() "no-signed-zeros-fp-math"
 ; CHECK-LABEL: @func_attr_no_signed_zeros_fp_math_false
 ; CHECK-SAME: attributes {no_signed_zeros_fp_math = false}
 declare void @func_attr_no_signed_zeros_fp_math_false() "no-signed-zeros-fp-math"="false"
+
+// -----
+
+; CHECK-LABEL: @noinline_attribute
+; CHECK-SAME: attributes {no_inline}
+declare void @noinline_attribute() noinline
+
+// -----
+
+; CHECK-LABEL: @alwaysinline_attribute
+; CHECK-SAME: attributes {always_inline}
+declare void @alwaysinline_attribute() alwaysinline
+
+// -----
+
+; CHECK-LABEL: @optnone_attribute
+; CHECK-SAME: attributes {no_inline, optimize_none}
+declare void @optnone_attribute() noinline optnone
diff --git a/mlir/test/Target/LLVMIR/llvmir-invalid.mlir b/mlir/test/Target/LLVMIR/llvmir-invalid.mlir
@@ -68,8 +68,8 @@ llvm.mlir.global internal constant @test([2.5, 7.4]) : !llvm.array<2 x f64>
 
 // -----
 
-// expected-error @below{{LLVM attribute 'noinline' does not expect a value}}
-llvm.func @passthrough_unexpected_value() attributes {passthrough = [["noinline", "42"]]}
+// expected-error @below{{LLVM attribute 'readonly' does not expect a value}}
+llvm.func @passthrough_unexpected_value() attributes {passthrough = [["readonly", "42"]]}
 
 // -----
 

diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -1730,12 +1730,11 @@ llvm.func @callFenceInst() {
 
 // CHECK-LABEL: @passthrough
 // CHECK: #[[ATTR_GROUP:[0-9]*]]
-llvm.func @passthrough() attributes {passthrough = ["noinline", ["alignstack", "4"], "null_pointer_is_valid", ["foo", "bar"]]} {
+llvm.func @passthrough() attributes {passthrough = [["alignstack", "4"], "null_pointer_is_valid", ["foo", "bar"]]} {
   llvm.return
 }
 
 // CHECK: attributes #[[ATTR_GROUP]] = {
-// CHECK-DAG: noinline
 // CHECK-DAG: alignstack=4
 // CHECK-DAG: null_pointer_is_valid
 // CHECK-DAG: "foo"="bar"
@@ -2401,3 +2400,36 @@ llvm.linker_options ["/DEFAULTLIB:", "libcmtd"]
 
 // CHECK: @big_ = common global [4294967296 x i8] zeroinitializer
 llvm.mlir.global common @big_(dense<0> : vector<4294967296xi8>) {addr_space = 0 : i32} : !llvm.array<4294967296 x i8>
+
+// -----
+
+// CHECK-LABEL: @no_inline
+// CHECK-SAME: #[[ATTRS:[0-9]+]]
+llvm.func @no_inline() attributes { no_inline } {
+  llvm.return
+}
+
+// CHECK: #[[ATTRS]]
+// CHECK-SAME: noinline
+
+// -----
+
+// CHECK-LABEL: @always_inline
+// CHECK-SAME: #[[ATTRS:[0-9]+]]
+llvm.func @always_inline() attributes { always_inline } {
+  llvm.return
+}
+
+// CHECK: #[[ATTRS]]
+// CHECK-SAME: alwaysinline
+
+// -----
+
+// CHECK-LABEL: @optimize_none
+// CHECK-SAME: #[[ATTRS:[0-9]+]]
+llvm.func @optimize_none() attributes { no_inline, optimize_none } {
+  llvm.return
+}
+
+// CHECK: #[[ATTRS]]
+// CHECK-SAME: optnone