[mlir][llvm] Add llvm.target_features features attribute #71510
Conversation
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-llvm Author: Benjamin Maxwell (MacDue) ChangesThis 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: 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. // 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. Full diff: https://github.com/llvm/llvm-project/pull/71510.diff 8 Files Affected:
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
+}
|
There was a problem hiding this comment.
All this string manipulation and sorting adds seemingly unnecessary complexity. Have you considered the attribute containing an array of StringAttr instead? Common target features can have their StringAttr preallocated in the dialect object, so the search would be just comparing a bunch of pointers. Less common features can still rely on string comparison. It will also give this a more MLIResque syntax.
|
|
||
| let builders = [ | ||
| TypeBuilder<(ins "::llvm::ArrayRef<TargetFeature>": $features)>, | ||
| TypeBuilder<(ins "StringRef": $features)> |
There was a problem hiding this comment.
| TypeBuilder<(ins "StringRef": $features)> | |
| TypeBuilder<(ins "::llvm::StringRef": $features)> |
| if (!bool(*this)) | ||
| return false; // Allows checking null target features. |
There was a problem hiding this comment.
It's unclear to me how this works, looks like dereferencing a null pointer.
There was a problem hiding this comment.
this is never null, it's the attribute impl that can be null. So this is calling the explicit bool() operator on the attribute to see if it has an impl:
| if (!parentFunction) | ||
| return {}; | ||
| return parentFunction.getOperation()->getAttrOfType<TargetFeaturesAttr>( | ||
| "target_features"); |
There was a problem hiding this comment.
Please factor out the magic string into a named StringLiteral in the dialect.
The LLVM attribute is just a string, so to convert between the representations will still require the string manipulation. I can experiment with other representations, but I think having the custom attribute leads to a little nicer API as the related utility functions can be attached to the attribute type. |
LLVM IR attributes are notoriously unstructured, we are trying to fix that a bit in the dialect. The string manipulation can be restricted to translation where it is unavoidable, but MLIR-only manipulation can more structured. Note that I'm not proposing to use an ArrayAttr containing StringAttr, but a custom attribute parameterized by an array of StringAttr. |
MacDue
force-pushed
the
llvm_target_features_attr
branch
from
9cab335 to
4a121a0
Compare
|
@ftynse , thanks for reviewing this 🙏🏻 . We chatted about this offline and realised that instead of using target attributes at this level, we might be able to control what lowerings to use higher up the stack (in our case that would be in IREE). That should be sufficient for SME's outer-product instructions that this was originally intended for - we will basically try to configure tiling/vectorisation in IREE depending on the available target extensions/attributes. I still believe that we will need this for things like My suggestion would be to label this as "in progress/draft" and park it until we:
That should help us to identify a suitable motivating example. WDYT? |
|
Sounds good to me, I'll move this to draft for now :) |
|
Depends on your time availability. I'm almost always happy to help land missing features in MLIR, but it's probably better if these features solve a problem for you. Though I'm still happy to land them even if they don't, just for the sake of you (and others) not being blocked by their absence when the need arises. |
MacDue
force-pushed
the
llvm_target_features_attr
branch
from
1c5e31f to
17b3017
Compare
Thank you! Makes sense to me, I've resolved the conflicts and rebased now (sorry it took a little while). |
| if (llvm::Attribute attr = func->getFnAttribute("target-features"); | ||
| attr.isStringAttribute()) { | ||
| funcOp.setTargetFeaturesAttr( | ||
| LLVM::TargetFeaturesAttr::get(context, attr.getValueAsString())); |
There was a problem hiding this comment.
nit: personally would find it clearer if attr.isStringAttribute() was in a separate if, threw me off a little.
| if (llvm::Attribute attr = func->getFnAttribute("target-features"); | |
| attr.isStringAttribute()) { | |
| funcOp.setTargetFeaturesAttr( | |
| LLVM::TargetFeaturesAttr::get(context, attr.getValueAsString())); | |
| if (llvm::Attribute attr = func->getFnAttribute("target-features")) { | |
| if (attr.isStringAttribute()) | |
| funcOp.setTargetFeaturesAttr( | |
| LLVM::TargetFeaturesAttr::get(context, attr.getValueAsString())); | |
| } |
There was a problem hiding this comment.
I think it's LLVM style to avoid extra nesting, so I'm going keep this :)
MacDue
force-pushed
the
llvm_target_features_attr
branch
2 times, most recently
from
1956a90 to
6fdaa20
Compare
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.
MacDue
force-pushed
the
llvm_target_features_attr
branch
from
6fdaa20 to
bea4aaf
Compare
|
@MacDue Thank you for your patch. I would like you to ask if you are planning to add target-cpu attribute to LLVM dialect as well? |
I've not got any plans to; feel free to post a patch for that yourself :) |
This patch adds a target_features (TargetFeaturesAttr) to the LLVM dialect to allow setting and querying the features in use on a function.
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:
Then within a later rewrite the attribute can be checked, and used to make lowering decisions.
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.