From 737564bb397a2d6a1127fab4cf5c47b479dfe128 Mon Sep 17 00:00:00 2001 From: Tapasvi Patel Date: Wed, 30 Oct 2024 16:36:26 +0000 Subject: [PATCH] #1108: Added ability to dump golden data into flatbuffer file for intermediate ops and output using python infra for generating ttir modules --- include/ttmlir/Target/Common/debug_info.fbs | 20 ++ .../Target/TTMetal/TTMetalToFlatbuffer.h | 6 +- include/ttmlir/Target/TTMetal/binary.fbs | 2 + include/ttmlir/Target/TTNN/TTNNToFlatbuffer.h | 9 +- .../ttmlir/Target/Utils/MLIRToFlatbuffer.h | 25 ++ lib/Target/TTMetal/TTMetalToFlatbuffer.cpp | 10 +- .../TTMetalToFlatbufferRegistration.cpp | 4 +- lib/Target/TTNN/TTNNToFlatbuffer.cpp | 27 +- .../TTNN/TTNNToFlatbufferRegistration.cpp | 5 +- python/Passes.cpp | 33 ++- python/test_infra/test_ttir_ops_ttmetal.py | 12 +- python/test_infra/test_ttir_ops_ttnn.py | 50 ++-- python/test_infra/test_utils.py | 235 +++-------------- python/test_infra/ttir_builder.py | 236 ++++-------------- python/ttmlir/dialects/ttnn.py | 1 - runtime/tools/python/ttrt/common/util.py | 31 +++ 16 files changed, 265 insertions(+), 441 deletions(-) diff --git a/include/ttmlir/Target/Common/debug_info.fbs b/include/ttmlir/Target/Common/debug_info.fbs index be3b8e28e9..468fb24f1a 100644 --- a/include/ttmlir/Target/Common/debug_info.fbs +++ b/include/ttmlir/Target/Common/debug_info.fbs @@ -1,5 +1,24 @@ +include "Common/types.fbs"; + namespace tt.target; +table GoldenTensor { + name: string; + shape: [int64]; + stride: [int64]; + dtype: tt.target.DataType; + data: [uint8]; +} + +table GoldenKV { + key: string; + value: GoldenTensor; +} + +table GoldenInfo { + golden_map: [GoldenKV]; +} + table MLIR { name: string; source: string; @@ -8,4 +27,5 @@ table MLIR { table DebugInfo { mlir: MLIR; cpp: string; + golden_info: GoldenInfo; } diff --git a/include/ttmlir/Target/TTMetal/TTMetalToFlatbuffer.h b/include/ttmlir/Target/TTMetal/TTMetalToFlatbuffer.h index 33908ebd8c..a6f8b80855 100644 --- a/include/ttmlir/Target/TTMetal/TTMetalToFlatbuffer.h +++ b/include/ttmlir/Target/TTMetal/TTMetalToFlatbuffer.h @@ -7,13 +7,15 @@ #include "mlir/IR/Operation.h" #include "mlir/Support/LogicalResult.h" +#include "ttmlir/Target/Utils/MLIRToFlatbuffer.h" namespace mlir::tt::ttmetal { // Translates a TTMetal operation to a flatbuffer and writes it to the given // stream. -LogicalResult translateTTMetalToFlatbuffer(Operation *op, - llvm::raw_ostream &os); +LogicalResult translateTTMetalToFlatbuffer( + Operation *op, llvm::raw_ostream &os, + std::unordered_map goldenMap = {}); } // namespace mlir::tt::ttmetal #endif diff --git a/include/ttmlir/Target/TTMetal/binary.fbs b/include/ttmlir/Target/TTMetal/binary.fbs index c3ca5bdda6..47a41407de 100644 --- a/include/ttmlir/Target/TTMetal/binary.fbs +++ b/include/ttmlir/Target/TTMetal/binary.fbs @@ -1,5 +1,6 @@ include "Common/types.fbs"; include "Common/version.fbs"; +include "Common/debug_info.fbs"; include "command.fbs"; namespace tt.target.metal; @@ -15,6 +16,7 @@ table Program { inputs: [TensorRef]; outputs: [TensorRef]; device_programs: [DeviceProgram]; + debug_info: DebugInfo; } table TTMetalBinary { diff --git a/include/ttmlir/Target/TTNN/TTNNToFlatbuffer.h b/include/ttmlir/Target/TTNN/TTNNToFlatbuffer.h index bc2009dd44..77cc45cfa5 100644 --- a/include/ttmlir/Target/TTNN/TTNNToFlatbuffer.h +++ b/include/ttmlir/Target/TTNN/TTNNToFlatbuffer.h @@ -7,16 +7,21 @@ #include "mlir/IR/Operation.h" #include "mlir/Support/LogicalResult.h" +#include "ttmlir/Target/Utils/MLIRToFlatbuffer.h" namespace mlir::tt::ttnn { // Convert a TTNNIR operation to a flatbuffer -std::shared_ptr ttnnToFlatbuffer(Operation *op); +std::shared_ptr +ttnnToFlatbuffer(Operation *op, + std::unordered_map goldenMap = {}); // Convert a TTNNIR operation to a flatbuffer // This function signature is required in order to register the conversion in // mlir translation framework -LogicalResult translateTTNNToFlatbuffer(Operation *op, llvm::raw_ostream &os); +LogicalResult translateTTNNToFlatbuffer( + Operation *op, llvm::raw_ostream &os, + std::unordered_map goldenMap = {}); } // namespace mlir::tt::ttnn #endif diff --git a/include/ttmlir/Target/Utils/MLIRToFlatbuffer.h b/include/ttmlir/Target/Utils/MLIRToFlatbuffer.h index 7f6356196d..de1362676f 100644 --- a/include/ttmlir/Target/Utils/MLIRToFlatbuffer.h +++ b/include/ttmlir/Target/Utils/MLIRToFlatbuffer.h @@ -5,6 +5,7 @@ #ifndef TTMLIR_TARGET_UTILS_MLIRTOFLATBUFFER_H #define TTMLIR_TARGET_UTILS_MLIRTOFLATBUFFER_H +#include #include #include "flatbuffers/flatbuffers.h" @@ -15,6 +16,30 @@ #include "ttmlir/Utils.h" namespace mlir::tt { +struct GoldenTensor { + std::string name; + std::vector shape; + std::vector strides; + ::tt::target::DataType dtype; + std::uint8_t *data; + + GoldenTensor(std::string name, std::vector shape, + std::vector strides, ::tt::target::DataType dtype, + std::uint8_t *data) + : name(name), shape(shape), strides(strides), dtype(dtype), data(data) {} + + std::vector convertDataToVector() { + int totalDataSize = std::accumulate(this->shape.begin(), this->shape.end(), + 1, std::multiplies()) * + sizeof(float); + + std::vector dataVec(totalDataSize); + std::memcpy(dataVec.data(), this->data, totalDataSize); + + return dataVec; + } +}; + inline ::tt::target::OOBVal toFlatbuffer(FlatbufferObjectCache &, OOBVal oobVal) { switch (oobVal) { diff --git a/lib/Target/TTMetal/TTMetalToFlatbuffer.cpp b/lib/Target/TTMetal/TTMetalToFlatbuffer.cpp index 33d033a4f1..fcdd04f061 100644 --- a/lib/Target/TTMetal/TTMetalToFlatbuffer.cpp +++ b/lib/Target/TTMetal/TTMetalToFlatbuffer.cpp @@ -198,7 +198,8 @@ Value getOperandThroughDPSOps(Value value) { return value; } -static std::shared_ptr translateModuleToFlatbuffer(Operation *op) { +static std::shared_ptr translateModuleToFlatbuffer( + Operation *op, std::unordered_map goldenMap) { ::flatbuffers::FlatBufferBuilder fbb; FlatbufferObjectCache cache(&fbb); @@ -372,9 +373,10 @@ static std::shared_ptr translateModuleToFlatbuffer(Operation *op) { return serializedBinary; } -LogicalResult translateTTMetalToFlatbuffer(Operation *op, - llvm::raw_ostream &os) { - std::shared_ptr data = translateModuleToFlatbuffer(op); +LogicalResult translateTTMetalToFlatbuffer( + Operation *op, llvm::raw_ostream &os, + std::unordered_map goldenMap) { + std::shared_ptr data = translateModuleToFlatbuffer(op, goldenMap); std::size_t size = ::flatbuffers::GetSizePrefixedBufferLength( static_cast(data.get())); os.write(reinterpret_cast(data.get()), size); diff --git a/lib/Target/TTMetal/TTMetalToFlatbufferRegistration.cpp b/lib/Target/TTMetal/TTMetalToFlatbufferRegistration.cpp index b713da66aa..ea2ccad4a9 100644 --- a/lib/Target/TTMetal/TTMetalToFlatbufferRegistration.cpp +++ b/lib/Target/TTMetal/TTMetalToFlatbufferRegistration.cpp @@ -18,7 +18,9 @@ namespace mlir::tt::ttmetal { void registerTTMetalToFlatbuffer() { TranslateFromMLIRRegistration reg( "ttmetal-to-flatbuffer", "translate ttmetal dialect to flatbuffer", - translateTTMetalToFlatbuffer /* function */, + [](Operation *op, llvm::raw_ostream &os) -> LogicalResult { + return translateTTMetalToFlatbuffer(op, os, {}); + }, [](DialectRegistry ®istry) { registry.insert ttnnToFlatbuffer(Operation *op) { +std::shared_ptr +ttnnToFlatbuffer(Operation *op, + std::unordered_map goldenMap) { ModuleOp module = dyn_cast(op); assert(module && "Expected ModuleOp as top level operation"); @@ -728,7 +730,22 @@ std::shared_ptr ttnnToFlatbuffer(Operation *op) { auto result = mlir::tt::ttnn::emitTTNNAsCpp(module, os); (void)result; - auto debugInfo = ::tt::target::CreateDebugInfoDirect(fbb, mlir, cpp.c_str()); + std::vector<::flatbuffers::Offset<::tt::target::GoldenKV>> goldenKVList; + goldenKVList.reserve(goldenMap.size()); + + for (auto element : goldenMap) { + std::vector dataTensor = element.second.convertDataToVector(); + auto goldenTensor = ::tt::target::CreateGoldenTensorDirect( + fbb, element.second.name.c_str(), &element.second.shape, + &element.second.strides, element.second.dtype, &dataTensor); + auto goldenKV = ::tt::target::CreateGoldenKVDirect( + fbb, element.first.c_str(), goldenTensor); + goldenKVList.push_back(goldenKV); + } + + auto goldenInfo = ::tt::target::CreateGoldenInfoDirect(fbb, &goldenKVList); + auto debugInfo = + ::tt::target::CreateDebugInfoDirect(fbb, mlir, cpp.c_str(), goldenInfo); std::vector<::flatbuffers::Offset<::tt::target::ttnn::Program>> programs; module->walk([&](func::FuncOp func) { @@ -756,8 +773,10 @@ std::shared_ptr ttnnToFlatbuffer(Operation *op) { return bufferPtr; } -LogicalResult translateTTNNToFlatbuffer(Operation *op, llvm::raw_ostream &os) { - std::shared_ptr data = ttnnToFlatbuffer(op); +LogicalResult translateTTNNToFlatbuffer( + Operation *op, llvm::raw_ostream &os, + std::unordered_map goldenMap) { + std::shared_ptr data = ttnnToFlatbuffer(op, goldenMap); std::size_t size = ::flatbuffers::GetSizePrefixedBufferLength( static_cast(data.get())); os.write(reinterpret_cast(data.get()), size); diff --git a/lib/Target/TTNN/TTNNToFlatbufferRegistration.cpp b/lib/Target/TTNN/TTNNToFlatbufferRegistration.cpp index f572728ff8..243c2b7e37 100644 --- a/lib/Target/TTNN/TTNNToFlatbufferRegistration.cpp +++ b/lib/Target/TTNN/TTNNToFlatbufferRegistration.cpp @@ -18,7 +18,10 @@ namespace mlir::tt::ttnn { void registerTTNNToFlatbuffer() { TranslateFromMLIRRegistration reg( "ttnn-to-flatbuffer", "translate ttnn to flatbuffer", - translateTTNNToFlatbuffer /* function */, [](DialectRegistry ®istry) { + [](Operation *op, llvm::raw_ostream &os) -> LogicalResult { + return translateTTNNToFlatbuffer(op, os, {}); + }, + [](DialectRegistry ®istry) { // clang-format off registry.insert goldenMap) { mlir::Operation *moduleOp = unwrap(mlirModuleGetOperation(module)); std::error_code fileError; @@ -178,15 +179,16 @@ void populatePassesModule(py::module &m) { ". Error: " + fileError.message()); } - if (mlir::failed( - mlir::tt::ttnn::translateTTNNToFlatbuffer(moduleOp, file))) { + if (mlir::failed(mlir::tt::ttnn::translateTTNNToFlatbuffer( + moduleOp, file, goldenMap))) { throw std::runtime_error("Failed to write flatbuffer to file: " + filepath); } }); m.def("ttmetal_to_flatbuffer_file", - [](MlirModule module, std::string &filepath) { + [](MlirModule module, std::string &filepath, + std::unordered_map goldenMap) { mlir::Operation *moduleOp = unwrap(mlirModuleGetOperation(module)); std::error_code fileError; llvm::raw_fd_ostream file(filepath, fileError); @@ -195,11 +197,32 @@ void populatePassesModule(py::module &m) { ". Error: " + fileError.message()); } if (mlir::failed(mlir::tt::ttmetal::translateTTMetalToFlatbuffer( - moduleOp, file))) { + moduleOp, file, goldenMap))) { throw std::runtime_error("Failed to write flatbuffer to file: " + filepath); } }); + + py::enum_<::tt::target::DataType>(m, "DataType") + .value("Float32", ::tt::target::DataType::Float32) + .value("Float16", ::tt::target::DataType::Float16); + + py::class_(m, "GoldenTensor") + .def(py::init, std::vector, + ::tt::target::DataType, std::uint8_t *>()) + .def_readwrite("name", &mlir::tt::GoldenTensor::name) + .def_readwrite("shape", &mlir::tt::GoldenTensor::shape) + .def_readwrite("strides", &mlir::tt::GoldenTensor::strides) + .def_readwrite("dtype", &mlir::tt::GoldenTensor::dtype) + .def_readwrite("data", &mlir::tt::GoldenTensor::data); + + m.def("create_golden_tensor", + [](std::string name, std::vector shape, + std::vector strides, ::tt::target::DataType dtype, + std::uintptr_t ptr) { + return mlir::tt::GoldenTensor(name, shape, strides, dtype, + reinterpret_cast(ptr)); + }); } } // namespace mlir::ttmlir::python diff --git a/python/test_infra/test_ttir_ops_ttmetal.py b/python/test_infra/test_ttir_ops_ttmetal.py index c166c3519a..83c56c1573 100644 --- a/python/test_infra/test_ttir_ops_ttmetal.py +++ b/python/test_infra/test_ttir_ops_ttmetal.py @@ -9,9 +9,7 @@ from ttmlir.test_utils import ( compile_as_mlir_module, - translate_ttnn_to_flatbuffer, - ttir_to_ttnn, - translate_ttmetal_to_flatbuffer, + ttmetal_to_flatbuffer, ttir_to_ttmetal, ) from ttmlir.ttir_builder import Operand, TTIRBuilder @@ -19,7 +17,7 @@ system_desc_path = os.getenv("SYSTEM_DESC_PATH", "") -@translate_ttmetal_to_flatbuffer(output_file_name="test_exp.ttm") +@ttmetal_to_flatbuffer(output_file_name="test_exp.ttm") @ttir_to_ttmetal( output_file_name="test_exp.mlir", system_desc_path=f"{system_desc_path}" ) @@ -28,7 +26,7 @@ def test_exp_ttmetal(in0: Operand, builder: TTIRBuilder): return builder.exp(in0) -@translate_ttmetal_to_flatbuffer(output_file_name="test_add.ttm") +@ttmetal_to_flatbuffer(output_file_name="test_add.ttm") @ttir_to_ttmetal( output_file_name="test_add.mlir", system_desc_path=f"{system_desc_path}" ) @@ -37,7 +35,7 @@ def test_add_ttmetal(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.add(in0, in1) -@translate_ttmetal_to_flatbuffer(output_file_name="test_multiply.ttm") +@ttmetal_to_flatbuffer(output_file_name="test_multiply.ttm") @ttir_to_ttmetal( output_file_name="test_multiply.mlir", system_desc_path=f"{system_desc_path}" ) @@ -46,7 +44,7 @@ def test_multiply_ttmetal(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.multiply(in0, in1) -@translate_ttmetal_to_flatbuffer(output_file_name="test_arbitrary_op_chain.ttm") +@ttmetal_to_flatbuffer(output_file_name="test_arbitrary_op_chain.ttm") @ttir_to_ttmetal( output_file_name="test_arbitrary_op_chain.mlir", system_desc_path=f"{system_desc_path}", diff --git a/python/test_infra/test_ttir_ops_ttnn.py b/python/test_infra/test_ttir_ops_ttnn.py index 7e3ba001c0..960d5f07c4 100644 --- a/python/test_infra/test_ttir_ops_ttnn.py +++ b/python/test_infra/test_ttir_ops_ttnn.py @@ -9,31 +9,29 @@ from ttmlir.test_utils import ( compile_as_mlir_module, - translate_ttnn_to_flatbuffer, + ttnn_to_flatbuffer, ttir_to_ttnn, - translate_ttmetal_to_flatbuffer, - ttir_to_ttmetal, ) from ttmlir.ttir_builder import Operand, TTIRBuilder system_desc_path = os.getenv("SYSTEM_DESC_PATH", "") -@translate_ttnn_to_flatbuffer(output_file_name="test_exp.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_exp.ttnn") @ttir_to_ttnn(output_file_name="test_exp.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((128, 128)) def test_exp_ttnn(in0: Operand, builder: TTIRBuilder): return builder.exp(in0) -@translate_ttnn_to_flatbuffer(output_file_name="test_abs.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_abs.ttnn") @ttir_to_ttnn(output_file_name="test_abs.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((128, 128)) def test_abs_ttnn(in0: Operand, builder: TTIRBuilder): return builder.abs(in0) -@translate_ttnn_to_flatbuffer(output_file_name="test_logical_not.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_logical_not.ttnn") @ttir_to_ttnn( output_file_name="test_logical_not.mlir", system_desc_path=f"{system_desc_path}", @@ -43,28 +41,28 @@ def test_logical_not_ttnn(in0: Operand, builder: TTIRBuilder): return builder.logical_not(in0) -@translate_ttnn_to_flatbuffer(output_file_name="test_neg.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_neg.ttnn") @ttir_to_ttnn(output_file_name="test_neg.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((128, 128)) def test_neg_ttnn(in0: Operand, builder: TTIRBuilder): return builder.neg(in0) -@translate_ttnn_to_flatbuffer(output_file_name="test_relu.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_relu.ttnn") @ttir_to_ttnn(output_file_name="test_relu.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((128, 128)) def test_relu_ttnn(in0: Operand, builder: TTIRBuilder): return builder.relu(in0) -@translate_ttnn_to_flatbuffer(output_file_name="test_sqrt.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_sqrt.ttnn") @ttir_to_ttnn(output_file_name="test_sqrt.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((128, 128)) def test_sqrt_ttnn(in0: Operand, builder: TTIRBuilder): return builder.sqrt(in0) -@translate_ttnn_to_flatbuffer(output_file_name="test_rsqrt.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_rsqrt.ttnn") @ttir_to_ttnn( output_file_name="test_rsqrt.mlir", system_desc_path=f"{system_desc_path}" ) @@ -73,7 +71,7 @@ def test_rsqrt_ttnn(in0: Operand, builder: TTIRBuilder): return builder.rsqrt(in0) -@translate_ttnn_to_flatbuffer(output_file_name="test_sigmoid.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_sigmoid.ttnn") @ttir_to_ttnn( output_file_name="test_sigmoid.mlir", system_desc_path=f"{system_desc_path}" ) @@ -82,7 +80,7 @@ def test_sigmoid_ttnn(in0: Operand, builder: TTIRBuilder): return builder.sigmoid(in0) -@translate_ttnn_to_flatbuffer(output_file_name="test_reciprocal.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_reciprocal.ttnn") @ttir_to_ttnn( output_file_name="test_reciprocal.mlir", system_desc_path=f"{system_desc_path}" ) @@ -91,14 +89,14 @@ def test_reciprocal_ttnn(in0: Operand, builder: TTIRBuilder): return builder.reciprocal(in0) -@translate_ttnn_to_flatbuffer(output_file_name="test_add.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_add.ttnn") @ttir_to_ttnn(output_file_name="test_add.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((64, 128), (64, 128)) def test_add_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.add(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_multiply.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_multiply.ttnn") @ttir_to_ttnn( output_file_name="test_multiply.mlir", system_desc_path=f"{system_desc_path}" ) @@ -107,7 +105,7 @@ def test_multiply_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.multiply(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_logical_and.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_logical_and.ttnn") @ttir_to_ttnn( output_file_name="test_logical_and.mlir", system_desc_path=f"{system_desc_path}", @@ -117,7 +115,7 @@ def test_logical_and_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.logical_and(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_logical_or.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_logical_or.ttnn") @ttir_to_ttnn( output_file_name="test_logical_or.mlir", system_desc_path=f"{system_desc_path}" ) @@ -126,7 +124,7 @@ def test_logical_or_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.logical_or(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_subtract.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_subtract.ttnn") @ttir_to_ttnn( output_file_name="test_subtract.mlir", system_desc_path=f"{system_desc_path}" ) @@ -135,56 +133,56 @@ def test_subtract_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.subtract(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_eq.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_eq.ttnn") @ttir_to_ttnn(output_file_name="test_eq.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((64, 64), (64, 64)) def test_eq_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.eq(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_ne.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_ne.ttnn") @ttir_to_ttnn(output_file_name="test_ne.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((64, 64), (64, 64)) def test_ne_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.ne(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_ge.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_ge.ttnn") @ttir_to_ttnn(output_file_name="test_ge.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((64, 64), (64, 64)) def test_ge_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.ge(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_gt.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_gt.ttnn") @ttir_to_ttnn(output_file_name="test_gt.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((64, 64), (64, 64)) def test_gt_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.gt(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_le.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_le.ttnn") @ttir_to_ttnn(output_file_name="test_le.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((64, 64), (64, 64)) def test_le_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.le(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_lt.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_lt.ttnn") @ttir_to_ttnn(output_file_name="test_lt.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((64, 64), (64, 64)) def test_lt_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.lt(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_div.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_div.ttnn") @ttir_to_ttnn(output_file_name="test_div.mlir", system_desc_path=f"{system_desc_path}") @compile_as_mlir_module((64, 64), (64, 64)) def test_div_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.div(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_maximum.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_maximum.ttnn") @ttir_to_ttnn( output_file_name="test_maximum.mlir", system_desc_path=f"{system_desc_path}" ) @@ -193,7 +191,7 @@ def test_maximum_ttnn(in0: Operand, in1: Operand, builder: TTIRBuilder): return builder.maximum(in0, in1) -@translate_ttnn_to_flatbuffer(output_file_name="test_arbitrary_op_chain.ttnn") +@ttnn_to_flatbuffer(output_file_name="test_arbitrary_op_chain.ttnn") @ttir_to_ttnn( output_file_name="test_arbitrary_op_chain.mlir", system_desc_path=f"{system_desc_path}", diff --git a/python/test_infra/test_utils.py b/python/test_infra/test_utils.py index b5e6f36aa6..d99a3a17a9 100644 --- a/python/test_infra/test_utils.py +++ b/python/test_infra/test_utils.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 import os -from typing import Callable, Dict, Tuple +from typing import Callable, Dict, Tuple, List import torch from ttmlir.dialects import func @@ -28,64 +28,6 @@ def _dump_module(module: Module) -> None: print(module) -def _run_ttmlir_translate_ttmetal( - input_file_name: str, output_file_name: str = "ttmetal_fb.ttm" -): - """ - Util function running `ttmlir-translate` tool on a file containing dumped TTMetal - module. It produces flatbuffer file `output_file_name`. - """ - import subprocess - - res = subprocess.run( - " ".join( - [ - f"ttmlir-translate", - "--ttmetal-to-flatbuffer", - input_file_name, - "-o", - output_file_name, - ] - ), - shell=True, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - ) - assert ( - res.returncode == 0 - ), f"Running ttmlir-translate failed with: {res.stdout.decode('utf-8')}" - return res - - -def _run_ttmlir_translate_ttnn( - input_file_name: str, output_file_name: str = "ttnn_fb.ttnn" -): - """ - Util function running `ttmlir-translate` tool on a file containing dumped TTNN - module. It produces flatbuffer file `output_file_name`. - """ - import subprocess - - res = subprocess.run( - " ".join( - [ - f"ttmlir-translate", - "--ttnn-to-flatbuffer", - input_file_name, - "-o", - output_file_name, - ] - ), - shell=True, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - ) - assert ( - res.returncode == 0 - ), f"Running ttmlir-translate failed with: {res.stdout.decode('utf-8')}" - return res - - # ----- Decorators for doing passes and compiling to flatbuffer ----- @@ -171,8 +113,8 @@ def wrapper(): @func.func(*test_fn_input_types, name=test_fn.__name__) def decorated_func(*inputs): # Randomly generate golden tensors for function inputs. - for i in inputs: - builder.generate_and_store_random_golden(i) + for index, i in enumerate(inputs): + builder.generate_input_golden(i, index) return test_fn(*inputs, builder=builder) @@ -183,7 +125,7 @@ def decorated_func(*inputs): if module_dump: _dump_module(module) - return module + return module, builder return wrapper @@ -207,17 +149,16 @@ def ttir_to_ttnn( output_file_name: str Name of the output file. + + Returns + ------- + MLIR module containing MLIR op graph defined by decorated test function and instance of TTIRBuilder. """ def decorator(fn: Callable): def wrapper(*args, **kwargs): - # First, call the decorated function to get the MLIR module. - module = fn(*args, **kwargs) - - assert isinstance(module, Module), ( - f"Make sure this decorator is used on top of " - f"`compile_as_mlir_module` decorator." - ) + # First, call the decorated function to get the MLIR module and builder instance + module, builder = fn(*args, **kwargs) # Now, pass it through the TTIR to TTNN pipeline. Module gets # modified in place. @@ -232,7 +173,7 @@ def wrapper(*args, **kwargs): with open(output_file_name, "w") as f: f.write(str(module)) - return output_file_name + return module, builder return wrapper @@ -242,7 +183,6 @@ def wrapper(*args, **kwargs): def ttir_to_ttmetal( dump_to_file: bool = True, output_file_name: str = "test.mlir", - return_module: bool = False, system_desc_path: str = "", ): """ @@ -258,21 +198,15 @@ def ttir_to_ttmetal( output_file_name: str Name of the output file. - return_module: bool - Flag through which one chooses to return the generated module or name of the - file in which module was dumped (i.e. `output_file_name`). Exists only to - accommodate both `ttmetal_to_flatbuffer` and `translate_ttmetal_to_flatbuffer`. + Returns + ------- + MLIR module containing MLIR op graph defined by decorated test function and instance of TTIRBuilder. """ def decorator(fn: Callable): def wrapper(*args, **kwargs): # First, call the decorated function to get the MLIR module. - module = fn(*args, **kwargs) - - assert isinstance(module, Module), ( - f"Make sure this decorator is used on top of " - f"`compile_as_mlir_module` decorator." - ) + module, builder = fn(*args, **kwargs) # Now, pass it through the TTIR to TTMetal pipeline. Module gets # modified in place. @@ -287,152 +221,61 @@ def wrapper(*args, **kwargs): with open(output_file_name, "w") as f: f.write(str(module)) - return module if return_module else output_file_name + return module, builder return wrapper return decorator -def ttmetal_to_flatbuffer( - output_file_name: str = "ttmetal_fb.ttmg", golden_info: Dict[Operand, Golden] = None +def ttnn_to_flatbuffer( + output_file_name: str = "ttnn_fb.ttnn", ): """ - NOTE NOT WORKING, DO NOT USE. - - Converts TTMetal module to flatbuffer and saves to file, meant to be used as a - decorator on top of `ttir_to_ttmetal` decorator. Take note that `ttir_to_ttmetal` + Converts TTNN module to flatbuffer and saves to file, meant to be used as a + decorator on top of `ttir_to_ttnn` decorator. Take note that `ttir_to_ttnn` has to return module instead of file name if decorated with this decorator. - Wrapper around `ttmetal_to_flatbuffer_file` pybound pass. - - TODO Optional golden info is passed to be embedded in flatbuffer as well. - - TODO Decorating a test function with this, i.e. calling - `ttmetal_to_flatbuffer_file` will result in - - 'LLVM ERROR: Building op `emitc.constant` but it isn't known in this MLIRContext: - the dialect may not be loaded or this operation hasn't been added by the dialect.' - - To circumvent this, `ttmlir-translate` is run on file that - `ttir_to_ttmetal_backend_pipeline` produces to generate TTMetal flatbuffer file, - which this decorator was supposed to generate. Use `translate_ttmetal_to_flatbuffer` - to achieve this, and make `ttir_to_ttmetal` return file name instead of module. + Wrapper around `ttnn_to_flatbuffer_file` pybound pass. """ def decorator(test_fn: Callable): def wrapper(*args, **kwargs): - # Get the TTMetal module by calling the wrapped function. - module = test_fn(*args, **kwargs) - - assert isinstance(module, Module), ( - f"Make sure `ttir_to_ttmetal` which was decorated with this function " - f"returns module, not file name." - ) + # Get the TTNN module by calling the wrapped function. + module, builder = test_fn(*args, **kwargs) # Convert to flatbuffer file. - ttmetal_to_flatbuffer_file(module, output_file_name) - - print("`ttmetal_to_flatbuffer_file` passed successfully.") - - return wrapper - - return decorator - + ttnn_to_flatbuffer_file(module, output_file_name, builder.get_golden_map()) -def translate_ttmetal_to_flatbuffer(output_file_name: str = "ttmetal_fb.ttm"): - """ - NOTE Substitutes `ttmetal_to_flatbuffer` decorator. - - By running `ttmlir-translate` on input file, it produces TTMetal flatbuffer file - `output_file_name`, meant to be used as a decorator on top of `ttir_to_ttmetal` - decorator. Take note that `ttir_to_ttmetal` has to return file name instead of - module if decorated with this decorator. - - Wrapper around `ttmlir-translate` call. - - Example - ------- - - ```python - @translate_ttmetal_to_flatbuffer(output_file_name="ttmetal_fb_test_add.ttm") - @ttir_to_ttmetal(dump_to_file=True, output_file_name="test_add.mlir", return_module=False) - @compile_as_mlir_module((32, 32), (32, 32)) - def test_add(in0: Operand, in1: Operand, builder: TTIRBuilder): - # CHECK: %0 = tensor.empty() : tensor<32x32xf32> - # CHECK: %1 = "ttir.add"(%arg0, %arg1, %0) - # CHECK: return %1 : tensor<32x32xf32> - - return builder.add(in0, in1) - ``` - """ - - def decorator(fn: Callable): - def wrapper(*args, **kwargs): - input_file_name = fn(*args, **kwargs) - - assert isinstance(input_file_name, str) and os.path.isfile( - input_file_name - ), ( - f"Make sure `ttir_to_ttmetal` which was decorated with this function " - f"returns file name, not module." - ) - - res = _run_ttmlir_translate_ttmetal(input_file_name, output_file_name) - - print( - f"Flatbuffer file for TTMetalBinary {output_file_name} successfully generated." - ) - - return res.returncode + print("`ttnn_to_flatbuffer_file` passed successfully.") return wrapper return decorator -def translate_ttnn_to_flatbuffer(output_file_name: str = "ttnn_fb.ttnn"): +def ttmetal_to_flatbuffer( + output_file_name: str = "ttmetal_fb.ttmg", +): """ + Converts TTMetal module to flatbuffer and saves to file, meant to be used as a + decorator on top of `ttir_to_ttmetal` decorator. Take note that `ttir_to_ttmetal` + has to return module instead of file name if decorated with this decorator. - By running `ttmlir-translate` on input file, it produces TTNN flatbuffer file - `output_file_name`, meant to be used as a decorator on top of `ttir_to_ttnn` - decorator. - - Wrapper around `ttmlir-translate` call. - - Example - ------- - - ```python - @translate_ttnn_to_flatbuffer(output_file_name="ttnn_fb_test_add.ttm") - @ttir_to_ttnn(dump_to_file=True, output_file_name="test_add.mlir") - @compile_as_mlir_module((32, 32), (32, 32)) - def test_add(in0: Operand, in1: Operand, builder: TTIRBuilder): - # CHECK: %0 = tensor.empty() : tensor<32x32xf32> - # CHECK: %1 = "ttir.add"(%arg0, %arg1, %0) - # CHECK: return %1 : tensor<32x32xf32> - - return builder.add(in0, in1) - ``` + Wrapper around `ttmetal_to_flatbuffer_file` pybound pass. """ - def decorator(fn: Callable): + def decorator(test_fn: Callable): def wrapper(*args, **kwargs): - input_file_name = fn(*args, **kwargs) - assert isinstance(input_file_name, str) and os.path.isfile( - input_file_name - ), ( - f"Make sure `ttir_to_ttnn` which was decorated with this function " - f"returns file name, not module." - ) - - res = _run_ttmlir_translate_ttnn(input_file_name, output_file_name) + # Get the TTMetal module by calling the wrapped function. + module, builder = test_fn(*args, **kwargs) - print( - f"Flatbuffer file for TTNNBinary {output_file_name} successfully generated." + # Convert to flatbuffer file. + ttmetal_to_flatbuffer_file( + module, output_file_name, builder.get_golden_map() ) - return res.returncode + print("`ttmetal_to_flatbuffer_file` passed successfully.") return wrapper diff --git a/python/test_infra/ttir_builder.py b/python/test_infra/ttir_builder.py index 122cf576de..8234fb3d1a 100644 --- a/python/test_infra/ttir_builder.py +++ b/python/test_infra/ttir_builder.py @@ -8,6 +8,7 @@ from typing import List, Optional, Union, Tuple, Callable, Dict from ttmlir.ir import * from ttmlir.dialects import ttir, tt, func, tensor +from ttmlir.passes import create_golden_tensor, DataType import torch # Alias for operands of ops which can be either BlockArguments, Values, or other @@ -57,12 +58,22 @@ def __init__(self, ctx: Context, location: Location): # graph. self._goldens: Dict[Operand, Golden] = {} + # global ID of operations + self._global_id = -1 + + # id to golden map + self.id_golden_map = {} + # ----- Public helpers ----- @property def goldens(self) -> Dict: return self._goldens + def get_next_global_id(self) -> int: + self._global_id += 1 + return self._global_id + def print_goldens(self) -> None: """ Prints saved operands and their respective goldens in descriptive form @@ -82,15 +93,39 @@ def get_shape(self, input: Operand) -> Shape: """Retrieves shape of operand which is expected to be a shaped type.""" return self._get_type(input).shape - def generate_and_store_random_golden(self, operand: Operand) -> None: + def generate_and_store_random_golden(self, operand: Operand) -> Golden: """ Generates random tensor of `operand`s shape, assigns it to a golden, and maps `operand` to that golden. + + Returns generated golden. """ seed = self._get_seed() random_tensor = self._generate_random_tensor(self.get_shape(operand), seed) golden = Golden(random_tensor, seed) self._store_golden(operand, golden) + return golden + + def generate_input_golden(self, operand: Operand, index: int) -> None: + """ + Generates random tensor of `input`s shape, assigns it to a golden, + and maps `input` to that golden. + """ + self.id_golden_map[f"input_{index}"] = self.generate_and_store_random_golden( + operand + ) + + def get_golden_map(self) -> Dict: + golden_info = {} + for name, golden_tensor in self.id_golden_map.items(): + golden_info[name] = create_golden_tensor( + name, + list(golden_tensor.tensor.shape), + list(golden_tensor.tensor.stride()), + DataType.Float32, + golden_tensor.tensor.data_ptr(), + ) + return golden_info # ----- Private helpers ----- @@ -236,16 +271,22 @@ def empty( # ----- TTIR op factories ----- def eltwise_proxy( - self, op_golden_function, op_ttir_function, inputs: List[Operand] + self, + op_golden_function: Callable, + op_ttir_function: Callable, + inputs: List[Operand], ) -> OpView: with self._ctx, self._loc: output = self.empty(self.get_shape(inputs[0])) + id = self.get_next_global_id() + op = op_ttir_function( [self._get_type(output)], inputs, [output], self._get_operand_constraint_attr(3), + loc=Location.name(str(id)), ) goldens = [] @@ -253,6 +294,7 @@ def eltwise_proxy( goldens.append(self._get_golden_tensor(input)) golden = Golden(op_golden_function(*goldens)) + self.id_golden_map[str(id)] = golden self._store_golden(op, golden) self._override_golden(output, golden) @@ -323,193 +365,3 @@ def div(self, in0: Operand, in1: Operand) -> OpView: def maximum(self, in0: Operand, in1: Operand) -> OpView: return self.eltwise_proxy(torch.maximum, ttir.MaximumOp, [in0, in1]) - - -def compile_as_mlir_module( - *inputs_shapes: Tuple[Shape], - module_dump: bool = True, -): - """ - Decorator to define a MLIR module specified as a python function. - - It will wrap decorated test function in a MLIR FuncOp wrapped in a MLIR - module, and tie arguments of that FuncOp to test function inputs. It will - also pass a `TTIRBuilder` object as the last argument of test function. - - Arguments - --------- - inputs_shapes: Tuple[Shape] - Shapes of the respective ranked tensor inputs of the test function. - - module_dump: bool - Set to True if printout of generated MLIR module is wished. - - golden_dump: bool - Set to True if printout of generated goldens is wished. - - Example - ------- - - ```python - @compile_as_mlir_module((32, 32), (32, 32)) - def test_add(in0: Operand, in1: Operand, builder: TTIRBuilder): - return builder.add(in0, in1) - - - test_add() # NOTE Called without arguments. - ``` - - which returns - - ``` - #any = #tt.operand_constraint<...> - module { - func.func @test_add( - %arg0: tensor<32x32xf32>, - %arg1: tensor<32x32xf32> - ) -> tensor<32x32xf32> { - %0 = tensor.empty() : tensor<32x32xf32> - %1 = "ttir.add"(%arg0, %arg1, %0) ... - return %1 : tensor<32x32xf32> - } - } - ``` - - Check out: - https://github.com/llvm/llvm-project/blob/main/mlir/test/python/dialects/tensor.py - """ - - def decorator(test_fn: Callable): - # test_fn should be called with no args. - def wrapper(): - ctx = Context() - loc = Location.unknown(ctx) - # Instantiate builder which is passed as the last argument to - # `test_fn` so the user can use it to build ops. - builder = TTIRBuilder(ctx, loc) - - with ctx, loc: - test_fn_input_types = [ - builder.ranked_tensor_type(input_shape) - for input_shape in inputs_shapes - ] - - # Wrap everything in a mlir module. - module = Module.create() - - with InsertionPoint(module.body): - # Wrap everything in a mlir function. - @func.func(*test_fn_input_types, name=test_fn.__name__) - def decorated_func(*inputs): - # Randomly generate golden tensors for function inputs. - for i in inputs: - builder.generate_and_store_random_golden(i) - - return test_fn(*inputs, builder=builder) - - if module_dump: - print(module) - - if golden_dump: - builder.print_goldens() - - return module - - return wrapper - - return decorator - - -def compile_as_mlir_module( - *inputs_shapes: Tuple[Shape], - module_dump: bool = True, -): - """ - Decorator to define a MLIR module specified as a python function. - - It will wrap decorated test function in a MLIR FuncOp wrapped in a MLIR - module, and tie arguments of that FuncOp to test function inputs. It will - also pass a `TTIRBuilder` object as the last argument of test function. - - Arguments - --------- - inputs_shapes: Tuple[Shape] - Shapes of the respective ranked tensor inputs of the test function. - - module_dump: bool - Set to True if printout of generated MLIR module is wished. - - golden_dump: bool - Set to True if printout of generated goldens is wished. - - Example - ------- - - ```python - @compile_as_mlir_module((32, 32), (32, 32)) - def test_add(in0: Operand, in1: Operand, builder: TTIRBuilder): - return builder.add(in0, in1) - - - test_add() # NOTE Called without arguments. - ``` - - which returns - - ``` - #any = #tt.operand_constraint<...> - module { - func.func @test_add( - %arg0: tensor<32x32xf32>, - %arg1: tensor<32x32xf32> - ) -> tensor<32x32xf32> { - %0 = tensor.empty() : tensor<32x32xf32> - %1 = "ttir.add"(%arg0, %arg1, %0) ... - return %1 : tensor<32x32xf32> - } - } - ``` - - Check out: - https://github.com/llvm/llvm-project/blob/main/mlir/test/python/dialects/tensor.py - """ - - def decorator(test_fn: Callable): - # test_fn should be called with no args. - def wrapper(): - ctx = Context() - loc = Location.unknown(ctx) - # Instantiate builder which is passed as the last argument to - # `test_fn` so the user can use it to build ops. - builder = TTIRBuilder(ctx, loc) - - with ctx, loc: - test_fn_input_types = [ - builder.ranked_tensor_type(input_shape) - for input_shape in inputs_shapes - ] - - # Wrap everything in a mlir module. - module = Module.create() - - with InsertionPoint(module.body): - # Wrap everything in a mlir function. - @func.func(*test_fn_input_types, name=test_fn.__name__) - def decorated_func(*inputs): - # Randomly generate golden tensors for function inputs. - for i in inputs: - builder.generate_and_store_random_golden(i) - - return test_fn(*inputs, builder=builder) - - if module_dump: - print(module) - - if golden_dump: - builder.print_goldens() - - return module - - return wrapper - - return decorator diff --git a/python/ttmlir/dialects/ttnn.py b/python/ttmlir/dialects/ttnn.py index 659938cf66..d81f58111a 100644 --- a/python/ttmlir/dialects/ttnn.py +++ b/python/ttmlir/dialects/ttnn.py @@ -3,5 +3,4 @@ # 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 diff --git a/runtime/tools/python/ttrt/common/util.py b/runtime/tools/python/ttrt/common/util.py index 751fa07663..ebbf1d6d72 100644 --- a/runtime/tools/python/ttrt/common/util.py +++ b/runtime/tools/python/ttrt/common/util.py @@ -522,11 +522,28 @@ def get_ttsys_file_extension(): return Flatbuffer.ttsys_file_extension +class Golden: + def __init__(self, tensor_id, tensor_shape, tensor_stride, tensor_data): + self.tensor_id = tensor_id + self.tensor_shape = tensor_shape + self.tensor_stride = tensor_stride + self.tensor_data = tensor_data + + def get_golden_tensor(self): + tensor_byte_data = bytes(self.tensor_data) + float_data = np.frombuffer(tensor_byte_data, dtype=np.float32) + golden_tensor = torch.tensor(float_data, dtype=torch.float32).reshape( + self.tensor_shape + ) + return golden_tensor + + class Binary(Flatbuffer): def __init__(self, logger, file_manager, file_path, capsule=None): super().__init__(logger, file_manager, file_path, capsule=capsule) import ttrt.binary + import torch if not capsule: self.fbb = ttrt.binary.load_binary_from_path(file_path) @@ -540,6 +557,20 @@ def __init__(self, logger, file_manager, file_path, capsule=None): program = Binary.Program(i, self.fbb_dict["programs"][i]) self.programs.append(program) + # populate golden tensors if they exist + if "debug_info" in self.fbb_dict["programs"][i]: + golden_info_list = self.fbb_dict["programs"][i]["debug_info"][ + "golden_info" + ]["golden_map"] + + for golden_tensor_dict in golden_info_list: + Golden( + golden_tensor_dict["key"], + golden_tensor_dict["value"]["shape"], + golden_tensor_dict["value"]["stride"], + golden_tensor_dict["value"]["data"], + ) + def check_system_desc(self, query): import ttrt.binary