diff --git a/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt b/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt index 4f8382e8e6e6b..816352fa6ce45 100644 --- a/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt +++ b/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) diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td index 1a08bafda54ee..b4d27f74aa257 100644 --- a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td +++ b/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"; diff --git a/mlir/include/mlir/Dialect/DLTI/Traits.h b/mlir/include/mlir/Dialect/DLTI/Traits.h index edfbdffbd1ba1..c28a904823778 100644 --- a/mlir/include/mlir/Dialect/DLTI/Traits.h +++ b/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 diff --git a/mlir/include/mlir/Dialect/DLTI/Transforms/CMakeLists.txt b/mlir/include/mlir/Dialect/DLTI/Transforms/CMakeLists.txt new file mode 100644 index 0000000000000..5ee4926345f74 --- /dev/null +++ b/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) diff --git a/mlir/include/mlir/Dialect/DLTI/Transforms/Passes.h b/mlir/include/mlir/Dialect/DLTI/Transforms/Passes.h new file mode 100644 index 0000000000000..9a0a536e3fac9 --- /dev/null +++ b/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 diff --git a/mlir/include/mlir/Dialect/DLTI/Transforms/Passes.td b/mlir/include/mlir/Dialect/DLTI/Transforms/Passes.td new file mode 100644 index 0000000000000..63c284d7a55d6 --- /dev/null +++ b/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 diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td index 68095b7bf5c59..1d20c7d2c7351 100644 --- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td +++ b/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, NoRegionArguments, SymbolTable, Symbol] # GraphRegionNoTerminator.traits> { let summary = "A top level compilation unit containing code to be run on a GPU."; let description = [{ diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h b/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h index a9270c6f52344..8bd4ab16c593a 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h +++ b/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" diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td index 0a6e66919f021..42b28fb574789 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td +++ b/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" diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h b/mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h index c2a82ffc1c43c..fa1131a463e1a 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h +++ b/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 ///// diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td index 186a4f53f93cb..80b5d89325ab5 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td +++ b/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. diff --git a/mlir/include/mlir/InitAllPasses.h b/mlir/include/mlir/InitAllPasses.h index dd8b292a87344..08f5563ff722f 100644 --- a/mlir/include/mlir/InitAllPasses.h +++ b/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(); diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h index ff40bfc4bee41..74d36e93cbab6 100644 --- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h +++ b/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; // Using explicit SmallVector size because we cannot infer the size from the @@ -305,6 +306,13 @@ class DataLayout { mutable std::optional 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 diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td index c5973d4252b0a..76f5d0e650bed 100644 --- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td +++ b/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,6 +427,12 @@ 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.", @@ -394,6 +440,26 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> { /*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 " diff --git a/mlir/include/mlir/Target/LLVM/Target.h b/mlir/include/mlir/Target/LLVM/Target.h new file mode 100644 index 0000000000000..3b6e9a601d25d --- /dev/null +++ b/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 +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 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 getTargetInfo(StringRef triple, StringRef chip, + StringRef features); + + /// Constructs the target info from `target`. + static FailureOr 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 targetMachine; +}; +} // namespace mlir + +#endif // MLIR_TARGET_LLVM_TARGET_H diff --git a/mlir/lib/Dialect/DLTI/CMakeLists.txt b/mlir/lib/Dialect/DLTI/CMakeLists.txt index 7691a8a10a37a..d0f49ed9b6f10 100644 --- a/mlir/lib/Dialect/DLTI/CMakeLists.txt +++ b/mlir/lib/Dialect/DLTI/CMakeLists.txt @@ -1,3 +1,4 @@ +add_subdirectory(Transforms) add_subdirectory(TransformOps) add_mlir_dialect_library(MLIRDLTIDialect DLTI.cpp diff --git a/mlir/lib/Dialect/DLTI/Traits.cpp b/mlir/lib/Dialect/DLTI/Traits.cpp index 34f2dd5896083..1c99ef3a7c72f 100644 --- a/mlir/lib/Dialect/DLTI/Traits.cpp +++ b/mlir/lib/Dialect/DLTI/Traits.cpp @@ -28,7 +28,26 @@ DataLayoutSpecInterface mlir::impl::getDataLayoutSpec(Operation *op) { DLTIDialect::kDataLayoutAttrName); } +void mlir::impl::setDataLayoutSpec(Operation *op, + DataLayoutSpecInterface spec) { + return op->setAttr(DLTIDialect::kDataLayoutAttrName, spec); +} + TargetSystemSpecInterface mlir::impl::getTargetSystemSpec(Operation *op) { return op->getAttrOfType( DLTIDialect::kTargetSystemDescAttrName); } + +void mlir::impl::setTargetSystemSpec(Operation *op, + TargetSystemSpecInterface spec) { + return op->setAttr(DLTIDialect::kTargetSystemDescAttrName, spec); +} + +TargetAttrInterface mlir::impl::getTargetAttr(Operation *op) { + return op->getAttrOfType( + DLTIDialect::kTargetDescAttrName); +} + +void mlir::impl::setTargetAttr(Operation *op, TargetAttrInterface target) { + return op->setAttr(DLTIDialect::kTargetDescAttrName, target); +} diff --git a/mlir/lib/Dialect/DLTI/Transforms/CMakeLists.txt b/mlir/lib/Dialect/DLTI/Transforms/CMakeLists.txt new file mode 100644 index 0000000000000..52394988471a1 --- /dev/null +++ b/mlir/lib/Dialect/DLTI/Transforms/CMakeLists.txt @@ -0,0 +1,12 @@ +add_mlir_dialect_library(MLIRDLTITransforms + SetTargetSpecs.cpp + + ADDITIONAL_HEADER_DIRS + ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/DLTI + + DEPENDS + MLIRDLTIPassIncGen + + LINK_LIBS PUBLIC + MLIRDataLayoutInterfaces + ) diff --git a/mlir/lib/Dialect/DLTI/Transforms/SetTargetSpecs.cpp b/mlir/lib/Dialect/DLTI/Transforms/SetTargetSpecs.cpp new file mode 100644 index 0000000000000..c002f2369937f --- /dev/null +++ b/mlir/lib/Dialect/DLTI/Transforms/SetTargetSpecs.cpp @@ -0,0 +1,52 @@ +//===- SetTargetSpecs.cpp - Sets target specs -----------------------------===// +// +// 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 implements the `DltiSetTargetSpecsFromTarget` pass. +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/DLTI/Transforms/Passes.h" + +#include "mlir/Dialect/DLTI/DLTI.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +#define GEN_PASS_DEF_DLTISETTARGETSPECSFROMTARGET +#include "mlir/Dialect/DLTI/Transforms/Passes.h.inc" +} // namespace mlir + +using namespace mlir; + +namespace { +struct SetTargetSpecs + : public impl::DltiSetTargetSpecsFromTargetBase { + using Base::Base; + + void runOnOperation() override { + if (failed(setTargetSpecsFromTarget(getOperation()))) + return signalPassFailure(); + } +}; +} // namespace + +LogicalResult mlir::setTargetSpecsFromTarget(Operation *op) { + auto dlOp = dyn_cast(op); + if (!dlOp) + return op->emitError("Op doesn't implement `DataLayoutOpInterface`."); + TargetAttrInterface target = dlOp.getTargetAttr(); + if (!target) + return op->emitError("Op doesn't have a target."); + TargetSpec spec; + if (failed(target.setTargetSpec(spec))) + return failure(); + if (spec.systemSpec) + dlOp.setTargetSystemSpec(spec.systemSpec); + if (spec.dataLayout) + dlOp.setDataLayoutSpec(spec.dataLayout); + return success(); +} diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp index f20126618060a..f8e08ebbd5daa 100644 --- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp +++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp @@ -1793,6 +1793,32 @@ void GPUModuleOp::setTargets(ArrayRef targets) { targetsAttr = ArrayAttr::get(getContext(), targetsVector); } +DataLayoutSpecInterface GPUModuleOp::getDataLayoutSpec() { + return mlir::impl::getDataLayoutSpec(getOperation()); +} + +void GPUModuleOp::setDataLayoutSpec(DataLayoutSpecInterface spec) { + return mlir::impl::setDataLayoutSpec(getOperation(), spec); +} + +TargetSystemSpecInterface GPUModuleOp::getTargetSystemSpec() { + return mlir::impl::getTargetSystemSpec(getOperation()); +} + +void GPUModuleOp::setTargetSystemSpec(TargetSystemSpecInterface spec) { + return mlir::impl::setTargetSystemSpec(getOperation(), spec); +} + +mlir::TargetAttrInterface GPUModuleOp::getTargetAttr() { + if (ArrayAttr targets = getTargetsAttr(); targets && targets.size() == 1) + return dyn_cast(targets[0]); + return nullptr; +} + +void GPUModuleOp::setTargetAttr(mlir::TargetAttrInterface target) { + getProperties().targets = ArrayAttr::get(target.getContext(), {target}); +} + //===----------------------------------------------------------------------===// // GPUBinaryOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/LLVMIR/CMakeLists.txt b/mlir/lib/Dialect/LLVMIR/CMakeLists.txt index c9a3b97294562..bd58983ec0ae6 100644 --- a/mlir/lib/Dialect/LLVMIR/CMakeLists.txt +++ b/mlir/lib/Dialect/LLVMIR/CMakeLists.txt @@ -60,6 +60,7 @@ add_mlir_dialect_library(MLIRNVVMDialect LINK_LIBS PUBLIC MLIRIR MLIRLLVMDialect + MLIRDataLayoutInterfaces MLIRSideEffectInterfaces MLIRInferIntRangeInterface ) @@ -83,6 +84,7 @@ add_mlir_dialect_library(MLIRROCDLDialect LINK_LIBS PUBLIC MLIRIR MLIRLLVMDialect + MLIRDataLayoutInterfaces MLIRSideEffectInterfaces ) diff --git a/mlir/lib/Target/LLVM/CMakeLists.txt b/mlir/lib/Target/LLVM/CMakeLists.txt index 83fbf7a5fe5f3..8370a0e93ac55 100644 --- a/mlir/lib/Target/LLVM/CMakeLists.txt +++ b/mlir/lib/Target/LLVM/CMakeLists.txt @@ -1,3 +1,8 @@ +set(LLVM_OPTIONAL_SOURCES + ModuleToObject.cpp + Target.cpp + ) + add_mlir_library(MLIRTargetLLVM ModuleToObject.cpp @@ -21,6 +26,18 @@ add_mlir_library(MLIRTargetLLVM MLIRTargetLLVMIRExport ) +add_mlir_library(MLIRTargetInfo + Target.cpp + + LINK_COMPONENTS + Core + MC + Support + Target + LINK_LIBS PUBLIC + MLIRDataLayoutInterfaces +) + if ("NVPTX" IN_LIST LLVM_TARGETS_TO_BUILD) set(NVPTX_LIBS NVPTXCodeGen @@ -47,6 +64,8 @@ add_mlir_dialect_library(MLIRNVVMTarget MLIRGPUDialect MLIRTargetLLVM MLIRNVVMToLLVMIRTranslation + MLIRTargetInfo + MLIRTargetLLVMIRImport ) if ("NVPTX" IN_LIST LLVM_TARGETS_TO_BUILD) diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp index 914a349696617..65cb2d302e582 100644 --- a/mlir/lib/Target/LLVM/NVVM/Target.cpp +++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp @@ -21,10 +21,12 @@ #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/DialectResourceBlobManager.h" #include "mlir/Target/LLVM/NVVM/Utils.h" +#include "mlir/Target/LLVM/Target.h" #include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h" #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h" #include "mlir/Target/LLVMIR/Export.h" +#include "mlir/Target/LLVMIR/Import.h" #include "llvm/ADT/ScopeExit.h" #include "llvm/Config/Targets.h" @@ -66,6 +68,16 @@ class NVVMTargetAttrImpl const SmallVector &object, const gpu::TargetOptions &options) const; }; + +// Implementation of the `::mlir::TargetAttrInterface` model. +class NVVMTargetInfoAttrImpl + : public mlir::TargetAttrInterface::FallbackModel { +public: + StringRef getTargetTriple(Attribute attribute) const; + StringRef getTargetChip(Attribute attribute) const; + std::string getTargetFeatures(Attribute attribute) const; + LogicalResult setTargetSpec(Attribute attribute, TargetSpec &spec) const; +}; } // namespace // Register the NVVM dialect, the NVVM translation & the target interface. @@ -73,6 +85,7 @@ void mlir::NVVM::registerNVVMTargetInterfaceExternalModels( DialectRegistry ®istry) { registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) { NVVMTargetAttr::attachInterface(*ctx); + NVVMTargetAttr::attachInterface(*ctx); }); } @@ -94,6 +107,19 @@ StringRef mlir::NVVM::getCUDAToolkitPath() { return __DEFAULT_CUDATOOLKIT_PATH__; } +static void initializeTarget() { + static llvm::once_flag initializeBackendOnce; + llvm::call_once(initializeBackendOnce, []() { + // If the `NVPTX` LLVM target was built, initialize it. +#if LLVM_HAS_NVPTX_TARGET + LLVMInitializeNVPTXTarget(); + LLVMInitializeNVPTXTargetInfo(); + LLVMInitializeNVPTXTargetMC(); + LLVMInitializeNVPTXAsmPrinter(); +#endif + }); +} + SerializeGPUModuleBase::SerializeGPUModuleBase( Operation &module, NVVMTargetAttr target, const gpu::TargetOptions &targetOptions) @@ -118,18 +144,7 @@ SerializeGPUModuleBase::SerializeGPUModuleBase( (void)appendStandardLibs(); } -void SerializeGPUModuleBase::init() { - static llvm::once_flag initializeBackendOnce; - llvm::call_once(initializeBackendOnce, []() { - // If the `NVPTX` LLVM target was built, initialize it. -#if LLVM_HAS_NVPTX_TARGET - LLVMInitializeNVPTXTarget(); - LLVMInitializeNVPTXTargetInfo(); - LLVMInitializeNVPTXTargetMC(); - LLVMInitializeNVPTXAsmPrinter(); -#endif - }); -} +void SerializeGPUModuleBase::init() { initializeTarget(); } NVVMTargetAttr SerializeGPUModuleBase::getTarget() const { return target; } @@ -800,3 +815,28 @@ NVVMTargetAttrImpl::createObject(Attribute attribute, Operation *module, builder.getStringAttr(StringRef(object.data(), object.size())), objectProps, /*kernels=*/nullptr); } + +StringRef NVVMTargetInfoAttrImpl::getTargetTriple(Attribute attribute) const { + return cast(attribute).getTriple(); +} + +StringRef NVVMTargetInfoAttrImpl::getTargetChip(Attribute attribute) const { + return cast(attribute).getChip(); +} + +std::string +NVVMTargetInfoAttrImpl::getTargetFeatures(Attribute attribute) const { + return cast(attribute).getFeatures().str(); +} + +LogicalResult NVVMTargetInfoAttrImpl::setTargetSpec(Attribute attribute, + TargetSpec &spec) const { + initializeTarget(); + FailureOr info = + TargetInfo::getTargetInfo(cast(attribute)); + if (failed(info)) + return failure(); + spec.dataLayout = + translateDataLayout(info->getDataLayout(), attribute.getContext()); + return success(spec.dataLayout != nullptr); +} diff --git a/mlir/lib/Target/LLVM/Target.cpp b/mlir/lib/Target/LLVM/Target.cpp new file mode 100644 index 0000000000000..89309a2320813 --- /dev/null +++ b/mlir/lib/Target/LLVM/Target.cpp @@ -0,0 +1,89 @@ +//===- Target.cpp - 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 defines utilities to interact with LLVM targets by querying an MLIR +// target. +// +//===----------------------------------------------------------------------===// + +#include "mlir/Target/LLVM/Target.h" + +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Support/Debug.h" +#include "llvm/Target/TargetMachine.h" + +#define DEBUG_TYPE "mlir-llvm-target" + +using namespace mlir; + +llvm::Triple getTargetTriple(TargetAttrInterface target) { + return llvm::Triple(target.getTargetTriple()); +} + +static FailureOr getLLVMTarget(StringRef triple) { + std::string error; + const llvm::Target *target = + llvm::TargetRegistry::lookupTarget(triple, error); + if (error.empty()) + return target; + LLVM_DEBUG({ + llvm::dbgs() << "Failed to retrieve the target with: `" << error << "`\n"; + }); + return failure(); +} + +FailureOr +mlir::getLLVMTarget(TargetAttrInterface target) { + return ::getLLVMTarget(target.getTargetTriple()); +} + +FailureOr +mlir::getLLVMDataLayout(StringRef triple, StringRef chip, StringRef features) { + FailureOr target = ::getLLVMTarget(triple); + if (failed(target)) + return failure(); + std::unique_ptr tgtMachine( + (*target)->createTargetMachine(llvm::Triple(triple), chip, features, {}, + {})); + return tgtMachine->createDataLayout(); +} + +FailureOr TargetInfo::getTargetInfo(StringRef triple, + StringRef chip, + StringRef features) { + FailureOr llvmTgt = ::getLLVMTarget(triple); + if (failed(llvmTgt)) + return failure(); + return FailureOr(TargetInfo((*llvmTgt)->createTargetMachine( + llvm::Triple(triple), chip, features, {}, {}))); +} + +TargetInfo::TargetInfo(llvm::TargetMachine *targetMachine) + : targetMachine(targetMachine) {} + +TargetInfo::~TargetInfo() = default; + +StringRef TargetInfo::getTargetChip() const { + return targetMachine->getTargetCPU(); +} + +StringRef TargetInfo::getTargetFeatures() const { + return targetMachine->getTargetFeatureString(); +} + +const llvm::Triple &TargetInfo::getTriple() const { + return targetMachine->getTargetTriple(); +} + +const llvm::Target &TargetInfo::getTarget() const { + return targetMachine->getTarget(); +} + +const llvm::DataLayout TargetInfo::getDataLayout() const { + return targetMachine->createDataLayout(); +} diff --git a/mlir/test/Dialect/DLTI/set-target-spec.mlir b/mlir/test/Dialect/DLTI/set-target-spec.mlir new file mode 100644 index 0000000000000..e7c5a91b6c5ff --- /dev/null +++ b/mlir/test/Dialect/DLTI/set-target-spec.mlir @@ -0,0 +1,26 @@ +// REQUIRES: host-supports-nvptx +// RUN: mlir-opt %s --pass-pipeline="builtin.module(gpu.module(dlti-set-target-specs))" | FileCheck %s + +module attributes {gpu.container_module} { + // CHECK-LABEL:gpu.module @kernel_module1 + // CHECK: dlti.dl_spec = #dlti.dl_spec< + // CHECK-SAME: !llvm.ptr<6> = dense<32> : vector<4xi64>, + // CHECK-SAME: i64 = dense<64> : vector<2xi64>, + // CHECK-SAME: i128 = dense<128> : vector<2xi64>, + // CHECK-SAME: !llvm.ptr = dense<64> : vector<4xi64>, + // CHECK-SAME: i1 = dense<8> : vector<2xi64>, + // CHECK-SAME: i8 = dense<8> : vector<2xi64>, + // CHECK-SAME: i16 = dense<16> : vector<2xi64>, + // CHECK-SAME: i32 = dense<32> : vector<2xi64>, + // CHECK-SAME: f16 = dense<16> : vector<2xi64>, + // CHECK-SAME: f64 = dense<64> : vector<2xi64>, + // CHECK-SAME: f128 = dense<128> : vector<2xi64>, + // CHECK-SAME: "dlti.endianness" = "little">} + gpu.module @kernel_module1 [#nvvm.target] { + llvm.func @kernel(%arg0: i32, %arg1: !llvm.ptr, + %arg2: !llvm.ptr, %arg3: i64, %arg4: i64, + %arg5: i64) attributes {gpu.kernel} { + llvm.return + } + } +} diff --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp index fd81f3021aa9b..55e61a67739a9 100644 --- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp +++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp @@ -35,7 +35,7 @@ constexpr static llvm::StringLiteral kGlobalKeyName = "dltest.global_memory_space"; constexpr static llvm::StringLiteral kStackAlignmentKeyName = "dltest.stack_alignment"; - +constexpr static llvm::StringLiteral kTargetAttrName = "dltest.target"; constexpr static llvm::StringLiteral kTargetSystemDescAttrName = "dl_target_sys_desc_test.target_system_spec"; @@ -273,11 +273,27 @@ struct OpWithLayout : public Op { return getOperation()->getAttrOfType(kAttrName); } + void setDataLayoutSpec(DataLayoutSpecInterface spec) { + return getOperation()->setAttr(kAttrName, spec); + } + TargetSystemSpecInterface getTargetSystemSpec() { return getOperation()->getAttrOfType( kTargetSystemDescAttrName); } + void setTargetSystemSpec(TargetSystemSpecInterface spec) { + return getOperation()->setAttr(kTargetSystemDescAttrName, spec); + } + + TargetAttrInterface getTargetAttr() { + return getOperation()->getAttrOfType(kTargetAttrName); + } + + void setTargetAttr(TargetAttrInterface target) { + return getOperation()->setAttr(kTargetAttrName, target); + } + static llvm::TypeSize getTypeSizeInBits(Type type, const DataLayout &dataLayout, DataLayoutEntryListRef params) { @@ -325,11 +341,27 @@ struct OpWith7BitByte return getOperation()->getAttrOfType(kAttrName); } + void setDataLayoutSpec(DataLayoutSpecInterface spec) { + return getOperation()->setAttr(kAttrName, spec); + } + TargetSystemSpecInterface getTargetSystemSpec() { return getOperation()->getAttrOfType( kTargetSystemDescAttrName); } + void setTargetSystemSpec(TargetSystemSpecInterface spec) { + return getOperation()->setAttr(kTargetSystemDescAttrName, spec); + } + + TargetAttrInterface getTargetAttr() { + return getOperation()->getAttrOfType(kTargetAttrName); + } + + void setTargetAttr(TargetAttrInterface target) { + return getOperation()->setAttr(kTargetAttrName, target); + } + // Bytes are assumed to be 7-bit here. static llvm::TypeSize getTypeSize(Type type, const DataLayout &dataLayout, DataLayoutEntryListRef params) {