[Mlir-commits] [mlir] [mlir][llvm] Add llvm.target_features features attribute (PR #71510)

Benjamin Maxwell llvmlistbot at llvm.org
Tue Nov 7 04:22:32 PST 2023


https://github.com/MacDue updated https://github.com/llvm/llvm-project/pull/71510

>From 442bbec51c217a30d040e6e1fa6cad50264fe259 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Tue, 7 Nov 2023 10:38:52 +0000
Subject: [PATCH 1/2] [mlir][llvm] Add llvm.target_features features attribute

This patch adds a target_features (TargetFeaturesAttr) to the LLVM
dialect to allow setting and querying the features in use on a function.

The features are stored as a sorted list rather plain string to allow
efficiently checking a function's features.

The motivation for this comes from the Arm SME dialect where we would
like a convenient way to check what variants of an operation are
available based on the CPU features.

Intended usage:

The target_features attribute is populated manually or by a pass:

```mlir
func.func @example() attributes {
   target_features = #llvm.target_features<"+sme,+sve,+sme-f64f64">
} {
 // ...
}
```

Then within a later rewrite the attribute can be checked, and used to
make lowering decisions.

```c++
// Finds the "target_features" attribute on the parent
// FunctionOpInterface.
auto targetFeatures = LLVM::TargetFeaturesAttr::featuresAt(op);

// Check a feature.
// Returns false if targetFeatures is null or the feature is not in
// the list.
if (!targetFeatures.contains("+sme-f64f64"))
    return failure();
```

For now, this is rather simple just checks if the exact feature is in
the list, though it could be possible to extend with implied features
using information from LLVM.
---
 .../mlir/Dialect/LLVMIR/LLVMAttrDefs.td       | 53 ++++++++++++++
 mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h  | 25 +++++++
 mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td   |  3 +-
 mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp      | 69 +++++++++++++++++++
 mlir/lib/Target/LLVMIR/ModuleImport.cpp       | 25 +++++--
 mlir/lib/Target/LLVMIR/ModuleTranslation.cpp  |  3 +
 .../Target/LLVMIR/Import/target-features.ll   |  9 +++
 mlir/test/Target/LLVMIR/target-features.mlir  |  7 ++
 8 files changed, 186 insertions(+), 8 deletions(-)
 create mode 100644 mlir/test/Target/LLVMIR/Import/target-features.ll
 create mode 100644 mlir/test/Target/LLVMIR/target-features.mlir

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index 0d9af88157a8ae1..e0f4518abf3c01f 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -860,4 +860,57 @@ def LLVM_VScaleRangeAttr : LLVM_Attr<"VScaleRange", "vscale_range"> {
     "IntegerAttr":$maxRange);
   let assemblyFormat = "`<` struct(params) `>`";
 }
+
+//===----------------------------------------------------------------------===//
+// TargetFeaturesAttr
+//===----------------------------------------------------------------------===//
+
+def LLVM_TargetFeaturesAttr : LLVM_Attr<"TargetFeatures", "target_features"> {
+  let summary = "LLVM target features attribute";
+
+  let description = [{
+    Represents the LLVM target features in a manner that is efficient to query.
+
+    Example:
+    ```mlir
+    #llvm.target_features<"+sme,+sve,+sme-f64f64">
+    ```
+
+    Then within a pass or rewrite the features active at an op can be queried:
+
+    ```c++
+    auto targetFeatures = LLVM::TargetFeaturesAttr::featuresAt(op);
+
+    if (!targetFeatures.contains("+sme-f64f64"))
+      return failure();
+    ```
+  }];
+
+  let parameters = (ins
+    ArrayRefOfSelfAllocationParameter<"TargetFeature", "">: $features);
+
+  let builders = [
+    TypeBuilder<(ins "::llvm::ArrayRef<TargetFeature>": $features)>,
+    TypeBuilder<(ins "StringRef": $features)>
+  ];
+
+  let extraClassDeclaration = [{
+    /// Checks if a feature is contained within the features list.
+    bool contains(TargetFeature) const;
+    bool contains(StringRef feature) const {
+      return contains(TargetFeature{feature});
+    }
+
+    /// Returns the list of features as an LLVM-compatible string.
+    std::string getFeaturesString() const;
+
+    /// Finds the target features on the parent FunctionOpInterface.
+    /// Note: This assumes the attribute is called "target_features".
+    static TargetFeaturesAttr featuresAt(Operation* op);
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+  let skipDefaultBuilders = 1;
+}
+
 #endif // LLVMIR_ATTRDEFS
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
index c370bfa2b733d65..68a3d1f74175bda 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
@@ -74,6 +74,31 @@ class TBAANodeAttr : public Attribute {
   }
 };
 
+/// This struct represents a single LLVM target feature.
+struct TargetFeature {
+  StringRef feature;
+
+  // Support allocating this struct into MLIR storage to ensure the feature
+  // string remains valid.
+  TargetFeature allocateInto(TypeStorageAllocator &alloc) const {
+    return TargetFeature{alloc.copyInto(feature)};
+  }
+
+  operator StringRef() const { return feature; }
+
+  bool operator==(TargetFeature const &other) const {
+    return feature == other.feature;
+  }
+
+  bool operator<(TargetFeature const &other) const {
+    return feature < other.feature;
+  }
+};
+
+inline llvm::hash_code hash_value(const TargetFeature &feature) {
+  return llvm::hash_value(feature.feature);
+}
+
 // Inline the LLVM generated Linkage enum and utility.
 // This is only necessary to isolate the "enum generated code" from the
 // attribute definition itself.
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 7928bb940ad8bae..a6338485a52d9f5 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -1388,7 +1388,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<LLVM_TargetFeaturesAttr>:$target_features
   );
 
   let regions = (region AnyRegion:$body);
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
index 3d45ab8fac4d705..e97c30487210cd8 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
@@ -14,10 +14,13 @@
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/BinaryFormat/Dwarf.h"
+#include <algorithm>
 #include <optional>
+#include <set>
 
 using namespace mlir;
 using namespace mlir::LLVM;
@@ -109,3 +112,69 @@ bool MemoryEffectsAttr::isReadWrite() {
     return false;
   return true;
 }
+
+//===----------------------------------------------------------------------===//
+// TargetFeaturesAttr
+//===----------------------------------------------------------------------===//
+
+Attribute TargetFeaturesAttr::parse(mlir::AsmParser &parser, mlir::Type) {
+  std::string targetFeatures;
+  if (parser.parseLess() || parser.parseString(&targetFeatures) ||
+      parser.parseGreater())
+    return {};
+  return get(parser.getContext(), targetFeatures);
+}
+
+void TargetFeaturesAttr::print(mlir::AsmPrinter &printer) const {
+  printer << "<\"";
+  llvm::interleave(
+      getFeatures(), printer,
+      [&](auto &feature) { printer << StringRef(feature); }, ",");
+  printer << "\">";
+}
+
+TargetFeaturesAttr
+TargetFeaturesAttr::get(MLIRContext *context,
+                        llvm::ArrayRef<TargetFeature> featuresRef) {
+  // Sort and de-duplicate target features.
+  std::set<TargetFeature> features(featuresRef.begin(), featuresRef.end());
+  return Base::get(context, llvm::to_vector(features));
+}
+
+TargetFeaturesAttr TargetFeaturesAttr::get(MLIRContext *context,
+                                           StringRef targetFeatures) {
+  SmallVector<StringRef> features;
+  StringRef{targetFeatures}.split(features, ',', /*MaxSplit=*/-1,
+                                  /*KeepEmpty=*/false);
+  return get(context, llvm::map_to_vector(features, [](StringRef feature) {
+               return TargetFeature{feature};
+             }));
+}
+
+bool TargetFeaturesAttr::contains(TargetFeature feature) const {
+  if (!bool(*this))
+    return false; // Allows checking null target features.
+  ArrayRef<TargetFeature> features = getFeatures();
+  // Note: The attribute getter ensures the feature list is sorted.
+  return std::binary_search(features.begin(), features.end(), feature);
+}
+
+std::string TargetFeaturesAttr::getFeaturesString() const {
+  std::string features;
+  bool first = true;
+  for (TargetFeature feature : getFeatures()) {
+    if (!first)
+      features += ",";
+    features += StringRef(feature);
+    first = false;
+  }
+  return features;
+}
+
+TargetFeaturesAttr TargetFeaturesAttr::featuresAt(Operation *op) {
+  auto parentFunction = op->getParentOfType<FunctionOpInterface>();
+  if (!parentFunction)
+    return {};
+  return parentFunction.getOperation()->getAttrOfType<TargetFeaturesAttr>(
+      "target_features");
+}
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index e3562049cd81c76..5eed08f371d8089 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1554,6 +1554,15 @@ static void processMemoryEffects(llvm::Function *func, LLVMFuncOp funcOp) {
   funcOp.setMemoryAttr(memAttr);
 }
 
+// List of LLVM IR attributes that map to explicit attribute on the MLIR
+// LLVMFuncOp.
+static constexpr std::array ExplicitAttributes{
+    StringLiteral("aarch64_pstate_sm_enabled"),
+    StringLiteral("aarch64_pstate_sm_body"),
+    StringLiteral("vscale_range"),
+    StringLiteral("target-features"),
+};
+
 static void processPassthroughAttrs(llvm::Function *func, LLVMFuncOp funcOp) {
   MLIRContext *context = funcOp.getContext();
   SmallVector<Attribute> passthroughs;
@@ -1579,11 +1588,8 @@ static void processPassthroughAttrs(llvm::Function *func, LLVMFuncOp funcOp) {
       attrName = llvm::Attribute::getNameFromAttrKind(attr.getKindAsEnum());
     auto keyAttr = StringAttr::get(context, attrName);
 
-    // Skip the aarch64_pstate_sm_<body|enabled> since the LLVMFuncOp has an
-    // explicit attribute.
-    // Also skip the vscale_range, it is also an explicit attribute.
-    if (attrName == "aarch64_pstate_sm_enabled" ||
-        attrName == "aarch64_pstate_sm_body" || attrName == "vscale_range")
+    // Skip attributes that map to an explicit attribute on the LLVMFuncOp.
+    if (llvm::is_contained(ExplicitAttributes, attrName))
       continue;
 
     if (attr.isStringAttribute()) {
@@ -1623,14 +1629,19 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
     funcOp.setArmStreaming(true);
   else if (func->hasFnAttribute("aarch64_pstate_sm_body"))
     funcOp.setArmLocallyStreaming(true);
-  llvm::Attribute attr = func->getFnAttribute(llvm::Attribute::VScaleRange);
-  if (attr.isValid()) {
+  if (llvm::Attribute attr = func->getFnAttribute(llvm::Attribute::VScaleRange);
+      attr.isValid()) {
     MLIRContext *context = funcOp.getContext();
     auto intTy = IntegerType::get(context, 32);
     funcOp.setVscaleRangeAttr(LLVM::VScaleRangeAttr::get(
         context, IntegerAttr::get(intTy, attr.getVScaleRangeMin()),
         IntegerAttr::get(intTy, attr.getVScaleRangeMax().value_or(0))));
   }
+  if (llvm::Attribute attr = func->getFnAttribute("target-features");
+      attr.isStringAttribute()) {
+    funcOp.setTargetFeaturesAttr(
+        LLVM::TargetFeaturesAttr::get(context, attr.getValueAsString()));
+  }
 }
 
 DictionaryAttr
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 7312388bc9b4dd2..3282857a6eb2991 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -890,6 +890,9 @@ LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
   else if (func.getArmLocallyStreaming())
     llvmFunc->addFnAttr("aarch64_pstate_sm_body");
 
+  if (auto targetFeatures = func.getTargetFeatures())
+    llvmFunc->addFnAttr("target-features", targetFeatures->getFeaturesString());
+
   if (auto attr = func.getVscaleRange())
     llvmFunc->addFnAttr(llvm::Attribute::getWithVScaleRangeArgs(
         getLLVMContext(), attr->getMinRange().getInt(),
diff --git a/mlir/test/Target/LLVMIR/Import/target-features.ll b/mlir/test/Target/LLVMIR/Import/target-features.ll
new file mode 100644
index 000000000000000..39e9a1204d3e022
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/Import/target-features.ll
@@ -0,0 +1,9 @@
+; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s
+
+; CHECK-LABEL: llvm.func @target_features()
+; CHECK-SAME: #llvm.target_features<"+sme,+sme-f64f64,+sve">
+define void @target_features() #0 {
+  ret void
+}
+
+attributes #0 = { "target-features"="+sme,+sme-f64f64,+sve" }
diff --git a/mlir/test/Target/LLVMIR/target-features.mlir b/mlir/test/Target/LLVMIR/target-features.mlir
new file mode 100644
index 000000000000000..02c07d27ca3cd84
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/target-features.mlir
@@ -0,0 +1,7 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: define void @target_features
+// CHECK: attributes #{{.*}} = { "target-features"="+sme,+sme-f64f64,+sve" }
+llvm.func @target_features() attributes { target_features = #llvm.target_features<"+sme,+sve,+sme-f64f64"> } {
+  llvm.return
+}

>From 9cab3352bd3e54a04af9976a6fe95daa6d6bf4fb Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Tue, 7 Nov 2023 12:21:26 +0000
Subject: [PATCH 2/2] Fixups

---
 .../mlir/Dialect/LLVMIR/LLVMAttrDefs.td       | 11 +++++++---
 mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp      | 20 ++++++++-----------
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index e0f4518abf3c01f..0246d795c0644ad 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -887,11 +887,11 @@ def LLVM_TargetFeaturesAttr : LLVM_Attr<"TargetFeatures", "target_features"> {
   }];
 
   let parameters = (ins
-    ArrayRefOfSelfAllocationParameter<"TargetFeature", "">: $features);
+    ArrayRefOfSelfAllocationParameter<"TargetFeature", "">:$features);
 
   let builders = [
-    TypeBuilder<(ins "::llvm::ArrayRef<TargetFeature>": $features)>,
-    TypeBuilder<(ins "StringRef": $features)>
+    TypeBuilder<(ins "::llvm::ArrayRef<TargetFeature>":$features)>,
+    TypeBuilder<(ins "::llvm::StringRef":$features)>
   ];
 
   let extraClassDeclaration = [{
@@ -907,6 +907,11 @@ def LLVM_TargetFeaturesAttr : LLVM_Attr<"TargetFeatures", "target_features"> {
     /// Finds the target features on the parent FunctionOpInterface.
     /// Note: This assumes the attribute is called "target_features".
     static TargetFeaturesAttr featuresAt(Operation* op);
+
+    /// Canonical name for this attribute within MLIR.
+    static constexpr StringLiteral attributeName() {
+      return StringLiteral("target_features");
+    }
   }];
 
   let hasCustomAssemblyFormat = 1;
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
index e97c30487210cd8..f510f949f9cfe4b 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp
@@ -144,8 +144,8 @@ TargetFeaturesAttr::get(MLIRContext *context,
 TargetFeaturesAttr TargetFeaturesAttr::get(MLIRContext *context,
                                            StringRef targetFeatures) {
   SmallVector<StringRef> features;
-  StringRef{targetFeatures}.split(features, ',', /*MaxSplit=*/-1,
-                                  /*KeepEmpty=*/false);
+  targetFeatures.split(features, ',', /*MaxSplit=*/-1,
+                       /*KeepEmpty=*/false);
   return get(context, llvm::map_to_vector(features, [](StringRef feature) {
                return TargetFeature{feature};
              }));
@@ -160,15 +160,11 @@ bool TargetFeaturesAttr::contains(TargetFeature feature) const {
 }
 
 std::string TargetFeaturesAttr::getFeaturesString() const {
-  std::string features;
-  bool first = true;
-  for (TargetFeature feature : getFeatures()) {
-    if (!first)
-      features += ",";
-    features += StringRef(feature);
-    first = false;
-  }
-  return features;
+  std::string featuresString;
+  llvm::raw_string_ostream ss(featuresString);
+  llvm::interleave(
+      getFeatures(), ss, [&](auto &feature) { ss << StringRef(feature); }, ",");
+  return ss.str();
 }
 
 TargetFeaturesAttr TargetFeaturesAttr::featuresAt(Operation *op) {
@@ -176,5 +172,5 @@ TargetFeaturesAttr TargetFeaturesAttr::featuresAt(Operation *op) {
   if (!parentFunction)
     return {};
   return parentFunction.getOperation()->getAttrOfType<TargetFeaturesAttr>(
-      "target_features");
+      attributeName());
 }



More information about the Mlir-commits mailing list