From 46612ac10a069e01137ac0325fbf849c2546008c Mon Sep 17 00:00:00 2001 From: Vlad Roubtsov Date: Wed, 5 Mar 2025 19:33:28 +0000 Subject: [PATCH] move getDpsOutputs() helper into tt::ttir and a dialect-specific Utils.h --- include/ttmlir/Dialect/TTIR/IR/TTIROps.h | 23 ++++----- include/ttmlir/Dialect/TTIR/IR/TTIROps.td | 12 ++--- include/ttmlir/Dialect/TTIR/IR/Utils.h | 59 +++++++++++++++++++++ include/ttmlir/Utils.h | 62 ++++------------------- 4 files changed, 86 insertions(+), 70 deletions(-) create mode 100644 include/ttmlir/Dialect/TTIR/IR/Utils.h diff --git a/include/ttmlir/Dialect/TTIR/IR/TTIROps.h b/include/ttmlir/Dialect/TTIR/IR/TTIROps.h index 62239a5000..6cba434256 100644 --- a/include/ttmlir/Dialect/TTIR/IR/TTIROps.h +++ b/include/ttmlir/Dialect/TTIR/IR/TTIROps.h @@ -8,18 +8,17 @@ #include "ttmlir/Dialect/TT/IR/TTOpsTypes.h" #include "ttmlir/Dialect/TTIR/IR/TTIROpsInterfaces.h" #include "ttmlir/Dialect/TTIR/IR/TTIRTraits.h" - -#include "ttmlir/Utils.h" - -#include "mlir/Bytecode/BytecodeOpInterface.h" -#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Interfaces/ControlFlowInterfaces.h" -#include "mlir/Interfaces/DestinationStyleOpInterface.h" -#include "mlir/Interfaces/InferTypeOpInterface.h" -#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "ttmlir/Dialect/TTIR/IR/Utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include namespace mlir::tt::ttir { diff --git a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td index 9cffa8bfa1..0cb5c9e199 100644 --- a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td +++ b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td @@ -24,7 +24,7 @@ class TTIR_DPSOp traits = []> : TTIR_Op { let extraClassDeclaration = [{ // base implementation that will detect getOutputMutable() or getOutputsMutable() - MutableOperandRange getDpsInitsMutable() { return ::ttmlir::utils::getDpsOutputs(this); } + MutableOperandRange getDpsInitsMutable() { return ttir::getDpsOutputs(this); } }]; } @@ -112,7 +112,7 @@ def TTIR_GenericOp : TTIR_BufferizableOp<"generic", [AttrSizedOperandSegments, N let hasVerifier = 1; let extraClassDeclaration = [{ - MutableOperandRange getDpsInitsMutable() { return ::ttmlir::utils::getDpsOutputs(this); } + MutableOperandRange getDpsInitsMutable() { return ttir::getDpsOutputs(this); } }]; } @@ -138,7 +138,7 @@ def TTIR_ToLayoutOp : TTIR_BufferizableOp<"to_layout"> { let results = (outs Variadic:$results); let extraClassDeclaration = [{ - MutableOperandRange getDpsInitsMutable() { return ::ttmlir::utils::getDpsOutputs(this); } + MutableOperandRange getDpsInitsMutable() { return ttir::getDpsOutputs(this); } struct CompoundComponents { bool isLayoutChange = false; @@ -781,7 +781,7 @@ class TTIR_ReductionOp traits = []> : let results = (outs AnyRankedTensor:$result); let extraClassDeclaration = [{ - MutableOperandRange getDpsInitsMutable() { return ::ttmlir::utils::getDpsOutputs(this); } + MutableOperandRange getDpsInitsMutable() { return ttir::getDpsOutputs(this); } void buildGenericRegion(::mlir::OpBuilder &opBuilder, ::mlir::Block* block); // Returns the indexing maps and iterator types for the reduction op. @@ -1826,7 +1826,7 @@ class TTIR_GenericElementwiseUnaryOp traits = []> : TTIR_ElementwiseUnaryOp { let extraClassDeclaration = [{ - MutableOperandRange getDpsInitsMutable() { return ::ttmlir::utils::getDpsOutputs(this); } + MutableOperandRange getDpsInitsMutable() { return ttir::getDpsOutputs(this); } void buildGenericRegion(::mlir::OpBuilder &opBuilder, ::mlir::Block* block); std::pair<::mlir::ArrayAttr, ::mlir::ArrayAttr> getIndexingMaps(Builder &builder) { @@ -1856,7 +1856,7 @@ class TTIR_GenericElementwiseBinaryOp traits = []> TTIR_ElementwiseBinaryOp { let extraClassDeclaration = [{ - MutableOperandRange getDpsInitsMutable() { return ::ttmlir::utils::getDpsOutputs(this); } + MutableOperandRange getDpsInitsMutable() { return ttir::getDpsOutputs(this); } void buildGenericRegion(::mlir::OpBuilder &opBuilder, ::mlir::Block* block); std::pair<::mlir::ArrayAttr, ::mlir::ArrayAttr> getIndexingMaps(Builder &builder) { diff --git a/include/ttmlir/Dialect/TTIR/IR/Utils.h b/include/ttmlir/Dialect/TTIR/IR/Utils.h new file mode 100644 index 0000000000..8e9dc2c711 --- /dev/null +++ b/include/ttmlir/Dialect/TTIR/IR/Utils.h @@ -0,0 +1,59 @@ +// SPDX-FileCopyrightText: (c) 20245 Tenstorrent AI ULC +// +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TTMLIR_DIALECT_TTIR_IR_UTILS_H +#define TTMLIR_DIALECT_TTIR_IR_UTILS_H + +#include + +#include + +namespace mlir::tt::ttir { + +// detect the presence of 'getOutputsMutable()' in 'Op': +template +inline constexpr bool has_variadic_outputs = false; + +template +inline constexpr bool has_variadic_outputs< + Op, std::void_t().getOutputsMutable())>> = true; + +namespace impl { + +template +struct getDpsOutputs { + static mlir::MutableOperandRange evaluate(Op *op) { + return op->getOutputMutable(); + } +}; + +template +struct getDpsOutputs>> { + static mlir::MutableOperandRange evaluate(Op *op) { + return op->getOutputsMutable(); + } +}; + +} // namespace impl + +// A helper for simplifying DPS tablegen derivations with 'arguments' of any +// form in {AnyRankedTensor:$output, Variadic:$outputs}. +// +// If a base tablegen 'class' adds this extra class declaration, derived 'def's +// don't need to overrride it just to switch from single to variadic type of +// '$outputs' (or vice versa): +// ... +// clang-format off +// let extraClassDeclaration = [{ +// MutableOperandRange getDpsInitsMutable() { return ttir::getDpsOutputs(this); } +// }] +// clang-format on +template +mlir::MutableOperandRange getDpsOutputs(Op *op) { + return impl::getDpsOutputs::evaluate(op); +} + +} // namespace mlir::tt::ttir + +#endif // TTMLIR_DIALECT_TTIR_IR_UTILS_H \ No newline at end of file diff --git a/include/ttmlir/Utils.h b/include/ttmlir/Utils.h index cd28d94e25..bef42610bd 100644 --- a/include/ttmlir/Utils.h +++ b/include/ttmlir/Utils.h @@ -5,19 +5,20 @@ #ifndef TTMLIR_UTILS_H #define TTMLIR_UTILS_H -#include "mlir-c/IR.h" -#include "mlir/CAPI/IR.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/IR/AffineMap.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringExtras.h" -#include "llvm/Support/Error.h" +#include +#include +#include +#include +#include +#include +#include +#include #include #include namespace ttmlir::utils { + template T alignUp(T ptr, T alignment) { return (ptr + alignment - 1) & ~(alignment - 1); @@ -459,49 +460,6 @@ OpTy replaceOpWithNewDPSOp(mlir::PatternRewriter &rewriter, mlir::Operation *op, return newOp; } -// detect the presence of 'getOutputsMutable()' in 'Op': -template -inline constexpr bool has_variadic_outputs = false; - -template -inline constexpr bool has_variadic_outputs< - Op, std::void_t().getOutputsMutable())>> = true; - -namespace impl { - -template -struct getDpsOutputs { - static mlir::MutableOperandRange evaluate(Op *op) { - return op->getOutputMutable(); - } -}; - -template -struct getDpsOutputs>> { - static mlir::MutableOperandRange evaluate(Op *op) { - return op->getOutputsMutable(); - } -}; - -} // namespace impl - -// A helper for simplifying DPS tablegen derivations with 'arguments' of any -// form in {AnyRankedTensor:$output, Variadic:$outputs}. -// -// If a base tablegen 'class' adds this extra class declaration, derived 'def's -// don't need to overrride it just to switch from single to variadic type of -// '$outputs' (or vice versa): -// ... -// clang-format off -// let extraClassDeclaration = [{ -// MutableOperandRange getDpsInitsMutable() { return ::ttmlir::utils::getDpsOutputs(this); } -// }] -// clang-format on -template -mlir::MutableOperandRange getDpsOutputs(Op *op) { - return impl::getDpsOutputs::evaluate(op); -} - } // namespace ttmlir::utils -#endif +#endif // TTMLIR_UTILS_H