[Mlir-commits] [mlir] [MLIR][LLVM] Add explicit target_cpu attribute to llvm.func (PR #78287)

[Sergio Afonso]

https://github.com/skatrak created https://github.com/llvm/llvm-project/pull/78287

This patch adds the target_cpu attribute to llvm.func MLIR operations and updates the translation to/from LLVM IR to match "target-cpu" function attributes.

>From ae33bcc9b9f8b6f4485a1084f6c0bf72c8fe1eb3 Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof at amd.com>
Date: Tue, 16 Jan 2024 12:24:36 +0000
Subject: [PATCH] [MLIR][LLVM] Add explicit target_cpu attribute to llvm.func

This patch adds the target_cpu attribute to llvm.func MLIR operations and
updates the translation to/from LLVM IR to match "target-cpu" function
attributes.
---
 mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td            |  1 +
 mlir/lib/Target/LLVMIR/ModuleImport.cpp                |  5 +++++
 mlir/lib/Target/LLVMIR/ModuleTranslation.cpp           |  3 +++
 mlir/test/Conversion/FuncToLLVM/convert-funcs.mlir     | 10 ++++++++++
 .../Transforms/one-shot-module-bufferize-analysis.mlir |  2 +-
 mlir/test/Target/LLVMIR/Import/target-cpu.ll           |  9 +++++++++
 mlir/test/Target/LLVMIR/target-cpu.mlir                |  7 +++++++
 7 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Target/LLVMIR/Import/target-cpu.ll
 create mode 100644 mlir/test/Target/LLVMIR/target-cpu.mlir

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 7f4e21e5af48fa5..01d476f530b1c57 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -1427,6 +1427,7 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     OptionalAttr<I64Attr>:$alignment,
     OptionalAttr<LLVM_VScaleRangeAttr>:$vscale_range,
     OptionalAttr<FramePointerKindAttr>:$frame_pointer,
+    OptionalAttr<StrAttr>:$target_cpu,
     OptionalAttr<LLVM_TargetFeaturesAttr>:$target_features
   );
 
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index 0eae42b86831b2e..565c16641688ec4 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1744,6 +1744,11 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
                                  .value()));
   }
 
+  if (llvm::Attribute attr = func->getFnAttribute("target-cpu");
+      attr.isStringAttribute()) {
+    funcOp.setTargetCpuAttr(StringAttr::get(context, attr.getValueAsString()));
+  }
+
   if (llvm::Attribute attr = func->getFnAttribute("target-features");
       attr.isStringAttribute()) {
     funcOp.setTargetFeaturesAttr(
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 2763a0fdd62aba1..1499a71c7c9a7ce 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -1104,6 +1104,9 @@ LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
   if (func.getArmPreservesZa())
     llvmFunc->addFnAttr("aarch64_pstate_za_preserved");
 
+  if (auto targetCpu = func.getTargetCpu())
+    llvmFunc->addFnAttr("target-cpu", *targetCpu);
+
   if (auto targetFeatures = func.getTargetFeatures())
     llvmFunc->addFnAttr("target-features", targetFeatures->getFeaturesString());
 
diff --git a/mlir/test/Conversion/FuncToLLVM/convert-funcs.mlir b/mlir/test/Conversion/FuncToLLVM/convert-funcs.mlir
index 765d8469f3c5618..dbbda8e1513a854 100644
--- a/mlir/test/Conversion/FuncToLLVM/convert-funcs.mlir
+++ b/mlir/test/Conversion/FuncToLLVM/convert-funcs.mlir
@@ -61,6 +61,16 @@ func.func @variadic_func(%arg0: i32) attributes { "func.varargs" = true } {
   return
 }
 
+// CHECK-LABEL: llvm.func @target_cpu()
+// CHECK-SAME: target_cpu = "gfx90a"
+func.func private @target_cpu() attributes { "target_cpu" = "gfx90a" }
+
+// CHECK-LABEL: llvm.func @target_features()
+// CHECK-SAME: target_features = #llvm.target_features<["+sme", "+sve"]>
+func.func private @target_features() attributes {
+  "target_features" = #llvm.target_features<["+sme", "+sve"]>
+}
+
 // -----
 
 // CHECK-LABEL: llvm.func @private_callee
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir
index a103e65affacd85..6e7b113aa35c0b1 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir
@@ -772,7 +772,7 @@ func.func @insert_slice_chain(
 // CHECK-SAME: bufferization.access = "none"
     %arg2: tensor<62x90xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
 // CHECK-SAME: bufferization.access = "write"
-  -> tensor<62x90xf32> attributes {passthrough = [["target-cpu", "skylake-avx512"], ["prefer-vector-width", "512"]]}
+  -> tensor<62x90xf32> attributes {passthrough = [["prefer-vector-width", "512"]], target_cpu = "skylake-avx512"}
 {
   %c0 = arith.constant 0 : index
   %cst = arith.constant 0.000000e+00 : f32
diff --git a/mlir/test/Target/LLVMIR/Import/target-cpu.ll b/mlir/test/Target/LLVMIR/Import/target-cpu.ll
new file mode 100644
index 000000000000000..84c0f96c267a11f
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/Import/target-cpu.ll
@@ -0,0 +1,9 @@
+; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s
+
+; CHECK-LABEL: llvm.func @target_cpu()
+; CHECK-SAME: target_cpu = "gfx90a"
+define void @target_cpu() #0 {
+  ret void
+}
+
+attributes #0 = { "target-cpu"="gfx90a" }
diff --git a/mlir/test/Target/LLVMIR/target-cpu.mlir b/mlir/test/Target/LLVMIR/target-cpu.mlir
new file mode 100644
index 000000000000000..80f8172143259a4
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/target-cpu.mlir
@@ -0,0 +1,7 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// CHECK: define void @target_cpu() #[[ATTRS:.*]] {
+// CHECK: attributes #[[ATTRS]] = { "target-cpu"="gfx90a" }
+llvm.func @target_cpu() attributes {target_cpu = "gfx90a"} {
+  llvm.return
+}