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

Author: MacDue

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

@@ -933,4 +933,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

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.

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

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

mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp

@@ -14,11 +14,14 @@
 #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 "llvm/IR/DebugInfoMetadata.h"
+#include <algorithm>
 #include <optional>
+#include <set>
 
 using namespace mlir;
 using namespace mlir::LLVM;
@@ -183,3 +186,69 @@ void printExpressionArg(AsmPrinter &printer, uint64_t opcode,
     i++;
   });
 }
+
+//===----------------------------------------------------------------------===//
+// 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");
+}

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -1627,6 +1627,7 @@ static constexpr std::array ExplicitAttributes{
     StringLiteral("aarch64_pstate_za_new"),
     StringLiteral("vscale_range"),
     StringLiteral("frame-pointer"),
+    StringLiteral("target-features"),
 };
 
 static void processPassthroughAttrs(llvm::Function *func, LLVMFuncOp funcOp) {
@@ -1717,6 +1718,12 @@ void ModuleImport::processFunctionAttributes(llvm::Function *func,
                                  stringRefFramePointerKind)
                                  .value()));
   }
+
+  if (llvm::Attribute attr = func->getFnAttribute("target-features");
+      attr.isStringAttribute()) {
+    funcOp.setTargetFeaturesAttr(
+        LLVM::TargetFeaturesAttr::get(context, attr.getValueAsString()));
+  }
 }
 
 DictionaryAttr

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -968,6 +968,9 @@ LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
   if (func.getArmNewZa())
     llvmFunc->addFnAttr("aarch64_pstate_za_new");
 
+  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(),

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

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