diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index ec2a9e823..7f26b6395 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -19,3 +19,4 @@ /python/ @nsmithtt /runtime/ @jnie-TT @kmabeeTT @AleksKnezevic @pilkicTT /runtime/tools/ @tapspatel +/tools/explorer/ @odjuricicTT @nobradovictt @vprajapati-tt diff --git a/.gitignore b/.gitignore index 76d9ed406..30672f9bf 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,4 @@ third_party/tt-metal .vscode/* .cache *pycache* +*.egg-info diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 3899b3775..e558d0567 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -1,2 +1,3 @@ add_subdirectory(ttmlir-opt) add_subdirectory(ttmlir-translate) +add_subdirectory(explorer) diff --git a/tools/explorer/CMakeLists.txt b/tools/explorer/CMakeLists.txt new file mode 100644 index 000000000..d58802c24 --- /dev/null +++ b/tools/explorer/CMakeLists.txt @@ -0,0 +1,15 @@ +set(TT_EXPLORER_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/run.py) +set(TTMLIR_BUILD_BIN_DIR ${TTMLIR_BINARY_DIR}/bin) + +add_custom_target(explorer + COMMENT "Building tt-explorer... ${TTMLIR_BIN_DIR}" + COMMAND cp -r ${CMAKE_CURRENT_SOURCE_DIR}/* ${CMAKE_CURRENT_BINARY_DIR} + COMMAND pip install ${CMAKE_CURRENT_BINARY_DIR}/tt_adapter + DEPENDS TTMLIRPythonModules +) + +add_custom_command(TARGET explorer POST_BUILD + COMMENT "Installing tt-explorer command..." + COMMAND ${CMAKE_COMMAND} -E copy ${TT_EXPLORER_SCRIPT} ${TTMLIR_BUILD_BIN_DIR}/tt-explorer + COMMAND ${CMAKE_COMMAND} -E echo "Done. Run: tt-explorer to start the server." +) diff --git a/tools/explorer/README.md b/tools/explorer/README.md new file mode 100644 index 000000000..a7fb4136a --- /dev/null +++ b/tools/explorer/README.md @@ -0,0 +1,27 @@ +# TT-explorer + +TT-explorer is a tool for MLIR graph visualization and applying optimizer overrides in order to easily experiment with model performance. + +TODO: add documentation from old tt-explorer repo + +## Build +```bash +source env/activate +cmake --build build -- explorer +``` + +## Usage +Start the server with: +```bash +tt-explorer +``` + +Then open http://localhost:8080 in the browser. + +#### Port Forwarding +P.S. +If using a remote machine make sure to forward the 8080 port. E.g: +```bash +ssh -L 8080:localhost:8080 user@remote-machine +``` +Or set the "Tt › Ird › Reservation: Ports" setting in vscode-ird. diff --git a/tools/explorer/run.py b/tools/explorer/run.py new file mode 100755 index 000000000..6e1f8efc4 --- /dev/null +++ b/tools/explorer/run.py @@ -0,0 +1,12 @@ +#!/usr/bin/env python +# +# # SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +# +# SPDX-License-Identifier: Apache-2.0 + +import model_explorer + +# TODO(odjuricic): Hack to make our extension default for .mlir files. +# This can be handled better when we switch to our model-explorer fork. +model_explorer.extension_manager.ExtensionManager.BUILTIN_ADAPTER_MODULES = [] +model_explorer.visualize(extensions=["tt_adapter"]) diff --git a/tools/explorer/test/models/forward_and_backward.mlir b/tools/explorer/test/models/forward_and_backward.mlir new file mode 100644 index 000000000..2d4769f49 --- /dev/null +++ b/tools/explorer/test/models/forward_and_backward.mlir @@ -0,0 +1,30 @@ +module @SimpleModel attributes {tt.system_desc = #tt.system_desc<[{arch = , grid = 8x8, l1_size = 1499136, num_dram_channels = 12, dram_channel_size = 1073741824, noc_l1_address_align_bytes = 16, pcie_address_align_bytes = 32, noc_dram_address_align_bytes = 32, l1_unreserved_base = 1024, erisc_l1_unreserved_base = 1024, dram_unreserved_base = 1024, dram_unreserved_end = 1073741824, physical_cores = {worker = [ 0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 1x0, 1x1, 1x2, 1x3, 1x4, 1x5, 1x6, 1x7, 2x0, 2x1, 2x2, 2x3, 2x4, 2x5, 2x6, 2x7, 3x0, 3x1, 3x2, 3x3, 3x4, 3x5, 3x6, 3x7, 4x0, 4x1, 4x2, 4x3, 4x4, 4x5, 4x6, 4x7, 5x0, 5x1, 5x2, 5x3, 5x4, 5x5, 5x6, 5x7, 6x0, 6x1, 6x2, 6x3, 6x4, 6x5, 6x6, 6x7, 7x0, 7x1, 7x2, 7x3, 7x4, 7x5, 7x6, 7x7] dram = [ 8x0, 9x0, 10x0, 8x1, 9x1, 10x1, 8x2, 9x2, 10x2, 8x3, 9x3, 10x3]}, supported_data_types = [, , , , , , , , , , , ], supported_tile_sizes = [ 4x16, 16x16, 32x16, 4x32, 16x32, 32x32]}], [0], [3 : i32], [ 0x0x0x0]>} { + func.func @forward(%arg0: tensor<1x784xf32> {ttir.name = "input_1"}, %arg1: tensor<10x784xf32> {ttir.name = "linear.weight"}, %arg2: tensor<10xf32> {ttir.name = "linear.bias"}) -> (tensor<1x10xf32> {ttir.name = "SimpleModel_472.output_softmax_1495"}) { + %0 = tensor.empty() : tensor<784x10xf32> + %1 = "ttir.transpose"(%arg1, %0) <{dim0 = -2 : si32, dim1 = -1 : si32, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint]}> : (tensor<10x784xf32>, tensor<784x10xf32>) -> tensor<784x10xf32> + %2 = tensor.empty() : tensor<1x10xf32> + %3 = "ttir.matmul"(%arg0, %1, %2) <{operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x784xf32>, tensor<784x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> + %4 = tensor.empty() : tensor<1x10xf32> + %5 = "ttir.add"(%3, %arg2, %4) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x10xf32>, tensor<10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> + %6 = tensor.empty() : tensor<1x10xf32> + %7 = "ttir.softmax"(%5, %6) <{dimension = -1 : si32, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> + return %7 : tensor<1x10xf32> + } + func.func @backward(%arg0: tensor<1x10xf32> {ttir.name = "loss_SimpleModel_472.output_softmax_1495"}, %arg1: tensor<1x10xf32> {ttir.name = "SimpleModel_472.output_softmax_1495"}, %arg2: tensor<1x784xf32> {ttir.name = "input_1"}) -> (tensor<1x10xf32> {ttir.name = "grad_acc_linear.bias_grad_accumulator"}, tensor<10x784xf32> {ttir.name = "grad_acc_linear.weight_grad_accumulator"}) { + %0 = tensor.empty() : tensor<1x10xf32> + %1 = "ttir.multiply"(%arg0, %arg1, %0) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x10xf32>, tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> + %2 = tensor.empty() : tensor<1x1xf32> + %3 = "ttir.sum"(%1, %2) <{keep_dim = true, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x10xf32>, tensor<1x1xf32>) -> tensor<1x1xf32> + %4 = tensor.empty() : tensor<1x10xf32> + %5 = "ttir.subtract"(%arg0, %3, %4) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x10xf32>, tensor<1x1xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> + %6 = tensor.empty() : tensor<1x10xf32> + %7 = "ttir.multiply"(%5, %arg1, %6) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x10xf32>, tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32> + %8 = tensor.empty() : tensor<784x1xf32> + %9 = "ttir.transpose"(%arg2, %8) <{dim0 = -2 : si32, dim1 = -1 : si32, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x784xf32>, tensor<784x1xf32>) -> tensor<784x1xf32> + %10 = tensor.empty() : tensor<784x10xf32> + %11 = "ttir.matmul"(%9, %7, %10) <{operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<784x1xf32>, tensor<1x10xf32>, tensor<784x10xf32>) -> tensor<784x10xf32> + %12 = tensor.empty() : tensor<10x784xf32> + %13 = "ttir.transpose"(%11, %12) <{dim0 = -2 : si32, dim1 = -1 : si32, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint]}> : (tensor<784x10xf32>, tensor<10x784xf32>) -> tensor<10x784xf32> + return %7, %13 : tensor<1x10xf32>, tensor<10x784xf32> + } +} diff --git a/tools/explorer/test/models/linear_autoencoder.mlir b/tools/explorer/test/models/linear_autoencoder.mlir new file mode 100644 index 000000000..0a700845d --- /dev/null +++ b/tools/explorer/test/models/linear_autoencoder.mlir @@ -0,0 +1,49 @@ +module @LinearAE attributes {tt.system_desc = #tt.system_desc<[{arch = , grid = 8x8, l1_size = 1499136, num_dram_channels = 12, dram_channel_size = 1073741824, noc_l1_address_align_bytes = 16, pcie_address_align_bytes = 32, noc_dram_address_align_bytes = 32, l1_unreserved_base = 1024, erisc_l1_unreserved_base = 1024, dram_unreserved_base = 1024, dram_unreserved_end = 1073741824, physical_cores = {worker = [ 0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 1x0, 1x1, 1x2, 1x3, 1x4, 1x5, 1x6, 1x7, 2x0, 2x1, 2x2, 2x3, 2x4, 2x5, 2x6, 2x7, 3x0, 3x1, 3x2, 3x3, 3x4, 3x5, 3x6, 3x7, 4x0, 4x1, 4x2, 4x3, 4x4, 4x5, 4x6, 4x7, 5x0, 5x1, 5x2, 5x3, 5x4, 5x5, 5x6, 5x7, 6x0, 6x1, 6x2, 6x3, 6x4, 6x5, 6x6, 6x7, 7x0, 7x1, 7x2, 7x3, 7x4, 7x5, 7x6, 7x7] dram = [ 8x0, 9x0, 10x0, 8x1, 9x1, 10x1, 8x2, 9x2, 10x2, 8x3, 9x3, 10x3]}, supported_data_types = [, , , , , , , , , , , ], supported_tile_sizes = [ 4x16, 16x16, 32x16, 4x32, 16x32, 32x32]}], [0], [3 : i32], [ 0x0x0x0]>} { + func.func @forward(%arg0: tensor<1x784xf32> {ttir.name = "input_1"}, %arg1: tensor<784x128xf32> {ttir.name = "encoder_lin1.weight"}, %arg2: tensor<128xf32> {ttir.name = "encoder_lin1.bias"}, %arg3: tensor<128x64xf32> {ttir.name = "encoder_lin2.weight"}, %arg4: tensor<64xf32> {ttir.name = "encoder_lin2.bias"}, %arg5: tensor<64x12xf32> {ttir.name = "encoder_lin3.weight"}, %arg6: tensor<12xf32> {ttir.name = "encoder_lin3.bias"}, %arg7: tensor<12x3xf32> {ttir.name = "encoder_lin4.weight"}, %arg8: tensor<3xf32> {ttir.name = "encoder_lin4.bias"}, %arg9: tensor<3x12xf32> {ttir.name = "decoder_lin1.weight"}, %arg10: tensor<12xf32> {ttir.name = "decoder_lin1.bias"}, %arg11: tensor<12x64xf32> {ttir.name = "decoder_lin2.weight"}, %arg12: tensor<64xf32> {ttir.name = "decoder_lin2.bias"}, %arg13: tensor<64x128xf32> {ttir.name = "decoder_lin3.weight"}, %arg14: tensor<128xf32> {ttir.name = "decoder_lin3.bias"}, %arg15: tensor<128x784xf32> {ttir.name = "decoder_lin4.weight"}, %arg16: tensor<784xf32> {ttir.name = "decoder_lin4.bias"}) -> (tensor<1x784xf32> {ttir.name = "LinearAE.output_add_29"}) { + %0 = tensor.empty() : tensor<1x128xf32> + %1 = "ttir.matmul"(%arg0, %arg1, %0) <{operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x784xf32>, tensor<784x128xf32>, tensor<1x128xf32>) -> tensor<1x128xf32> + %2 = tensor.empty() : tensor<1x128xf32> + %3 = "ttir.add"(%1, %arg2, %2) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x128xf32>, tensor<128xf32>, tensor<1x128xf32>) -> tensor<1x128xf32> + %4 = tensor.empty() : tensor<1x128xf32> + %5 = "ttir.relu"(%3, %4) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x128xf32>, tensor<1x128xf32>) -> tensor<1x128xf32> + %6 = tensor.empty() : tensor<1x64xf32> + %7 = "ttir.matmul"(%5, %arg3, %6) <{operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x128xf32>, tensor<128x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> + %8 = tensor.empty() : tensor<1x64xf32> + %9 = "ttir.add"(%7, %arg4, %8) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x64xf32>, tensor<64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> + %10 = tensor.empty() : tensor<1x64xf32> + %11 = "ttir.relu"(%9, %10) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> + %12 = tensor.empty() : tensor<1x12xf32> + %13 = "ttir.matmul"(%11, %arg5, %12) <{operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x64xf32>, tensor<64x12xf32>, tensor<1x12xf32>) -> tensor<1x12xf32> + %14 = tensor.empty() : tensor<1x12xf32> + %15 = "ttir.add"(%13, %arg6, %14) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x12xf32>, tensor<12xf32>, tensor<1x12xf32>) -> tensor<1x12xf32> + %16 = tensor.empty() : tensor<1x12xf32> + %17 = "ttir.relu"(%15, %16) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x12xf32>, tensor<1x12xf32>) -> tensor<1x12xf32> + %18 = tensor.empty() : tensor<1x3xf32> + %19 = "ttir.matmul"(%17, %arg7, %18) <{operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x12xf32>, tensor<12x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32> + %20 = tensor.empty() : tensor<1x3xf32> + %21 = "ttir.add"(%19, %arg8, %20) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x3xf32>, tensor<3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32> + %22 = tensor.empty() : tensor<1x12xf32> + %23 = "ttir.matmul"(%21, %arg9, %22) <{operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x3xf32>, tensor<3x12xf32>, tensor<1x12xf32>) -> tensor<1x12xf32> + %24 = tensor.empty() : tensor<1x12xf32> + %25 = "ttir.add"(%23, %arg10, %24) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x12xf32>, tensor<12xf32>, tensor<1x12xf32>) -> tensor<1x12xf32> + %26 = tensor.empty() : tensor<1x12xf32> + %27 = "ttir.relu"(%25, %26) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x12xf32>, tensor<1x12xf32>) -> tensor<1x12xf32> + %28 = tensor.empty() : tensor<1x64xf32> + %29 = "ttir.matmul"(%27, %arg11, %28) <{operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x12xf32>, tensor<12x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> + %30 = tensor.empty() : tensor<1x64xf32> + %31 = "ttir.add"(%29, %arg12, %30) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x64xf32>, tensor<64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> + %32 = tensor.empty() : tensor<1x64xf32> + %33 = "ttir.relu"(%31, %32) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x64xf32>, tensor<1x64xf32>) -> tensor<1x64xf32> + %34 = tensor.empty() : tensor<1x128xf32> + %35 = "ttir.matmul"(%33, %arg13, %34) <{operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x64xf32>, tensor<64x128xf32>, tensor<1x128xf32>) -> tensor<1x128xf32> + %36 = tensor.empty() : tensor<1x128xf32> + %37 = "ttir.add"(%35, %arg14, %36) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x128xf32>, tensor<128xf32>, tensor<1x128xf32>) -> tensor<1x128xf32> + %38 = tensor.empty() : tensor<1x128xf32> + %39 = "ttir.relu"(%37, %38) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x128xf32>, tensor<1x128xf32>) -> tensor<1x128xf32> + %40 = tensor.empty() : tensor<1x784xf32> + %41 = "ttir.matmul"(%39, %arg15, %40) <{operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x128xf32>, tensor<128x784xf32>, tensor<1x784xf32>) -> tensor<1x784xf32> + %42 = tensor.empty() : tensor<1x784xf32> + %43 = "ttir.add"(%41, %arg16, %42) <{operandSegmentSizes = array, operand_constraints = [#tt.operand_constraint, #tt.operand_constraint, #tt.operand_constraint]}> : (tensor<1x784xf32>, tensor<784xf32>, tensor<1x784xf32>) -> tensor<1x784xf32> + return %43 : tensor<1x784xf32> + } +} diff --git a/tools/explorer/test/run_tests.py b/tools/explorer/test/run_tests.py new file mode 100644 index 000000000..d8d0894fa --- /dev/null +++ b/tools/explorer/test/run_tests.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +# +# SPDX-License-Identifier: Apache-2.0 +import model_explorer + +import requests +import time +import multiprocessing +import pytest +import glob + +HOST = "localhost" +PORT = 8002 +CONVERT_URL = "http://" + HOST + ":" + str(PORT) + "/apipost/v1/send_command" +TEST_LOAD_MODEL_PATHS = [ + "test/ttmlir/Dialect/TTNN/mnist_sharding.mlir", + "tools/explorer/test/models/*.mlir", +] + + +def get_test_files(): + files = [] + for path in TEST_LOAD_MODEL_PATHS: + files.extend(glob.glob(path)) + return files + + +@pytest.fixture(scope="session", autouse=True) +def start_server(request): + server_thread = multiprocessing.Process( + target=model_explorer.visualize, + kwargs={"extensions": ["tt_adapter"], "host": HOST, "port": PORT}, + ) + server_thread.start() + time.sleep(1) + + request.addfinalizer(lambda: server_thread.terminate()) + + +@pytest.mark.parametrize("model_path", get_test_files()) +def test_load_model(model_path): + cmd = { + "extensionId": "tt_adapter", + "cmdId": "convert", + "modelPath": model_path, + "deleteAfterConversion": False, + "settings": {}, + } + + result = requests.post(CONVERT_URL, json=cmd) + assert result.ok + if "error" in result.json(): + print(result.json()) + assert False diff --git a/tools/explorer/tt_adapter/README.md b/tools/explorer/tt_adapter/README.md new file mode 100644 index 000000000..6194bcaee --- /dev/null +++ b/tools/explorer/tt_adapter/README.md @@ -0,0 +1,20 @@ +# tt-adapter +Model Explorer Adapter built for TT-MLIR outputs. Contains the logic for converting IRs into model explorer graphs. + +## Integration into model-explorer +Model-Explorer currently primarily supports loading extensions through the CLI. An example of a run call: + +```sh +model-explorer --extensions=tt_adapter +``` + +You should be able to see + +```sh +Loading extensions... + - ... + - Tenstorrent Adapter + - JSON adapter +``` + +in the command output to verify that it has been run. diff --git a/tools/explorer/tt_adapter/pyproject.toml b/tools/explorer/tt_adapter/pyproject.toml new file mode 100644 index 000000000..d69b549c6 --- /dev/null +++ b/tools/explorer/tt_adapter/pyproject.toml @@ -0,0 +1,11 @@ +[project] +name = "tt-adapter" +version = "0.0.1" +description = "Model Explorer Adapter built for TT-MLIR Compiled outputs." +readme = "README.md" +dependencies = [ + "ai-edge-model-explorer>=0.1" +] + +[tool.poetry] +packages = [{ include = "src/tt_adapter" }] diff --git a/tools/explorer/tt_adapter/src/tt_adapter/main.py b/tools/explorer/tt_adapter/src/tt_adapter/main.py new file mode 100644 index 000000000..b88754af6 --- /dev/null +++ b/tools/explorer/tt_adapter/src/tt_adapter/main.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +# +# SPDX-License-Identifier: Apache-2.0 +from typing import Dict +import model_explorer +import ttmlir +from . import ttir + + +class TTAdapter(model_explorer.Adapter): + metadata = model_explorer.AdapterMetadata( + id="tt_adapter", + name="Tenstorrent MLIR Adapter", + description="Adapter for Tenstorrent MLIR dialects used in the Forge compiler.", + source_repo="https://github.com/tenstorrent/tt-mlir/tree/main/tools/explorer/tt_adapter", + fileExts=["mlir", "ttir"], + ) + + # Required. + def __init__(self): + super().__init__() + + def convert( + self, model_path: str, settings: Dict + ) -> model_explorer.ModelExplorerGraphs: + with ttmlir.ir.Context() as ctx, open(model_path, "r") as model_file: + ttmlir.dialects.ttkernel.register_dialect(ctx) + ttmlir.dialects.ttir.register_dialect(ctx) + ttmlir.dialects.tt.register_dialect(ctx) + module = ttmlir.ir.Module.parse("".join(model_file.readlines()), ctx) + + # Convert TTIR to Model Explorer Graphs and Display/Return + graph = ttir.ttir_to_graph(module, ctx) + return {"graphs": [graph]} diff --git a/tools/explorer/tt_adapter/src/tt_adapter/ttir.py b/tools/explorer/tt_adapter/src/tt_adapter/ttir.py new file mode 100644 index 000000000..31522bad5 --- /dev/null +++ b/tools/explorer/tt_adapter/src/tt_adapter/ttir.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC +# +# SPDX-License-Identifier: Apache-2.0 +# Library to manipulate TTIR Modules + +from model_explorer import graph_builder +from ttmlir.dialects import tt, ttir, ttkernel +from collections import defaultdict + + +def get_loc_str(loc): + # TODO(odjuricic) Need to expose this in python bindings, if possible. + try: + res = str(loc).split('"')[1] + except: + res = "unknown" + return res + + +def create_id(op, name_dict): + name = get_loc_str(op.location) + name_num = name_dict[name] + id = name + "__" + str(name_num) + name_dict[name] += 1 + return id + + +def get_attrs(op): + result = [] + for attr in op.attributes: + result.append(graph_builder.KeyValue(key=attr.name, value=str(attr.attr))) + return result + + +def create_namespace(op): + name = get_loc_str(op.location) + if op.parent and op.parent.name != "builtin.module": + return create_namespace(op.parent) + "/" + name + return name + + +def get_layout_attrs(tensor): + attrs = [ + graph_builder.KeyValue(key="shape", value=str(tensor.type.shape)), + graph_builder.KeyValue( + key="element_type", + value=str(tensor.type.element_type), + ), + graph_builder.KeyValue(key="rank", value=str(tensor.type.rank)), + ] + + if hasattr(tensor.type, "encoding") and tensor.type.encoding: + layout = tt.ir.LayoutAttr.getLayout(tensor.type) + attrs.extend( + [ + graph_builder.KeyValue( + key="Memory Space", + value=str(tt.MemorySpace(layout.memory_space_as_int)), + ), + graph_builder.KeyValue( + key="Memory Layout", + value=str(tt.TensorMemoryLayout(layout.memory_layout_as_int)), + ), + graph_builder.KeyValue( + key="Grid Shape", + value=str(list(layout.grid_attr.shape)), + ), + ] + ) + + return attrs + + +def ttir_to_graph(module, ctx): + # Can assume that to-layout pass has already been run on the module. + name_dict = defaultdict(int) + output_connections = defaultdict(int) + graph = graph_builder.Graph(id="ttir-graph") + + op_to_graph_node = dict() + + for op in module.body.operations: + append_later = [] + for region in op.regions: + for block in region.blocks: + for op in block.operations: + # Create all the nodes and constants in the first pass. + graph_node = graph_builder.GraphNode( + id=create_id(op, name_dict), + label=op.name, + namespace=create_namespace(op), + attrs=get_attrs(op), + ) + + if op.name == "tensor.empty": + append_later.append(graph_node) + else: + graph.nodes.append(graph_node) + + op_to_graph_node[op] = graph_node + + for operand in op.operands: + if operand.owner == block and operand not in op_to_graph_node: + # This is a constant and we need to create a node for it. + operand_node = graph_builder.GraphNode( + id=create_id(op, name_dict), + label=operand.get_name(), + namespace=create_namespace(op), + ) + graph.nodes.append(operand_node) + op_to_graph_node[operand] = operand_node + + # This puts the node at the far right when viewing which is a bit more consistant with it being the last operand. + for node in append_later: + graph.nodes.append(node) + + for op in block.operations: + # Create all the edges in the second pass. + for operand_index, operand in enumerate(op.operands): + if operand.owner == block: + source_node = op_to_graph_node[operand] + else: + source_node = op_to_graph_node[operand.owner] + + target_node = op_to_graph_node[op] + + target_node.incomingEdges.append( + graph_builder.IncomingEdge( + sourceNodeId=source_node.id, + sourceNodeOutputId=output_connections[source_node.id], + targetNodeInputId=operand_index, + ) + ) + + output_attrs = get_layout_attrs(operand) + source_node.outputsMetadata.append( + graph_builder.MetadataItem( + id=str(output_connections[source_node.id]), + attrs=[ + graph_builder.KeyValue( + key="__tensor_tag", value=target_node.label + ), + ] + + output_attrs, + ) + ) + output_connections[source_node.id] += 1 + + return graph