DOCSEARCHbyAlgolia
Submit Docs

tt-mlir

repository·main·Indexed 43 days ago

https://github.com/tenstorrent/tt-mlir

An open-source compiler framework built on LLVM's MLIR infrastructure designed to compile and optimize machine learning models for Tenstorrent AI accelerators, specifically targeting Wormhole and Blackhole architectures. It provides custom dialects and transformation passes to translate high-level models (such as StableHLO) into optimized hardware binaries for the TTNN ecosystem.

Tokens
104.5K
Snippets
247
Records
414
Agent score
55%

What's inside tt-mlir

  2. What is tt-mlir?

    main

    tt-mlir is an open-source compiler framework built on top of LLVM's MLIR infrastructure. It is designed to compile and optimize machine learning models for Tenstorrent's custom AI hardware, specifically targeting architectures like Wormhole and Blackhole.

    Key capabilities include:

    • Defining custom MLIR dialects and transformation passes tailored for Tenstorrent architectures.
    • Translating models from front-end dialects (such as StableHLO) into optimized binaries.
    • Targeting the TTNN ecosystem.
    • Supporting a wide range of AI workloads, including training, with primitives designed for multi-chip scaling.

  5. What is ttnn-standalone

    main
    ttnn-standalone is a post-compile tuning and debugging tool. It provides the scaffolding required to run C++ code (generated from compiled ML models like PyTorch, Jax, or ONNX) on device. This allows developers to use generated TTNN library op calls outside of the primary compiler environment.

  6. Module layout and API surface of d2m-jit simulator

    main

    The simulator is exposed through a shadow surface in tools/d2m-jit/sim.py. This module re-exports Layout, dtypes, and config, along with the simulator's host and kernel APIs.

    Public Simulator Surface (_src/sim/__init__.py):

    • Tensors/Blocks: SimTensor, SimBlock, to_layout, empty, zeros, full, tilize, untilize, view, view_layout, permute, to_host, reduction_layout.
    • Kernel/Execution: SimKernel (via run.py), core_index, remote_load, remote_store, Semaphore.
    • Operations: SIM_OPS and SIM_METHODS providing Torch backings for in-kernel syntax.

    Backend Switching: To switch between the device and the simulator, use the config.backend field or the D2M_JIT_BACKEND environment variable.

  7. Understand the limitations of the d2m-jit simulator

    main

    The d2m-jit simulator is designed for functional verification and serves as a fast inner loop and a golden oracle for on-device tests. It is important to note its specific non-goals to avoid misusing it for hardware-specific analysis:

    • Functional, not Performance: It models values, not timing. It does not provide cycle counts, bandwidth estimates, or L1 pressure. For performance modeling, use the device profiler hooks via config.insert_profiler_traces.
    • Sequential, not Concurrent: It does not detect races or deadlocks because it uses sequential execution, which hides real hardware synchronization bugs.
    • Approximate, not Bit-Exact: In default mode, it is not bit-exact to silicon. Even in quirk mode, bit-exactness is only best-effort.
    • Verification, not Replacement: It is intended to be used alongside on-device tests, not as a replacement for them.

  8. Understand the TTMLIR project structure

    main

    The tt-mlir repository is organized into several key functional areas:

    • env: Environment setup for building dependencies like LLVM and Flatbuffers.
    • include/ttmlir: Public headers for the TTMLIR library. This includes:
      • Dialect: MLIR dialect interfaces and definitions (organized into IR, Passes, and Transforms).
      • Target: Flatbuffer schema definitions that define the binary interface between the compiler and the runtime.
    • lib: The core TTMLIR library implementation, including the CAPI (C API used for Python bindings) and Dialect implementations.
    • runtime: The device runtime implementation. It contains public headers in include/tt/runtime, the implementation in lib, and Python bindings in tools/python (where ttrt is implemented).
    • tools/ttmlir-opt: The TTMLIR optimizer driver tool.
    • test: The project test suite.

  9. Explore related Tenstorrent projects

    main

    The tt-mlir ecosystem is part of a broader suite of Tenstorrent tools. Depending on your specific workflow (e.g., XLA integration, ONNX support, or direct Metalium usage), you may need to interact with these related repositories:

    • tt-forge: General purpose forge tools.
    • tt-xla: XLA (Accelerated Linear Algebra) integration.
    • tt-metalium: Based on tt-metal, providing lower-level hardware access.
    • tt-forge-onnx: Specialized ONNX support via forge.
    • tt-tvm: TVM (Tensor Virtual Machine) integration.

  10. Understand d2m-jit kernel capabilities and status

    main

    The d2m-jit kernel-feasibility audit defines what operations can currently be authored as kernels.

    Available Today (✅):

    • Per-tile eltwise: 41 unary and 13 binary operations (via _eltwise_block wrapping d2m.tile_*).
    • Ternary selection: where for mask-and-blend operations.
    • Matmul (@, d2m.matmul): Currently limited to single-tile-per-shard with a zeros-prefilled accumulator.
    • DMA Shortcuts: remote_load and remote_store for per-shard access.
    • Control Flow: scf.for and scf.if within kernel bodies.
    • Synchronization: async, await, and semaphores for multi-thread sync.
    • Views: view, view_layout, and permute (metadata reinterpretation, no data movement).
    • Layouts: tilize, untilize, and to_layout for conversions; zeros, full, and empty for host-side fill.
    • Reductions: reduce_sum, reduce_max, and reduce_mean (per-tile row/col reductions).
    • Broadcasting: tile_bcast (row/col/2d shorthands).

    Limited or Blocked (🟡/🔴):

    • Matmul: Requires zeros-prefilled accumulator; single-tile-per-shard only.
    • In-kernel typecast: Currently host-side only via tilize(dtype=...).
    • Per-tile transpose: Not exposed (use logical permute via views instead).
    • Softmax/LayerNorm: Wide-row implementations are blocked by missing cross-core gather/redistribute ops or accumulator handling.

  11. Explore the tt-mlir tool ecosystem

    main

    The tt-mlir ecosystem includes several specialized components for different stages of the compilation and optimization pipeline:

    • Optimizers: Dedicated logic for improving MLIR representations.
    • PyKernel: Tools related to Python kernel execution/definition.
    • ttnn-jit: Just-In-Time compilation capabilities for TTNN.
    • Python Bindings: Interfaces for using tt-mlir components within Python workflows.
    • Flatbuffers: Support for data serialization using Flatbuffers.

  12. What the TTNN Optimizer does

    main

    The TTNN optimizer optimizes TTNN operations for Tenstorrent hardware by focusing on two primary objectives:

    1. Maximizing L1 memory usage: It attempts to keep intermediate tensors in the fast on-chip L1 memory (SRAM) to avoid high-latency DRAM round-trips.
    2. Optimizing op-specific configurations: It selects the best parameters for operations, such as optimal block sizes for Conv2d or specific activation handling.

    The current implementation uses a greedy, two-pass architecture:

    • Pass 1: GreedyMemoryLayoutPropagation: Enumerates valid tensor layouts for each operation and greedily propagates these configurations through the graph, inserting reshards where necessary.
    • Pass 2: GreedyL1SpillManagement: Manages the L1 memory budget by spilling tensors to DRAM when the L1 capacity is exceeded.

    Note: The previous chain-based optimizer (including TTNNOptimizer, DFShardingPolicy, and ShardSolver) has been removed and is superseded by this pass-based design.