From e09ed7358d7d9ee5b543ae2a1a0fa59ad25efe16 Mon Sep 17 00:00:00 2001 From: vprajapati-tt Date: Fri, 18 Oct 2024 15:44:13 +0000 Subject: [PATCH] Added TTNN Bindings updated imports Added TTNN Ops, Attrs, & Enums to Python Bindings --- include/ttmlir-c/Dialects.h | 1 + include/ttmlir-c/TTNNAttrs.h | 48 ++++++++ include/ttmlir/Bindings/Python/TTMLIRModule.h | 1 + lib/CAPI/CMakeLists.txt | 2 + lib/CAPI/Dialects.cpp | 2 + lib/CAPI/TTNNAttrs.cpp | 66 +++++++++++ python/CMakeLists.txt | 13 ++- python/TTMLIRModule.cpp | 5 + python/TTNNModule.cpp | 104 ++++++++++++++++++ .../{TTEnumBindings.td => TTEnumBinding.td} | 0 python/ttmlir/dialects/TTNNBinding.td | 10 ++ python/ttmlir/dialects/TTNNEnumBinding.td | 10 ++ python/ttmlir/dialects/ttnn.py | 7 ++ 13 files changed, 268 insertions(+), 1 deletion(-) create mode 100644 include/ttmlir-c/TTNNAttrs.h create mode 100644 lib/CAPI/TTNNAttrs.cpp create mode 100644 python/TTNNModule.cpp rename python/ttmlir/dialects/{TTEnumBindings.td => TTEnumBinding.td} (100%) create mode 100644 python/ttmlir/dialects/TTNNBinding.td create mode 100644 python/ttmlir/dialects/TTNNEnumBinding.td create mode 100644 python/ttmlir/dialects/ttnn.py diff --git a/include/ttmlir-c/Dialects.h b/include/ttmlir-c/Dialects.h index e0fca87907..32ba4eba95 100644 --- a/include/ttmlir-c/Dialects.h +++ b/include/ttmlir-c/Dialects.h @@ -14,6 +14,7 @@ extern "C" { MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(TT, tt); MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(TTIR, ttir); MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(TTKernel, ttkernel); +MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(TTNN, ttnn); #ifdef __cplusplus } diff --git a/include/ttmlir-c/TTNNAttrs.h b/include/ttmlir-c/TTNNAttrs.h new file mode 100644 index 0000000000..480deac237 --- /dev/null +++ b/include/ttmlir-c/TTNNAttrs.h @@ -0,0 +1,48 @@ +// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +// +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TTMLIR_C_TTNNATTRS_H +#define TTMLIR_C_TTNNATTRS_H + +#include "ttmlir-c/Dialects.h" + +#ifdef __cplusplus +extern "C" { +#endif + +MLIR_CAPI_EXPORTED MlirAttribute ttmlirTTNNCoreRangeAttrGet(MlirContext ctx, + int64_t *offset, + size_t offsetSize, + int64_t *size, + size_t sizeSize); + +MLIR_CAPI_EXPORTED MlirAttribute ttmlirTTNNCoreRangeArrayAttrGet( + MlirContext ctx, MlirAttribute *coreRangeAttrs, size_t coreRangeAttrsSize); + +MLIR_CAPI_EXPORTED MlirAttribute ttmlirTTNNLayoutAttrGet(MlirContext ctx, + uint32_t layout); + +MLIR_CAPI_EXPORTED MlirAttribute ttmlirTTNNTensorMemoryLayoutAttrGet( + MlirContext ctx, uint32_t tensorMemoryLayout); + +MLIR_CAPI_EXPORTED MlirAttribute +ttmlirTTNNBufferTypeAttrGet(MlirContext ctx, uint32_t bufferType); + +MLIR_CAPI_EXPORTED MlirAttribute ttmlirTTNNMemoryConfigAttrGet( + MlirContext ctx, MlirAttribute tensorMemoryLayoutAttr, + MlirAttribute bufferTypeAttr); + +MLIR_CAPI_EXPORTED MlirAttribute ttmlirTTNNShapeAttrGet(MlirContext ctx, + int64_t *shape, + size_t shapeSize); + +MLIR_CAPI_EXPORTED MlirAttribute ttmlirTTNNMeshShapeAttrGet(MlirContext ctx, + int64_t y, + int64_t x); + +#ifdef __cplusplus +} +#endif + +#endif // TTMLIR_C_TTNNATTRS_H diff --git a/include/ttmlir/Bindings/Python/TTMLIRModule.h b/include/ttmlir/Bindings/Python/TTMLIRModule.h index 88900ef881..0b1ed0aabb 100644 --- a/include/ttmlir/Bindings/Python/TTMLIRModule.h +++ b/include/ttmlir/Bindings/Python/TTMLIRModule.h @@ -26,6 +26,7 @@ namespace py = pybind11; namespace mlir::ttmlir::python { void populateTTModule(py::module &m); void populateTTKernelModule(py::module &m); +void populateTTNNModule(py::module &m); void populateOverridesModule(py::module &m); void populatePassesModule(py::module &m); } // namespace mlir::ttmlir::python diff --git a/lib/CAPI/CMakeLists.txt b/lib/CAPI/CMakeLists.txt index d3c6752b56..fbe108a19a 100644 --- a/lib/CAPI/CMakeLists.txt +++ b/lib/CAPI/CMakeLists.txt @@ -3,6 +3,7 @@ add_mlir_public_c_api_library(TTMLIRCAPI TTKernelTypes.cpp TTAttrs.cpp TTTypes.cpp + TTNNAttrs.cpp ADDITIONAL_HEADER_DIRS ${PROJECT_SOURCE_DIR}/include/ttmlir-c/ @@ -16,6 +17,7 @@ add_mlir_public_c_api_library(TTMLIRCAPI MLIRTTDialect MLIRTTIRDialect MLIRTTKernelDialect + MLIRTTNNDialect MLIRTTIRTransforms MLIRTTIRAnalysis ) diff --git a/lib/CAPI/Dialects.cpp b/lib/CAPI/Dialects.cpp index f9a3604e90..83687ae8be 100644 --- a/lib/CAPI/Dialects.cpp +++ b/lib/CAPI/Dialects.cpp @@ -8,8 +8,10 @@ #include "ttmlir/Dialect/TT/IR/TT.h" #include "ttmlir/Dialect/TTIR/IR/TTIR.h" #include "ttmlir/Dialect/TTKernel/IR/TTKernel.h" +#include "ttmlir/Dialect/TTNN/IR/TTNN.h" MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(TT, tt, mlir::tt::TTDialect) MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(TTIR, ttir, mlir::tt::ttir::TTIRDialect) MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(TTKernel, ttkernel, mlir::tt::ttkernel::TTKernelDialect) +MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(TTNN, ttnn, mlir::tt::ttnn::TTNNDialect) diff --git a/lib/CAPI/TTNNAttrs.cpp b/lib/CAPI/TTNNAttrs.cpp new file mode 100644 index 0000000000..512bd8fae8 --- /dev/null +++ b/lib/CAPI/TTNNAttrs.cpp @@ -0,0 +1,66 @@ +// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +// +// SPDX-License-Identifier: Apache-2.0 + +#include "ttmlir-c/TTNNAttrs.h" +#include "mlir/CAPI/IR.h" +#include "mlir/CAPI/Support.h" + +#include "ttmlir/Dialect/TTNN/IR/TTNNOpsAttrs.h" + +namespace mlir::tt::ttnn { + +MlirAttribute ttmlirTTNNCoreRangeAttrGet(MlirContext ctx, int64_t *offset, + size_t offsetSize, int64_t *size, + size_t sizeSize) { + return wrap(CoreRangeAttr::get(unwrap(ctx), {offset, offset + offsetSize}, + {size, size + sizeSize})); +} + +MlirAttribute ttmlirTTNNCoreRangeArrayAttrGet(MlirContext ctx, + MlirAttribute *coreRangeAttrs, + size_t coreRangeAttrsSize) { + std::vector coreRanges; + for (size_t i = 0; i < coreRangeAttrsSize; i++) { + coreRanges.push_back(mlir::cast(unwrap(coreRangeAttrs[i]))); + } + return wrap(ArrayAttr::get(unwrap(ctx), coreRanges)); +} + +MlirAttribute ttmlirTTNNLayoutAttrGet(MlirContext ctx, uint32_t layout) { + return wrap(LayoutAttr::get(unwrap(ctx), static_cast(layout))); +} + +MlirAttribute ttmlirTTNNTensorMemoryLayoutAttrGet(MlirContext ctx, + uint32_t tensorMemoryLayout) { + return wrap(TensorMemoryLayoutAttr::get( + unwrap(ctx), static_cast(tensorMemoryLayout))); +} + +MlirAttribute ttmlirTTNNBufferTypeAttrGet(MlirContext ctx, + uint32_t bufferType) { + return wrap( + BufferTypeAttr::get(unwrap(ctx), static_cast(bufferType))); +} + +MlirAttribute +ttmlirTTNNMemoryConfigAttrGet(MlirContext ctx, + MlirAttribute tensorMemoryLayoutAttr, + MlirAttribute bufferTypeAttr) { + return wrap(MemoryConfigAttr::get( + unwrap(ctx), + mlir::cast(unwrap(tensorMemoryLayoutAttr)), + mlir::cast(unwrap(bufferTypeAttr)))); +} + +MlirAttribute ttmlirTTNNShapeAttrGet(MlirContext ctx, int64_t *shape, + size_t shapeSize) { + return wrap(ShapeAttr::get(unwrap(ctx), {shape, shape + shapeSize})); +} + +MlirAttribute ttmlirTTNNMeshShapeAttrGet(MlirContext ctx, int64_t y, + int64_t x) { + return wrap(MeshShapeAttr::get(unwrap(ctx), y, x)); +} + +} // namespace mlir::tt::ttnn diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 03fa9e87a9..9d90b1c9c2 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -25,7 +25,7 @@ declare_mlir_dialect_python_bindings( ROOT_DIR "${TTMLIR_PYTHON_ROOT_DIR}" TD_FILE dialects/TTBinding.td GEN_ENUM_BINDINGS ON - GEN_ENUM_TD_FILE dialects/TTEnumBindings.td + GEN_ENUM_TD_FILE dialects/TTEnumBinding.td SOURCES dialects/tt.py DIALECT_NAME tt ) @@ -46,6 +46,16 @@ declare_mlir_dialect_python_bindings( DIALECT_NAME ttkernel ) +declare_mlir_dialect_python_bindings( + ADD_TO_PARENT TTMLIRPythonSources.Dialects + ROOT_DIR "${TTMLIR_PYTHON_ROOT_DIR}" + TD_FILE dialects/TTNNBinding.td + GEN_ENUM_BINDINGS ON + GEN_ENUM_TD_FILE dialects/TTNNEnumBinding.td + SOURCES dialects/ttnn.py + DIALECT_NAME ttnn +) + declare_mlir_python_sources(TTMLIRPythonSources.Overrides ROOT_DIR "${TTMLIR_PYTHON_ROOT_DIR}" ADD_TO_PARENT TTMLIRPythonSources @@ -71,6 +81,7 @@ declare_mlir_python_extension(TTMLIRPythonExtensions.Main TTMLIRModule.cpp TTModule.cpp TTKernelModule.cpp + TTNNModule.cpp Overrides.cpp Passes.cpp EMBED_CAPI_LINK_LIBS diff --git a/python/TTMLIRModule.cpp b/python/TTMLIRModule.cpp index d9f62ae7b8..4d259c216f 100644 --- a/python/TTMLIRModule.cpp +++ b/python/TTMLIRModule.cpp @@ -13,13 +13,16 @@ PYBIND11_MODULE(_ttmlir, m) { MlirDialectHandle tt_handle = mlirGetDialectHandle__tt__(); MlirDialectHandle ttir_handle = mlirGetDialectHandle__ttir__(); MlirDialectHandle ttkernel_handle = mlirGetDialectHandle__ttkernel__(); + MlirDialectHandle ttnn_handle = mlirGetDialectHandle__ttnn__(); mlirDialectHandleRegisterDialect(tt_handle, context); mlirDialectHandleRegisterDialect(ttir_handle, context); mlirDialectHandleRegisterDialect(ttkernel_handle, context); + mlirDialectHandleRegisterDialect(ttnn_handle, context); if (load) { mlirDialectHandleLoadDialect(tt_handle, context); mlirDialectHandleLoadDialect(ttir_handle, context); mlirDialectHandleLoadDialect(ttkernel_handle, context); + mlirDialectHandleLoadDialect(ttnn_handle, context); } }, py::arg("context"), py::arg("load") = true); @@ -28,6 +31,8 @@ PYBIND11_MODULE(_ttmlir, m) { mlir::ttmlir::python::populateTTModule(tt_ir); auto ttkernel_ir = m.def_submodule("ttkernel_ir", "TTKernel IR Bindings"); mlir::ttmlir::python::populateTTKernelModule(ttkernel_ir); + auto ttnn_ir = m.def_submodule("ttnn_ir", "TTNN IR Bindings"); + mlir::ttmlir::python::populateTTNNModule(ttnn_ir); auto overrides = m.def_submodule("overrides", "Python-Bound Overrides"); mlir::ttmlir::python::populateOverridesModule(overrides); auto passes = diff --git a/python/TTNNModule.cpp b/python/TTNNModule.cpp new file mode 100644 index 0000000000..7a55c2d0a1 --- /dev/null +++ b/python/TTNNModule.cpp @@ -0,0 +1,104 @@ +// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +// +// SPDX-License-Identifier: Apache-2.0 + +#include "ttmlir/Bindings/Python/TTMLIRModule.h" + +namespace mlir::ttmlir::python { +void populateTTNNModule(py::module &m) { + + py::class_(m, "CoreRangeAttr") + .def_static("get", + [](MlirContext ctx, std::vector offset, + std::vector size) { + return wrap(tt::ttnn::CoreRangeAttr::get(unwrap(ctx), + offset, size)); + }) + .def_static( + "get_with_grid", + [](MlirContext ctx, MlirAttribute grid, std::vector offset) { + llvm::SmallVector offsetVec{0, 0}; + if (offset.size() == 2 && not(offset[0] == 0 && offset[1] == 0)) { + offsetVec[0] = offset[0]; + offsetVec[1] = offset[1]; + } + return wrap(tt::ttnn::CoreRangeAttr::get( + unwrap(ctx), mlir::cast(unwrap(grid)), + offsetVec)); + }, + py::arg("ctx"), py::arg("grid"), + py::arg("offset") = std::vector{0, 0}); + py::class_(m, "LayoutAttr") + .def_static("get", + [](MlirContext ctx, uint32_t layout) { + return wrap(tt::ttnn::LayoutAttr::get( + unwrap(ctx), static_cast(layout))); + }) + .def_property_readonly("value", [](tt::ttnn::LayoutAttr self) { + return static_cast(self.getValue()); + }); + py::class_(m, "TensorMemoryLayoutAttr") + .def_static("get", + [](MlirContext ctx, uint32_t tensorMemoryLayout) { + return wrap(tt::ttnn::TensorMemoryLayoutAttr::get( + unwrap(ctx), static_cast( + tensorMemoryLayout))); + }) + .def_property_readonly("value", + [](tt::ttnn::TensorMemoryLayoutAttr self) { + return static_cast(self.getValue()); + }); + py::class_(m, "BufferTypeAttr") + .def_static( + "get", + [](MlirContext ctx, uint32_t bufferType) { + return wrap(tt::ttnn::BufferTypeAttr::get( + unwrap(ctx), static_cast(bufferType))); + }) + .def_property_readonly("value", [](tt::ttnn::BufferTypeAttr self) { + return static_cast(self.getValue()); + }); + py::class_(m, "MemoryConfigAttr") + .def_static("get", + [](MlirContext ctx, + tt::ttnn::TensorMemoryLayoutAttr tensorMemoryLayoutAttr, + tt::ttnn::BufferTypeAttr bufferTypeAttr) { + return wrap(tt::ttnn::MemoryConfigAttr::get( + unwrap(ctx), tensorMemoryLayoutAttr, bufferTypeAttr)); + }) + .def_static( + "get_by_value", + [](MlirContext ctx, uint32_t tensorMemoryLayout, + uint32_t bufferType) { + return wrap(tt::ttnn::MemoryConfigAttr::get( + unwrap(ctx), + tt::ttnn::TensorMemoryLayoutAttr::get( + unwrap(ctx), static_cast( + tensorMemoryLayout)), + tt::ttnn::BufferTypeAttr::get( + unwrap(ctx), + static_cast(bufferType)))); + }) + .def_property_readonly("tensor_memory_layout", + &tt::ttnn::MemoryConfigAttr::getTensorMemoryLayout) + .def_property_readonly("buffer_type", + &tt::ttnn::MemoryConfigAttr::getBufferType); + py::class_(m, "ShapeAttr") + .def_static("get", + [](MlirContext ctx, std::vector shape) { + return wrap(tt::ttnn::ShapeAttr::get(unwrap(ctx), shape)); + }) + .def_property_readonly("shape", [](tt::ttnn::ShapeAttr self) { + return std::vector(self.getShape().begin(), + self.getShape().end()); + }); + py::class_(m, "MeshShapeAttr") + .def_static("get", + [](MlirContext ctx, int64_t y, int64_t x) { + return wrap( + tt::ttnn::MeshShapeAttr::get(unwrap(ctx), y, x)); + }) + .def_property_readonly("y", &tt::ttnn::MeshShapeAttr::getY) + .def_property_readonly("x", &tt::ttnn::MeshShapeAttr::getX); +} +} // namespace mlir::ttmlir::python diff --git a/python/ttmlir/dialects/TTEnumBindings.td b/python/ttmlir/dialects/TTEnumBinding.td similarity index 100% rename from python/ttmlir/dialects/TTEnumBindings.td rename to python/ttmlir/dialects/TTEnumBinding.td diff --git a/python/ttmlir/dialects/TTNNBinding.td b/python/ttmlir/dialects/TTNNBinding.td new file mode 100644 index 0000000000..ece1633533 --- /dev/null +++ b/python/ttmlir/dialects/TTNNBinding.td @@ -0,0 +1,10 @@ +// SPDX-FileCopyrightText: © 2024 Tenstorrent Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#ifndef PYTHON_BINDINGS_TTMLIR_TTNNOPS +#define PYTHON_BINDINGS_TTMLIR_TTNNOPS + +include "ttmlir/Dialect/TTNN/IR/TTNNOps.td" + +#endif diff --git a/python/ttmlir/dialects/TTNNEnumBinding.td b/python/ttmlir/dialects/TTNNEnumBinding.td new file mode 100644 index 0000000000..3a482b3f0b --- /dev/null +++ b/python/ttmlir/dialects/TTNNEnumBinding.td @@ -0,0 +1,10 @@ +// SPDX-FileCopyrightText: © 2024 Tenstorrent Inc. +// +// SPDX-License-Identifier: Apache-2.0 + +#ifndef PYTHON_BINDINGS_TTMLIR_TTNNENUMS +#define PYTHON_BINDINGS_TTMLIR_TTNNENUMS + +include "ttmlir/Dialect/TTNN/IR/TTNNOpsEnums.td" + +#endif diff --git a/python/ttmlir/dialects/ttnn.py b/python/ttmlir/dialects/ttnn.py new file mode 100644 index 0000000000..659938cf66 --- /dev/null +++ b/python/ttmlir/dialects/ttnn.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +# +# SPDX-License-Identifier: Apache-2.0 + +from ._ttnn_ops_gen import * +from ._ttnn_enum_gen import * +from .._mlir_libs._ttmlir import register_dialect, ttnn_ir as ir