[mlir][interfaces] Add the TargetInfo attribute interface

mlir/include/mlir/Dialect/DLTI/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_subdirectory(TransformOps)
+add_subdirectory(Transforms)
 
 add_mlir_dialect(DLTI dlti)
 add_mlir_doc(DLTIAttrs DLTIDialect Dialects/ -gen-dialect-doc)

mlir/include/mlir/Dialect/DLTI/DLTIBase.td

@@ -39,6 +39,10 @@ def DLTI_Dialect : Dialect {
     constexpr const static ::llvm::StringLiteral
     kTargetDeviceDescAttrName = "dlti.target_device_spec";
 
+    // Top-level attribute name for target information.
+    constexpr const static ::llvm::StringLiteral
+    kTargetDescAttrName = "dlti.target";
+
     // Constants used in entries.
     constexpr const static ::llvm::StringLiteral
     kDataLayoutEndiannessKey = "dlti.endianness";

mlir/include/mlir/Dialect/DLTI/Traits.h

@@ -19,6 +19,10 @@ namespace impl {
 LogicalResult verifyHasDefaultDLTIDataLayoutTrait(Operation *op);
 DataLayoutSpecInterface getDataLayoutSpec(Operation *op);
 TargetSystemSpecInterface getTargetSystemSpec(Operation *op);
+TargetAttrInterface getTargetAttr(Operation *op);
+void setDataLayoutSpec(Operation *op, DataLayoutSpecInterface spec);
+void setTargetSystemSpec(Operation *op, TargetSystemSpecInterface spec);
+void setTargetAttr(Operation *op, TargetAttrInterface target);
 } // namespace impl
 
 /// Trait to be used by operations willing to use the implementation of the
@@ -39,11 +43,27 @@ class HasDefaultDLTIDataLayout
     return impl::getDataLayoutSpec(this->getOperation());
   }
 
+  /// Sets the data layout specification.
+  void setDataLayoutSpec(DataLayoutSpecInterface spec) {
+    impl::setDataLayoutSpec(this->getOperation(), spec);
+  }
   /// Returns the target system description specification as provided by DLTI
   /// dialect
   TargetSystemSpecInterface getTargetSystemSpec() {
     return impl::getTargetSystemSpec(this->getOperation());
   }
+  /// Sets the target system description specification.
+  void setTargetSystemSpec(TargetSystemSpecInterface spec) {
+    impl::setTargetSystemSpec(this->getOperation(), spec);
+  }
+  /// Returns the target information as provided by DLTI dialect.
+  TargetAttrInterface getTargetAttr() {
+    return impl::getTargetAttr(this->getOperation());
+  }
+  /// Sets the target information.
+  void setTargetAttr(TargetAttrInterface target) {
+    impl::setTargetAttr(this->getOperation(), target);
+  }
 };
 } // namespace mlir
 

mlir/include/mlir/Dialect/DLTI/Transforms/CMakeLists.txt

@@ -0,0 +1,7 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name DLTI)
+mlir_tablegen(Passes.capi.h.inc -gen-pass-capi-header --prefix DLTI)
+mlir_tablegen(Passes.capi.cpp.inc -gen-pass-capi-impl --prefix DLTI)
+add_public_tablegen_target(MLIRDLTIPassIncGen)
+
+add_mlir_doc(Passes DLTIPasses ./ -gen-pass-doc)

mlir/include/mlir/Dialect/DLTI/Transforms/Passes.h

@@ -0,0 +1,30 @@
+//===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines prototypes that expose pass constructors.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_DLTI_TRANSFORMS_PASSES_H
+#define MLIR_DIALECT_DLTI_TRANSFORMS_PASSES_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+#define GEN_PASS_DECL
+#include "mlir/Dialect/DLTI/Transforms/Passes.h.inc"
+
+/// Generate the code for registering passes.
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/DLTI/Transforms/Passes.h.inc"
+
+/// Sets the target specs using the target attached to the module.
+LogicalResult setTargetSpecsFromTarget(Operation *op);
+} // namespace mlir
+
+#endif // MLIR_DIALECT_DLTI_TRANSFORMS_PASSES_H

mlir/include/mlir/Dialect/DLTI/Transforms/Passes.td

@@ -0,0 +1,40 @@
+//===-- Passes.td - DLTI pass definition file --------------*- tablegen -*-===//
+//
+// 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_DIALECT_DLTI_PASSES
+#define MLIR_DIALECT_DLTI_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def DltiSetTargetSpecsFromTarget: Pass<"dlti-set-target-specs", ""> {
+  let summary = "Sets DLTI target specs using a target.";
+  let description = [{
+    This pass potentially sets the following DLTI target specs in the current
+    operation:
+    - The data layout.
+    - The target system spec.
+
+    Example:
+
+    ```mlir
+    // Given the following input:
+    builtin.module @module_1 attributes {dlti.target = #my.target} {...}
+    // After applying the pass:
+    builtin.module @module_1 attributes {
+      dlti.target = #my.target,
+      dlti.target_system_spec = #my.system_spec,
+      dlti.dl_spec = #my.dl_spec
+    } {...}
+    ```
+  }];
+  let dependentDialects = [
+    "::mlir::DLTIDialect"
+  ];
+}
+
+#endif // MLIR_DIALECT_DLTI_PASSES

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

@@ -1388,7 +1388,8 @@ def GPU_BarrierOp : GPU_Op<"barrier"> {
 }
 
 def GPU_GPUModuleOp : GPU_Op<"module", [
-      IsolatedFromAbove, DataLayoutOpInterface, HasDefaultDLTIDataLayout,
+      IsolatedFromAbove,
+      DeclareOpInterfaceMethods<DataLayoutOpInterface, ["getTargetAttr", "setTargetAttr"]>,
       NoRegionArguments, SymbolTable, Symbol] # GraphRegionNoTerminator.traits> {
   let summary = "A top level compilation unit containing code to be run on a GPU.";
   let description = [{

mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h

@@ -19,6 +19,7 @@
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Target/LLVMIR/ModuleTranslation.h"

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

@@ -17,6 +17,7 @@ include "mlir/IR/EnumAttr.td"
 include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/DataLayoutInterfaces.td"
 include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td"
 include "mlir/Interfaces/InferIntRangeInterface.td"
 include "mlir/Dialect/LLVMIR/LLVMTypes.td"

mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h

@@ -26,6 +26,7 @@
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 ///// Ops /////

mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

@@ -15,6 +15,7 @@
 
 include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
+include "mlir/Interfaces/DataLayoutInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 //===----------------------------------------------------------------------===//
@@ -1116,8 +1117,7 @@ def ROCDL_CvtSrFp8F32Op :
 // ROCDL target attribute.
 //===----------------------------------------------------------------------===//
 
-def ROCDL_TargetAttr :
-    ROCDL_Attr<"ROCDLTarget", "target"> {
+def ROCDL_TargettAttr : ROCDL_Attr<"ROCDLTarget", "target"> {
   let description = [{
     ROCDL target attribute for controlling compilation of AMDGPU targets. All
     parameters decay into default values if not present.

mlir/include/mlir/InitAllPasses.h

@@ -23,6 +23,7 @@
 #include "mlir/Dialect/Async/Passes.h"
 #include "mlir/Dialect/Bufferization/Pipelines/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/DLTI/Transforms/Passes.h"
 #include "mlir/Dialect/EmitC/Transforms/Passes.h"
 #include "mlir/Dialect/Func/Transforms/Passes.h"
 #include "mlir/Dialect/GPU/Pipelines/Passes.h"
@@ -75,6 +76,7 @@ inline void registerAllPasses() {
   bufferization::registerBufferizationPasses();
   func::registerFuncPasses();
   registerGPUPasses();
+  registerDLTIPasses();
   registerLinalgPasses();
   registerNVGPUPasses();
   registerSparseTensorPasses();

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

@@ -27,6 +27,7 @@ class DataLayout;
 class DataLayoutEntryInterface;
 class DLTIQueryInterface;
 class TargetDeviceSpecInterface;
+struct TargetSpec;
 class TargetSystemSpecInterface;
 using DataLayoutEntryKey = llvm::PointerUnion<Type, StringAttr>;
 // Using explicit SmallVector size because we cannot infer the size from the
@@ -305,6 +306,13 @@ class DataLayout {
   mutable std::optional<uint64_t> stackAlignment;
 };
 
+/// Helper struct for storing a target specification.
+struct TargetSpec {
+  /// Target system spec.
+  TargetSystemSpecInterface systemSpec;
+  /// Target data layout.
+  DataLayoutSpecInterface dataLayout;
+};
 } // namespace mlir
 
 #endif // MLIR_INTERFACES_DATALAYOUTINTERFACES_H

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

@@ -354,6 +354,46 @@ def TargetSystemSpecInterface : AttrInterface<"TargetSystemSpecInterface", [DLTI
   }];
 }
 
+def TargetAttrInterface : AttrInterface<"TargetAttrInterface"> {
+  let cppNamespace = "::mlir";
+
+  let description = [{
+    Attribute interface describing target information.
+
+    Target information attributes provide essential information on a
+    compilation target. This information includes the target triple identifier,
+    the target chip identifier, a string representation of the target features,
+    and the target spec data layout.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      /*description=*/"Returns the target triple identifier.",
+      /*retTy=*/"::llvm::StringRef",
+      /*methodName=*/"getTargetTriple",
+      /*args=*/(ins)
+    >,
+    InterfaceMethod<
+      /*description=*/"Returns the target chip identifier.",
+      /*retTy=*/"::llvm::StringRef",
+      /*methodName=*/"getTargetChip",
+      /*args=*/(ins)
+    >,
+    InterfaceMethod<
+      /*description=*/"Returns the target features as a string.",
+      /*retTy=*/"std::string",
+      /*methodName=*/"getTargetFeatures",
+      /*args=*/(ins)
+    >,
+    InterfaceMethod<
+      /*description=*/"Sets the target spec. Returns failure if there was a problem.",
+      /*retTy=*/"::llvm::LogicalResult",
+      /*methodName=*/"setTargetSpec",
+      /*args=*/(ins "::mlir::TargetSpec&":$spec)
+    >
+  ];
+}
+
 //===----------------------------------------------------------------------===//
 // Operation interface
 //===----------------------------------------------------------------------===//
@@ -387,13 +427,39 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
       /*methodName=*/"getDataLayoutSpec",
       /*args=*/(ins)
     >,
+    InterfaceMethod<
+      /*description=*/"Sets the data layout specification for this op.",
+      /*retTy=*/"void",
+      /*methodName=*/"setDataLayoutSpec",
+      /*args=*/(ins "::mlir::DataLayoutSpecInterface":$spec)
+    >,
     InterfaceMethod<
       /*description=*/"Returns the target system desc specification for this "
                       "op, or null if it does not exist.",
       /*retTy=*/"::mlir::TargetSystemSpecInterface",
       /*methodName=*/"getTargetSystemSpec",
       /*args=*/(ins)
     >,
+    InterfaceMethod<
+      /*description=*/"Sets the target system desc specification for this "
+                      "op.",
+      /*retTy=*/"void",
+      /*methodName=*/"setTargetSystemSpec",
+      /*args=*/(ins "::mlir::TargetSystemSpecInterface":$spec)
+    >,
+    InterfaceMethod<
+      /*description=*/"Returns the target attr for this op, or null if it "
+                      "does not exist.",
+      /*retTy=*/"::mlir::TargetAttrInterface",
+      /*methodName=*/"getTargetAttr",
+      /*args=*/(ins)
+    >,
+    InterfaceMethod<
+      /*description=*/"Sets the target attr for this",
+      /*retTy=*/"void",
+      /*methodName=*/"setTargetAttr",
+      /*args=*/(ins "::mlir::TargetAttrInterface":$target)
+    >,
     StaticInterfaceMethod<
       /*description=*/"Returns the size of the given type computed using the "
                       "relevant entries. The data layout object can be used "

mlir/include/mlir/Target/LLVM/Target.h

@@ -0,0 +1,84 @@
+//===- Target.h - Target information ----------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declare utilities to interact with LLVM targets by querying an MLIR
+// target.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TARGET_LLVM_TARGET_H
+#define MLIR_TARGET_LLVM_TARGET_H
+
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/TargetParser/Triple.h"
+
+namespace llvm {
+class Triple;
+class Target;
+class TargetMachine;
+} // namespace llvm
+
+namespace mlir {
+/// Given a target triple. chip and features returns the LLVM data layout.
+FailureOr<const llvm::DataLayout>
+getLLVMDataLayout(StringRef triple, StringRef chip, StringRef features);
+
+/// Returns the LLVM target triple held by `target`.
+llvm::Triple getTargetTriple(TargetAttrInterface target);
+
+/// Returns the LLVM target held by `target`.
+FailureOr<const llvm::Target *> getLLVMTarget(TargetAttrInterface target);
+
+/// Helper class for holding LLVM target information. Note: This class requires
+/// that the corresponding LLVM target has ben initialized.
+class TargetInfo {
+public:
+  TargetInfo(TargetInfo &&) = default;
+  TargetInfo(const TargetInfo &) = delete;
+  ~TargetInfo();
+  TargetInfo &operator=(TargetInfo &&) = default;
+  TargetInfo &operator=(const TargetInfo &) = delete;
+  /// Constructs the target info from `target`.
+  static FailureOr<TargetInfo> getTargetInfo(StringRef triple, StringRef chip,
+                                             StringRef features);
+
+  /// Constructs the target info from `target`.
+  static FailureOr<TargetInfo> getTargetInfo(TargetAttrInterface target) {
+    return getTargetInfo(target.getTargetTriple(), target.getTargetChip(),
+                         target.getTargetFeatures());
+  }
+
+  /// Returns the target chip.
+  StringRef getTargetChip() const;
+
+  /// Returns the target features.
+  StringRef getTargetFeatures() const;
+
+  /// Returns the target triple.
+  const llvm::Triple &getTriple() const;
+
+  /// Returns the target.
+  const llvm::Target &getTarget() const;
+
+  /// Returns the target machine.
+  const llvm::TargetMachine *getTargetMachine() const {
+    return targetMachine.get();
+  }
+
+  /// Returns the LLVM data layout for the corresponding target.
+  const llvm::DataLayout getDataLayout() const;
+
+private:
+  TargetInfo(llvm::TargetMachine *targetMachine);
+  /// The LLVM target machine.
+  mutable std::unique_ptr<llvm::TargetMachine> targetMachine;
+};
+} // namespace mlir
+
+#endif // MLIR_TARGET_LLVM_TARGET_H