From ae04c64346812f7a84c4f976988af9fef3113bf5 Mon Sep 17 00:00:00 2001 From: Armin Ale Date: Thu, 27 Feb 2025 19:44:48 +0000 Subject: [PATCH 1/2] add constraints and runtime APIs for transpose --- include/ttmlir/Dialect/TTNN/IR/TTNNOps.td | 4 +- include/ttmlir/OpModel/TTNN/TTNNOpModel.h | 17 +++++ lib/Dialect/TTNN/IR/TTNNOpModelInterface.cpp | 33 ++++++++++ lib/OpModel/TTNN/MetalHeaders.h | 1 + lib/OpModel/TTNN/TTNNOpModelLib.cpp | 65 ++++++++++++++++++++ 5 files changed, 119 insertions(+), 1 deletion(-) diff --git a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td index 3a45d05693..6c93512666 100644 --- a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td +++ b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td @@ -879,7 +879,9 @@ def TTNN_SoftmaxOp : TTNN_Op<"softmax", let hasVerifier = 1; } -def TTNN_TransposeOp : TTNN_Op<"transpose"> { +def TTNN_TransposeOp : TTNN_Op<"transpose", + [DeclareOpInterfaceMethods] + > { let summary = "Transpose op."; let description = [{ Transpose tensor along two given dimensions. diff --git a/include/ttmlir/OpModel/TTNN/TTNNOpModel.h b/include/ttmlir/OpModel/TTNN/TTNNOpModel.h index 27eff0ca5d..47b0144cd3 100644 --- a/include/ttmlir/OpModel/TTNN/TTNNOpModel.h +++ b/include/ttmlir/OpModel/TTNN/TTNNOpModel.h @@ -120,6 +120,23 @@ getOpRuntime(llvm::ArrayRef inputShape, }; // namespace ReshapeOpInterface +//===----------------------------------------------------------------------===// +// TransposeOp +//===----------------------------------------------------------------------===// + +namespace TransposeOpInterface { +llvm::Expected> +getOpConstraints(llvm::ArrayRef inputShape, + mlir::tt::ttnn::TTNNLayoutAttr inputLayout, const int dim0, + const int dim1, mlir::tt::ttnn::TTNNLayoutAttr outputLayout); + +llvm::Expected +getOpRuntime(llvm::ArrayRef inputShape, + mlir::tt::ttnn::TTNNLayoutAttr inputLayout, const int dim0, + const int dim1, mlir::tt::ttnn::TTNNLayoutAttr outputLayout); + +}; // namespace TransposeOpInterface + //===----------------------------------------------------------------------===// // MatmulOp //===----------------------------------------------------------------------===// diff --git a/lib/Dialect/TTNN/IR/TTNNOpModelInterface.cpp b/lib/Dialect/TTNN/IR/TTNNOpModelInterface.cpp index d1874441f6..b7ac625d6d 100644 --- a/lib/Dialect/TTNN/IR/TTNNOpModelInterface.cpp +++ b/lib/Dialect/TTNN/IR/TTNNOpModelInterface.cpp @@ -238,6 +238,39 @@ ReshapeOp::getOpRuntime(const std::vector &inputs, outputShape, output); } +//===----------------------------------------------------------------------===// +// TransposeOp - TTNN Op Model Interface +//===----------------------------------------------------------------------===// + +llvm::Expected> +TransposeOp::getOpConstraints(const std::vector &inputs, + const TTNNLayoutAttr &output) { + assert(inputs.size() == 1); + + const auto inputShape = + mlir::cast(getOperand().getType()).getShape(); + + llvm::Expected check = detail::checkDeviceWorkerGrid(getOperation()); + if (!check) { + return check.takeError(); + } + + return op_model::ttnn::TransposeOpInterface::getOpConstraints( + inputShape, inputs[0], getDim0(), getDim1(), output); +} + +llvm::Expected +TransposeOp::getOpRuntime(const std::vector &inputs, + const TTNNLayoutAttr &output) { + assert(inputs.size() == 1); + + const auto inputShape = + mlir::cast(getOperand().getType()).getShape(); + + return op_model::ttnn::TransposeOpInterface::getOpRuntime( + inputShape, inputs[0], getDim0(), getDim1(), output); +} + //===----------------------------------------------------------------------===// // MatmulOp - TTNN Op Model Interface //===----------------------------------------------------------------------===// diff --git a/lib/OpModel/TTNN/MetalHeaders.h b/lib/OpModel/TTNN/MetalHeaders.h index cbf929dd3b..9210a8db44 100644 --- a/lib/OpModel/TTNN/MetalHeaders.h +++ b/lib/OpModel/TTNN/MetalHeaders.h @@ -62,6 +62,7 @@ #include "ttnn/graph/graph_query_op_runtime.hpp" #include "ttnn/graph/graph_trace_utils.hpp" #include "ttnn/operations/data_movement/reshape_view/reshape.hpp" +#include "ttnn/operations/data_movement/transpose/transpose.hpp" #include "ttnn/operations/eltwise/binary/binary.hpp" #include "ttnn/operations/eltwise/unary/unary.hpp" #include "ttnn/operations/matmul/matmul.hpp" diff --git a/lib/OpModel/TTNN/TTNNOpModelLib.cpp b/lib/OpModel/TTNN/TTNNOpModelLib.cpp index f8b0de05ec..2562f8efac 100644 --- a/lib/OpModel/TTNN/TTNNOpModelLib.cpp +++ b/lib/OpModel/TTNN/TTNNOpModelLib.cpp @@ -566,6 +566,71 @@ ReshapeOpInterface::getOpRuntime(llvm::ArrayRef inputShape, #endif // TTMLIR_ENABLE_OPMODEL } +//===----------------------------------------------------------------------===// +// TransposeOp +//===----------------------------------------------------------------------===// +llvm::Expected> +TransposeOpInterface::getOpConstraints( + llvm::ArrayRef inputShape, + mlir::tt::ttnn::TTNNLayoutAttr inputLayout, const int dim0, const int dim1, + mlir::tt::ttnn::TTNNLayoutAttr outputLayout) { +#ifdef TTMLIR_ENABLE_OPMODEL + auto transposeOpQuery = [](llvm::ArrayRef inputShape, + mlir::tt::ttnn::TTNNLayoutAttr inputLayout, + const int dim0, const int dim1, + mlir::tt::ttnn::TTNNLayoutAttr outputLayout) { + // open device device, will close it at the end of function + ::tt::tt_metal::v0::IDevice *device = + SingletonDeviceContext::getInstance().getDevice(); + + // prepare io specs + const auto [inputSpec] = detail::convertToTensorSpec( + device, std::make_tuple(inputShape, inputLayout)); + + // run op constraint query + return ::ttnn::graph::query_op_constraints( + ::ttnn::transpose, device, inputSpec, dim0, dim1, + conversion::getMemoryConfig(outputLayout)); + }; + + return operation::getOpConstraints("TransposeOpInterface", transposeOpQuery, + inputShape, inputLayout, dim0, dim1, + outputLayout); +#else + return std::make_tuple(0, 0, 0); +#endif // TTMLIR_ENABLE_OPMODEL +} + +llvm::Expected TransposeOpInterface::getOpRuntime( + llvm::ArrayRef inputShape, + mlir::tt::ttnn::TTNNLayoutAttr inputLayout, const int dim0, const int dim1, + mlir::tt::ttnn::TTNNLayoutAttr outputLayout) { +#ifdef TTMLIR_ENABLE_OPMODEL + auto transposeOpQuery = [](llvm::ArrayRef inputShape, + mlir::tt::ttnn::TTNNLayoutAttr inputLayout, + const int dim0, const int dim1, + mlir::tt::ttnn::TTNNLayoutAttr outputLayout) { + // open device device, will close it at the end of function + ::tt::tt_metal::v0::IDevice *device = + SingletonDeviceContext::getInstance().getDevice(); + + // prepare io specs + const auto [inputSpec] = detail::convertToTensorSpec( + device, std::make_tuple(inputShape, inputLayout)); + + return ::ttnn::graph::query_op_runtime( + ::ttnn::transpose, device, inputSpec, dim0, dim1, + conversion::getMemoryConfig(outputLayout)); + }; + + return operation::getOpRuntime("TransposeOpInterface", transposeOpQuery, + inputShape, inputLayout, dim0, dim1, + outputLayout); +#else + return llvm::createStringError("Not Implemented"); +#endif // TTMLIR_ENABLE_OPMODEL +} + //===----------------------------------------------------------------------===// // MatmulOp //===----------------------------------------------------------------------===// From 7f924dafb1cf35fa5370e51fd9cc967f33d35804 Mon Sep 17 00:00:00 2001 From: Armin Ale Date: Thu, 27 Feb 2025 19:45:08 +0000 Subject: [PATCH 2/2] add unittests --- .../OpModel/TTNN/Lib/TestOpModelLib.cpp | 39 +++++++++++++++++++ .../OpModel/TTNN/Op/TestOpModelInterface.cpp | 35 ++++++++++++++++- 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/test/unittests/OpModel/TTNN/Lib/TestOpModelLib.cpp b/test/unittests/OpModel/TTNN/Lib/TestOpModelLib.cpp index 2dfd84c6e6..18d58f6ba2 100644 --- a/test/unittests/OpModel/TTNN/Lib/TestOpModelLib.cpp +++ b/test/unittests/OpModel/TTNN/Lib/TestOpModelLib.cpp @@ -324,6 +324,45 @@ TEST_F(OpModelTest, Reshape) { EXPECT_TRUE(runtimeExp.get() > 0); } +TEST_F(OpModelTest, Transpose) { + const llvm::SmallVector tensorShape = {workerCoresN300, 1024}; + const auto workerGrid = CreateWorkerGrid(gridShapeHwN300); + const mlir::tt::ttnn::TTNNLayoutAttr layoutDRAM = + CreateTiledLayout(tensorShape, mlir::tt::ttnn::BufferType::DRAM, + mlir::tt::ttnn::TensorMemoryLayout::Interleaved); + const mlir::tt::ttnn::TTNNLayoutAttr layoutL1 = + CreateTiledLayout(tensorShape, mlir::tt::ttnn::BufferType::L1, + mlir::tt::ttnn::TensorMemoryLayout::Interleaved); + auto legalExp = Device::getDeviceConstraints(workerGrid); + EXPECT_TRUE(static_cast(legalExp)); + + auto constraintsExp = TransposeOpInterface::getOpConstraints( + tensorShape, layoutDRAM, 0, 1, layoutDRAM); + EXPECT_TRUE(static_cast(constraintsExp)); + auto [cb_size, peak_size, output_size] = constraintsExp.get(); + EXPECT_EQ(cb_size, 8192); + EXPECT_EQ(output_size, 0); + EXPECT_EQ(peak_size, 0); + + auto runtimeExp = TransposeOpInterface::getOpRuntime(tensorShape, layoutDRAM, + 0, 1, layoutDRAM); + EXPECT_TRUE(static_cast(runtimeExp)); + EXPECT_TRUE(runtimeExp.get() > 0); + + constraintsExp = TransposeOpInterface::getOpConstraints( + tensorShape, layoutDRAM, 0, 1, layoutL1); + EXPECT_TRUE(static_cast(constraintsExp)); + std::tie(cb_size, peak_size, output_size) = constraintsExp.get(); + EXPECT_EQ(cb_size, 8192); + EXPECT_EQ(output_size, 2048); + EXPECT_EQ(peak_size, 2048); + + runtimeExp = TransposeOpInterface::getOpRuntime(tensorShape, layoutDRAM, 0, 1, + layoutL1); + EXPECT_TRUE(static_cast(runtimeExp)); + EXPECT_TRUE(runtimeExp.get() > 0); +} + TEST_F(OpModelTest, SoftmaxSharded) { const llvm::SmallVector tensorShape = {16 * workerCoresN300 * 32, 32}; diff --git a/test/unittests/OpModel/TTNN/Op/TestOpModelInterface.cpp b/test/unittests/OpModel/TTNN/Op/TestOpModelInterface.cpp index 27536e8bc6..e9efe1aa23 100644 --- a/test/unittests/OpModel/TTNN/Op/TestOpModelInterface.cpp +++ b/test/unittests/OpModel/TTNN/Op/TestOpModelInterface.cpp @@ -266,7 +266,7 @@ TEST_F(OpModelBase, reshapeOp) { reshape.setShapeAttr(builder.getArrayAttr(llvm::SmallVector{ builder.getI64IntegerAttr(64 * 4), builder.getI64IntegerAttr(1024 / 4)})); - // test mean Op interface + // test reshape Op interface auto constraintsExp = getOpConstraints(reshape.getOperation()); if (constraintsExp) { auto l1 = constraintsExp.get(); @@ -287,4 +287,37 @@ TEST_F(OpModelBase, reshapeOp) { } } +TEST_F(OpModelBase, transposeOp) { + // create TransposeOp + llvm::SmallVector tensorShapeA = {64, 1024}; + llvm::SmallVector tensorShapeO = {1024, 64}; + + auto input = createEmptyTensor(tensorShapeA); + auto output = createEmptyTensor(tensorShapeO); + + auto transpose = builder.create(builder.getUnknownLoc(), + output.getType(), input, 0, 1); + transpose->setAttr(DeviceAttr::name, getFakeDeviceAttr()); + + // test transpose Op interface + auto constraintsExp = getOpConstraints(transpose.getOperation()); + if (constraintsExp) { + auto l1 = constraintsExp.get(); + const auto &[cb_size, peak_size, output_size] = l1; + EXPECT_EQ(cb_size, 8192); + EXPECT_EQ(peak_size, 2048); + EXPECT_EQ(output_size, 2048); + } else { + FAIL() << "Missing L1 constraints; Error=" + << llvm::toString(constraintsExp.takeError()) << std::endl; + } + + auto runtimeExp = getOpRuntime(transpose.getOperation()); + if (runtimeExp) { + EXPECT_TRUE(runtimeExp.get() > 0); + } else { + FAIL() << llvm::toString(runtimeExp.takeError()); + } +} + } // namespace mlir::tt::ttnn