[MLIR] Add support for frame pointers in MLIR (#72145)

Add support for frame pointers in MLIR.

---------

Co-authored-by: Markus Böck <[email protected]>
Co-authored-by: Christian Ulmann <[email protected]>

Author: 3 people

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td

@@ -48,6 +48,15 @@ def LinkageAttr : LLVM_Attr<"Linkage", "linkage"> {
   let assemblyFormat = "`<` $linkage `>`";
 }
 
+//===----------------------------------------------------------------------===//
+// FramePointerKindAttr
+//===----------------------------------------------------------------------===//
+
+def FramePointerKindAttr : LLVM_Attr<"FramePointerKind", "framePointerKind"> {
+  let parameters = (ins "framePointerKind::FramePointerKind":$framePointerKind);
+  let assemblyFormat = "`<` $framePointerKind `>`";
+}
+
 //===----------------------------------------------------------------------===//
 // Loop Attributes
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/LLVMIR/LLVMEnums.td

@@ -587,7 +587,6 @@ def Linkage : DialectAttr<
           "::mlir::LLVM::LinkageAttr::get($_builder.getContext(), $0)";
 }
 
-
 //===----------------------------------------------------------------------===//
 // Comdat
 //===----------------------------------------------------------------------===//
@@ -664,4 +663,23 @@ def ModRefInfoEnum : LLVM_EnumAttr<
   let cppNamespace = "::mlir::LLVM";
 }
 
+//===----------------------------------------------------------------------===//
+// FramePointerKind
+//===----------------------------------------------------------------------===//
+
+def FramePointerKindNone
+    : LLVM_EnumAttrCase<"None", "none", "None", 0>;
+def FramePointerKindNonLeaf
+    : LLVM_EnumAttrCase<"NonLeaf", "non-leaf", "NonLeaf", 1>;
+def FramePointerKindAll
+    : LLVM_EnumAttrCase<"All", "all", "All", 2>;
+
+def FramePointerKindEnum : LLVM_EnumAttr<
+    "FramePointerKind",
+    "::llvm::FramePointerKind",
+    "LLVM FramePointerKind",
+    [FramePointerKindNone, FramePointerKindNonLeaf, FramePointerKindAll]> {
+  let cppNamespace = "::mlir::LLVM::framePointerKind";
+}
+
 #endif // LLVMIR_ENUMS

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

@@ -1393,7 +1393,8 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     OptionalAttr<StrAttr>:$section,
     OptionalAttr<UnnamedAddr>:$unnamed_addr,
     OptionalAttr<I64Attr>:$alignment,
-    OptionalAttr<LLVM_VScaleRangeAttr>:$vscale_range
+    OptionalAttr<LLVM_VScaleRangeAttr>:$vscale_range,
+    OptionalAttr<FramePointerKindAttr>:$frame_pointer
   );
 
   let regions = (region AnyRegion:$body);

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -1626,6 +1626,7 @@ static constexpr std::array ExplicitAttributes{
     StringLiteral("aarch64_pstate_sm_body"),
     StringLiteral("aarch64_pstate_za_new"),
     StringLiteral("vscale_range"),
+    StringLiteral("frame-pointer"),
 };
 
 static void processPassthroughAttrs(llvm::Function *func, LLVMFuncOp funcOp) {
@@ -1706,6 +1707,16 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
         context, IntegerAttr::get(intTy, attr.getVScaleRangeMin()),
         IntegerAttr::get(intTy, attr.getVScaleRangeMax().value_or(0))));
   }
+
+  // Process frame-pointer attribute.
+  if (func->hasFnAttribute("frame-pointer")) {
+    StringRef stringRefFramePointerKind =
+        func->getFnAttribute("frame-pointer").getValueAsString();
+    funcOp.setFramePointerAttr(LLVM::FramePointerKindAttr::get(
+        funcOp.getContext(), LLVM::framePointerKind::symbolizeFramePointerKind(
+                                 stringRefFramePointerKind)
+                                 .value()));
+  }
 }
 
 DictionaryAttr

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -973,6 +973,12 @@ LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
         getLLVMContext(), attr->getMinRange().getInt(),
         attr->getMaxRange().getInt()));
 
+  // Add function attribute frame-pointer, if found.
+  if (FramePointerKindAttr attr = func.getFramePointerAttr())
+    llvmFunc->addFnAttr("frame-pointer",
+                        LLVM::framePointerKind::stringifyFramePointerKind(
+                            (attr.getFramePointerKind())));
+
   // First, create all blocks so we can jump to them.
   llvm::LLVMContext &llvmContext = llvmFunc->getContext();
   for (auto &bb : func) {

mlir/test/Dialect/LLVMIR/func.mlir

@@ -249,6 +249,12 @@ module {
     // CHECK-SAME: vscale_range(1, 2)
     llvm.return
   }
+
+  // CHECK-LABEL: @frame_pointer_roundtrip()
+  // CHECK-SAME: attributes {frame_pointer = #llvm.framePointerKind<"non-leaf">}
+  llvm.func @frame_pointer_roundtrip() attributes {frame_pointer = #llvm.framePointerKind<"non-leaf">} {
+    llvm.return
+  }
 }
 
 // -----

mlir/test/Target/LLVMIR/Import/frame-pointer.ll

@@ -0,0 +1,8 @@
+; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s
+
+; CHECK-LABEL: llvm.func @frame_pointer_func
+; CHECK-SAME: attributes {frame_pointer = #llvm.framePointerKind<"non-leaf">}
+
+define void @frame_pointer_func() "frame-pointer"="non-leaf" {
+  ret void
+}

mlir/test/Target/LLVMIR/frame-pointer.mlir

@@ -0,0 +1,8 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: define void @frame_pointer_func() 
+// CHECK-SAME: #[[ATTRS:[0-9]+]]
+llvm.func @frame_pointer_func() attributes {frame_pointer = #llvm.framePointerKind<"non-leaf">}  {
+  llvm.return
+}
+// CHECK: attributes #[[ATTRS]] = { "frame-pointer"="non-leaf" }