[MLIR][LLVMIR][DLTI] Pass to update #llvm.target's features per relevant backend (#154938)

Modifies `#llvm.target<..., features = $FEATURES>` so that `$FEATURES`
is now an `#llvm.target_features<[...]>` attribute (rather than a
`StringAttr`). This enables the attribute to respond to DLTI queries for
the different target features.

The pass updates the `$FEATURES` attribute of the target attr at name
`llvm.target` in accordance with the (Sub)Target's features that the
relevant LLVM backend knows about.

---

DEMO:
```mlir
module attributes {llvm.target = #llvm.target<triple = "x86_64-unknown-linux",
                                              chip = "skylake"> } {
}
```
by way of `-llvm-target-to-target-features` turns into:
```mlir
module attributes {llvm.target = #llvm.target<triple = "x86_64-unknown-linux",
                                              chip = "skylake", 
                                              features = <["+64bit", "+64bit-mode", "+adx", "+aes", "+allow-light-256-bit", "+avx", "+avx2", "+bmi", "+bmi2", "+clflushopt", "+cmov", "+crc32", "+cx16", "+cx8", "+ermsb", "+f16c", "+false-deps-popcnt", "+fast-15bytenop", "+fast-gather", "+fast-scalar-fsqrt", "+fast-shld-rotate", "+fast-variable-crosslane-shuffle", "+fast-variable-perlane-shuffle", "+fast-vector-fsqrt", "+fma", "+fsgsbase", "+fxsr", "+idivq-to-divl", "+invpcid", "+lzcnt", "+macrofusion", "+mmx", "+movbe", "+no-bypass-delay-blend", "+no-bypass-delay-mov", "+no-bypass-delay-shuffle", "+nopl", "+pclmul", "+popcnt", "+prfchw", "+rdrnd", "+rdseed", "+sahf", "+slow-3ops-lea", "+sse", "+sse2", "+sse3", "+sse4.1", "+sse4.2", "+ssse3", "+vzeroupper", "+x87", "+xsave", "+xsavec", "+xsaveopt", "+xsaves"]>>} {
}
```

Author: rolfmorel

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

@@ -1246,7 +1246,8 @@ def LLVM_VScaleRangeAttr : LLVM_Attr<"VScaleRange", "vscale_range"> {
 // TargetFeaturesAttr
 //===----------------------------------------------------------------------===//
 
-def LLVM_TargetFeaturesAttr : LLVM_Attr<"TargetFeatures", "target_features">
+def LLVM_TargetFeaturesAttr : LLVM_Attr<"TargetFeatures", "target_features",
+                                        [DLTIQueryInterface]>
 {
   let summary = "LLVM target features attribute";
 
@@ -1299,14 +1300,17 @@ def LLVM_TargetFeaturesAttr : LLVM_Attr<"TargetFeatures", "target_features">
     static constexpr StringLiteral getAttributeName() {
       return StringLiteral("target_features");
     }
+
+    /// Returns the attribute associated with the key.
+    FailureOr<Attribute> query(DataLayoutEntryKey key);
   }];
 
   let assemblyFormat = "`<` `[` (`]`) : ($features^ `]`)? `>`";
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
-// LLVM_TargetAttr
+// TargetAttr
 //===----------------------------------------------------------------------===//
 
 def LLVM_TargetAttr : LLVM_Attr<"Target", "target",
@@ -1324,7 +1328,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/LLVMInterfaces.td

@@ -560,8 +560,10 @@ def LLVM_TargetAttrInterface
       /*args=*/(ins)
     >,
     InterfaceMethod<
-      /*description=*/"Returns the target features as a string.",
-      /*retTy=*/"StringAttr",
+      /*description=*/"Returns the target features as a TargetFeaturesAttr.",
+      /*retTy=*/"Attribute", // NB: will be a LLVM::TargetFeaturesAttr, though
+                             // need to work around a cyclic dependency on
+                             // LLVMInterfaces.td and LLVMAttrDefs.td.
       /*methodName=*/"getFeatures",
       /*args=*/(ins)
     >

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

@@ -12,7 +12,6 @@
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
-
 namespace LLVM {
 
 #define GEN_PASS_DECL

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

@@ -27,4 +27,20 @@ def LLVMTargetToDataLayout : Pass<"llvm-target-to-data-layout"> {
   ];
 }
 
+def LLVMTargetToTargetFeatures : Pass<"llvm-target-to-target-features"> {
+  let summary = "Update attached #llvm.target's features per the described target";
+  let description = [{
+    Obtain the TargetMachine specified by the attached #llvm.target's attributes
+    and obtain from it the full list of features of the selected target. Updates
+    the attached #llvm.target so that its features reflect the full list of
+    features.
+  }];
+  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/include/mlir/Target/LLVMIR/Transforms/TargetUtils.h

@@ -0,0 +1,35 @@
+//===- TargetUtils.h - Utils to obtain LLVM's TargetMachine and DataLayout ===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TARGET_LLVMIR_TRANSFORMS_TARGETUTILS_H
+#define MLIR_TARGET_LLVMIR_TRANSFORMS_TARGETUTILS_H
+
+#include "mlir/Dialect/LLVMIR/LLVMInterfaces.h"
+#include "llvm/Support/Threading.h"
+#include "llvm/Target/TargetMachine.h"
+
+namespace mlir {
+namespace LLVM {
+namespace detail {
+/// Idempotent helper to register/initialize all backends that LLVM has been
+/// configured to support. Only runs the first time it is called.
+void initializeBackendsOnce();
+
+/// Helper to obtain the TargetMachine specified by the properties of the
+/// TargetAttrInterface-implementing attribute.
+FailureOr<std::unique_ptr<llvm::TargetMachine>>
+getTargetMachine(mlir::LLVM::TargetAttrInterface attr);
+
+/// Helper to obtain the DataLayout of the target specified by the properties of
+/// the TargetAttrInterface-implementing attribute.
+FailureOr<llvm::DataLayout> getDataLayout(mlir::LLVM::TargetAttrInterface attr);
+} // namespace detail
+} // namespace LLVM
+} // namespace mlir
+
+#endif // MLIR_TARGET_LLVMIR_TRANSFORMS_TARGETUTILS_H

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

@@ -374,6 +374,43 @@ TargetFeaturesAttr TargetFeaturesAttr::featuresAt(Operation *op) {
       getAttributeName());
 }
 
+FailureOr<Attribute> TargetFeaturesAttr::query(DataLayoutEntryKey key) {
+  auto stringKey = dyn_cast<StringAttr>(key);
+  if (!stringKey)
+    return failure();
+
+  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();
+}
+
+//===----------------------------------------------------------------------===//
+// TargetAttr
+//===----------------------------------------------------------------------===//
+
+FailureOr<::mlir::Attribute> TargetAttr::query(DataLayoutEntryKey key) {
+  if (auto stringAttrKey = dyn_cast<StringAttr>(key)) {
+    if (stringAttrKey.getValue() == "triple")
+      return getTriple();
+    if (stringAttrKey.getValue() == "chip")
+      return getChip();
+    if (stringAttrKey.getValue() == "features" && getFeatures())
+      return getFeatures();
+  }
+  return failure();
+}
+
+//===----------------------------------------------------------------------===//
+// ModuleFlagAttr
+//===----------------------------------------------------------------------===//
+
 LogicalResult
 ModuleFlagAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                        LLVM::ModFlagBehavior flagBehavior, StringAttr key,
@@ -402,20 +439,3 @@ ModuleFlagAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                         "supported for unknown key '"
                      << key << "'";
 }
-
-//===----------------------------------------------------------------------===//
-// LLVM_TargetAttr
-//===----------------------------------------------------------------------===//
-
-FailureOr<::mlir::Attribute> TargetAttr::query(DataLayoutEntryKey key) {
-  if (auto stringAttrKey = dyn_cast<StringAttr>(key)) {
-    if (stringAttrKey.getValue() == "triple")
-      return getTriple();
-    if (stringAttrKey.getValue() == "chip")
-      return getChip();
-    if (stringAttrKey.getValue() == "features" && getFeatures())
-      return getFeatures();
-  }
-
-  return failure();
-}

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

@@ -1,5 +1,7 @@
 add_mlir_dialect_library(MLIRTargetLLVMIRTransforms
   TargetToDataLayout.cpp
+  TargetToTargetFeatures.cpp
+  TargetUtils.cpp
 
   DEPENDS
   MLIRTargetLLVMIRTransformsIncGen

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

@@ -5,21 +5,14 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+
 #include "mlir/Target/LLVMIR/Transforms/Passes.h"
+#include "mlir/Target/LLVMIR/Transforms/TargetUtils.h"
 
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Target/LLVMIR/Import.h"
 
-#include "llvm/MC/TargetRegistry.h"
-#include "llvm/Support/Debug.h"
-#include "llvm/Support/TargetSelect.h"
-#include "llvm/Target/TargetMachine.h"
-
-#define DEBUG_TYPE "mlir-llvm-target-to-data-layout"
-#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
-#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
-
 namespace mlir {
 namespace LLVM {
 #define GEN_PASS_DEF_LLVMTARGETTODATALAYOUT
@@ -29,35 +22,6 @@ namespace LLVM {
 
 using namespace mlir;
 
-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 error;
-  const llvm::Target *target =
-      llvm::TargetRegistry::lookupTarget(triple, error);
-  if (!target || !error.empty()) {
-    LDBG("Looking up target '" << triple << "' failed: " << error << "\n");
-    return failure();
-  }
-
-  return std::unique_ptr<llvm::TargetMachine>(target->createTargetMachine(
-      llvm::Triple(triple), chipAKAcpu, features, {}, {}));
-}
-
-static FailureOr<llvm::DataLayout>
-getDataLayout(LLVM::TargetAttrInterface attr) {
-  FailureOr<std::unique_ptr<llvm::TargetMachine>> targetMachine =
-      getTargetMachine(attr);
-  if (failed(targetMachine)) {
-    LDBG("Failed to retrieve the target machine for data layout.\n");
-    return failure();
-  }
-  return (targetMachine.value())->createDataLayout();
-}
-
 struct TargetToDataLayoutPass
     : public LLVM::impl::LLVMTargetToDataLayoutBase<TargetToDataLayoutPass> {
   using LLVM::impl::LLVMTargetToDataLayoutBase<
@@ -66,15 +30,8 @@ struct TargetToDataLayoutPass
   void runOnOperation() override {
     Operation *op = getOperation();
 
-    if (initializeLLVMTargets) {
-      static llvm::once_flag initializeBackendsOnce;
-      llvm::call_once(initializeBackendsOnce, []() {
-        // Ensure that the targets, that LLVM has been configured to support,
-        // are loaded into the TargetRegistry.
-        llvm::InitializeAllTargets();
-        llvm::InitializeAllTargetMCs();
-      });
-    }
+    if (initializeLLVMTargets)
+      LLVM::detail::initializeBackendsOnce();
 
     auto targetAttr = op->getAttrOfType<LLVM::TargetAttrInterface>(
         LLVM::LLVMDialect::getTargetAttrName());
@@ -85,7 +42,8 @@ struct TargetToDataLayoutPass
       return signalPassFailure();
     }
 
-    FailureOr<llvm::DataLayout> dataLayout = getDataLayout(targetAttr);
+    FailureOr<llvm::DataLayout> dataLayout =
+        LLVM::detail::getDataLayout(targetAttr);
     if (failed(dataLayout)) {
       op->emitError() << "failed to obtain llvm::DataLayout for " << targetAttr;
       return signalPassFailure();

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

@@ -0,0 +1,78 @@
+//===- TargetToTargetFeatures.cpp - extract features from TargetMachine ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Target/LLVMIR/Transforms/Passes.h"
+#include "mlir/Target/LLVMIR/Transforms/TargetUtils.h"
+
+#include "mlir/Dialect/DLTI/DLTI.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Target/LLVMIR/Import.h"
+
+#include "llvm/MC/MCSubtargetInfo.h"
+
+namespace mlir {
+namespace LLVM {
+#define GEN_PASS_DEF_LLVMTARGETTOTARGETFEATURES
+#include "mlir/Target/LLVMIR/Transforms/Passes.h.inc"
+} // namespace LLVM
+} // namespace mlir
+
+using namespace mlir;
+
+struct TargetToTargetFeaturesPass
+    : public LLVM::impl::LLVMTargetToTargetFeaturesBase<
+          TargetToTargetFeaturesPass> {
+  using LLVM::impl::LLVMTargetToTargetFeaturesBase<
+      TargetToTargetFeaturesPass>::LLVMTargetToTargetFeaturesBase;
+
+  void runOnOperation() override {
+    Operation *op = getOperation();
+
+    if (initializeLLVMTargets)
+      LLVM::detail::initializeBackendsOnce();
+
+    auto targetAttr = op->getAttrOfType<LLVM::TargetAttr>(
+        LLVM::LLVMDialect::getTargetAttrName());
+    if (!targetAttr) {
+      op->emitError() << "no LLVM::TargetAttr attribute at key \""
+                      << LLVM::LLVMDialect::getTargetAttrName() << "\"";
+      return signalPassFailure();
+    }
+
+    FailureOr<std::unique_ptr<llvm::TargetMachine>> targetMachine =
+        LLVM::detail::getTargetMachine(targetAttr);
+    if (failed(targetMachine)) {
+      op->emitError() << "failed to obtain llvm::TargetMachine for "
+                      << targetAttr;
+      return signalPassFailure();
+    }
+
+    llvm::MCSubtargetInfo const *subTargetInfo =
+        (*targetMachine)->getMCSubtargetInfo();
+
+    const std::vector<llvm::SubtargetFeatureKV> enabledFeatures =
+        subTargetInfo->getEnabledProcessorFeatures();
+
+    auto plussedFeatures = llvm::to_vector(
+        llvm::map_range(enabledFeatures, [](llvm::SubtargetFeatureKV feature) {
+          return std::string("+") + feature.Key;
+        }));
+
+    auto plussedFeaturesRefs = llvm::to_vector(llvm::map_range(
+        plussedFeatures, [](auto &it) { return StringRef(it.c_str()); }));
+
+    auto fullTargetFeaturesAttr =
+        LLVM::TargetFeaturesAttr::get(&getContext(), plussedFeaturesRefs);
+
+    auto updatedTargetAttr =
+        LLVM::TargetAttr::get(&getContext(), targetAttr.getTriple(),
+                              targetAttr.getChip(), fullTargetFeaturesAttr);
+
+    op->setAttr(LLVM::LLVMDialect::getTargetAttrName(), updatedTargetAttr);
+  }
+};