[MLIR][LLVMIR][DLTI] Pass for updating target features per TargetMachine

Author: rolfmorel

mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt

@@ -7,12 +7,26 @@ mlir_tablegen(LLVMOpsDialect.h.inc -gen-dialect-decls)
 mlir_tablegen(LLVMOpsDialect.cpp.inc -gen-dialect-defs)
 mlir_tablegen(LLVMOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(LLVMOpsEnums.cpp.inc -gen-enum-defs)
-mlir_tablegen(LLVMOpsAttrDefs.h.inc -gen-attrdef-decls
-              -attrdefs-dialect=llvm)
+#For LLVMOpsAttrDefs.h.inc, see below.
 mlir_tablegen(LLVMOpsAttrDefs.cpp.inc -gen-attrdef-defs
               -attrdefs-dialect=llvm)
 add_public_tablegen_target(MLIRLLVMOpsIncGen)
 
+# NB: Separate out LLVMOpsAttrDefs.h.inc generation as generating it
+#     through LLVMOps.td ends up defining LLVMTargetFeaturesAttr even
+#     though LLVMTargetFeaturesAttrDefs.* is responsible for that.
+set(LLVM_TARGET_DEFINITIONS LLVMAttrAndEnumDefs.td)
+mlir_tablegen(LLVMOpsAttrDefs.h.inc -gen-attrdef-decls -attrdefs-dialect=llvm)
+add_public_tablegen_target(MLIRLLVMAttrsIncGen)
+
+# NB: LLVMTargetFeaturesAttr is split out into its own file
+#     to break a recursive dependency: LLVMInterfaces depends
+#     on it, and other LLVMAttrs depending on LLVMInterfaces.
+set(LLVM_TARGET_DEFINITIONS LLVMTargetFeaturesAttrDefs.td)
+mlir_tablegen(LLVMTargetFeaturesAttrDefs.h.inc -gen-attrdef-decls)
+mlir_tablegen(LLVMTargetFeaturesAttrDefs.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(MLIRLLVMTargetFeaturesAttrsIncGen)
+
 set(LLVM_TARGET_DEFINITIONS LLVMTypes.td)
 mlir_tablegen(LLVMTypes.h.inc -gen-typedef-decls -typedefs-dialect=llvm)
 mlir_tablegen(LLVMTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=llvm)

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

@@ -1243,70 +1243,7 @@ def LLVM_VScaleRangeAttr : LLVM_Attr<"VScaleRange", "vscale_range"> {
 }
 
 //===----------------------------------------------------------------------===//
-// TargetFeaturesAttr
-//===----------------------------------------------------------------------===//
-
-def LLVM_TargetFeaturesAttr : LLVM_Attr<"TargetFeatures", "target_features">
-{
-  let summary = "LLVM target features attribute";
-
-  let description = [{
-    Represents the LLVM target features as a list that can be checked within
-    passes/rewrites.
-
-    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 OptionalArrayRefParameter<"StringAttr">:$features);
-
-  let builders = [
-    TypeBuilder<(ins "::llvm::StringRef":$features)>,
-    TypeBuilder<(ins "::llvm::ArrayRef<::llvm::StringRef>":$features)>
-  ];
-
-  let extraClassDeclaration = [{
-    /// Checks if a feature is contained within the features list.
-    /// Note: Using a StringAttr allows doing pointer-comparisons.
-    bool contains(::mlir::StringAttr feature) const;
-    bool contains(::llvm::StringRef feature) const;
-
-    bool nullOrEmpty() const {
-      // Checks if this attribute is null, or the features are empty.
-      return !bool(*this) || getFeatures().empty();
-    }
-
-    /// 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 name matches the return value of
-    /// `getAttributeName()`.
-    static TargetFeaturesAttr featuresAt(Operation* op);
-
-    /// Canonical name for this attribute within MLIR.
-    static constexpr StringLiteral getAttributeName() {
-      return StringLiteral("target_features");
-    }
-  }];
-
-  let assemblyFormat = "`<` `[` (`]`) : ($features^ `]`)? `>`";
-  let genVerifyDecl = 1;
-}
-
-//===----------------------------------------------------------------------===//
-// LLVM_TargetAttr
+// TargetAttr
 //===----------------------------------------------------------------------===//
 
 def LLVM_TargetAttr : LLVM_Attr<"Target", "target",
@@ -1324,7 +1261,7 @@ def LLVM_TargetAttr : LLVM_Attr<"Target", "target",
   }];
   let parameters = (ins "StringAttr":$triple,
                         "StringAttr":$chip,
-                        OptionalParameter<"StringAttr", "">:$features);
+                        OptionalParameter<"TargetFeaturesAttr", "">:$features);
 
   let assemblyFormat = [{`<` struct($triple, $chip, $features) `>`}];
 

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h

@@ -92,9 +92,14 @@ class TBAANodeAttr : public Attribute {
 using cconv::CConv;
 using linkage::Linkage;
 using tailcallkind::TailCallKind;
+
+class TargetFeaturesAttr;
 } // namespace LLVM
 } // namespace mlir
 
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/LLVMIR/LLVMTargetFeaturesAttrDefs.h.inc"
+
 #include "mlir/Dialect/LLVMIR/LLVMAttrInterfaces.h.inc"
 
 #define GET_ATTRDEF_CLASSES

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td

@@ -10,6 +10,7 @@
 #define LLVMIR_DIALECT
 
 include "mlir/IR/DialectBase.td"
+include "mlir/IR/AttrTypeBase.td"
 
 def LLVM_Dialect : Dialect {
   let name = "llvm";
@@ -127,4 +128,12 @@ def LLVM_Dialect : Dialect {
   }];
 }
 
+// All of the attributes will extend this class.
+class LLVM_Attr<string name, string attrMnemonic,
+                list<Trait> traits = [],
+                string baseCppClass = "::mlir::Attribute">
+    : AttrDef<LLVM_Dialect, name, traits, baseCppClass> {
+  let mnemonic = attrMnemonic;
+}
+
 #endif  // LLVMIR_DIALECT

mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td

@@ -561,7 +561,7 @@ def LLVM_TargetAttrInterface
     >,
     InterfaceMethod<
       /*description=*/"Returns the target features as a string.",
-      /*retTy=*/"StringAttr",
+      /*retTy=*/"LLVM::TargetFeaturesAttr",
       /*methodName=*/"getFeatures",
       /*args=*/(ins)
     >

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

@@ -14,6 +14,7 @@
 #define LLVMIR_OPS
 
 include "mlir/Dialect/LLVMIR/LLVMAttrDefs.td"
+include "mlir/Dialect/LLVMIR/LLVMTargetFeaturesAttrDefs.td"
 include "mlir/Dialect/LLVMIR/LLVMEnums.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/IR/EnumAttr.td"

mlir/include/mlir/Target/LLVMIR/Transforms/Passes.td

@@ -27,4 +27,18 @@ def LLVMTargetToDataLayout : Pass<"llvm-target-to-data-layout"> {
   ];
 }
 
+def LLVMFeaturesFromTarget : Pass<"llvm-features-from-target"> {
+  let summary = "TODO";
+  let dependentDialects = ["mlir::DLTIDialect"];
+  let description = [{
+    TODO
+  }];
+  let options = [
+    Option<"initializeLLVMTargets", "initialize-llvm-targets", "bool",
+           /*default=*/"true",
+           "Whether to pre-load all available target machines, that LLVM is "
+           "configured to support, into the TargetRegistry.">
+  ];
+}
+
 #endif // MLIR_TARGET_LLVMIR_TRANSFORMS_PASSES

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

@@ -403,6 +403,20 @@ ModuleFlagAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                      << key << "'";
 }
 
+FailureOr<Attribute> TargetFeaturesAttr::query(DataLayoutEntryKey key) {
+  if (auto stringKey = dyn_cast<StringAttr>(key)) {
+    if (contains(stringKey))
+      return UnitAttr::get(getContext());
+
+    if (contains((std::string("+") + stringKey.strref()).str()))
+      return BoolAttr::get(getContext(), true);
+
+    if (contains((std::string("-") + stringKey.strref()).str()))
+      return BoolAttr::get(getContext(), false);
+  }
+  return failure();
+}
+
 //===----------------------------------------------------------------------===//
 // LLVM_TargetAttr
 //===----------------------------------------------------------------------===//

mlir/lib/Target/LLVMIR/Transforms/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_mlir_dialect_library(MLIRTargetLLVMIRTransforms
   TargetToDataLayout.cpp
+  FeaturesFromTarget.cpp
 
   DEPENDS
   MLIRTargetLLVMIRTransformsIncGen

mlir/lib/Target/LLVMIR/Transforms/TargetToDataLayout.cpp

@@ -33,7 +33,8 @@ static FailureOr<std::unique_ptr<llvm::TargetMachine>>
 getTargetMachine(LLVM::TargetAttrInterface attr) {
   StringRef triple = attr.getTriple();
   StringRef chipAKAcpu = attr.getChip();
-  StringRef features = attr.getFeatures() ? attr.getFeatures().getValue() : "";
+  std::string features =
+      attr.getFeatures() ? attr.getFeatures().getFeaturesString() : "";
 
   std::string error;
   const llvm::Target *target =