[GH-ISSUE #3673] [Feature] ML compiler IR reverse engineering: MLIR bytecode, StableHLO, TOSA and XLA HLO graph reconstruction #1386

New issue

Open

opened 2026-10-16 21:04:57 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-16 21:04:57 +00:00

Owner

Originally created by @matbrgz on GitHub (Oct 14, 2026).
Original GitHub issue: https://github.com/morluto/rea/issues/3673

Summary

Machine learning compilers (such as Google XLA, IREE, TVM, PyTorch Inductor, and ONNX Runtime) compile high-level neural network graphs into intermediate representations (IR) using MLIR (Multi-Level Intermediate Representation) and hardware-specific dialects:

MLIR Bytecode: Compact binary serialization of MLIR modules;
StableHLO / MHLO: High-Level Optimizer dialect used across JAX, TensorFlow, OpenXLA, and PyTorch;
TOSA (Tensor Operator Set Architecture): Arm's standard neural network operator dialect;
XLA HLO: Computation graph IR compiled into target CPU/GPU/TPU machine code.

In compiled AI binaries, model artifacts, and accelerator kernels, the high-level framework information is lowered into these IR dialects before code generation.

While issue #3637 covers runtime model formats (ONNX, TFLite), it does not cover compiler intermediate representations.

This feature proposal introduces dedicated ML Compiler IR Reverse Engineering to REA, enabling automated parsing of MLIR bytecode, extraction of dialect operations (StableHLO, TOSA, Linalg, Affine, LLVM), and reconstruction of the tensor dataflow graph.

Technical Capabilities & Architecture

                                    +-----------------------------------------+
                                    | Input: MLIR Binary Bytecode (.mlirbc)   |
                                    |        or Compiled XLA HLO Artifact     |
                                    +--------------------+--------------------+
                                                         |
                         +-------------------------------+-------------------------------+
                         |                                                               |
                         v                                                               v
           +---------------------------+                                   +---------------------------+
           | MLIR Bytecode Stream      |                                   | Dialect Registry &        |
           | Reader (Sections, Strings)|                                   | Operation Classifier      |
           +-------------+-------------+                                   +-------------+-------------+
                         |                                                               |
                         +-------------------------------+-------------------------------+
                                                         |
                                                         v
                                     +---------------------------------------+
                                     | Region, Block & Operation Traversal   |
                                     | (Blocks, Values, Operands, Results)   |
                                     +-------------------+-------------------+
                                                         |
                         +-------------------------------+-------------------------------+
                         |                                                               |
                         v                                                               v
           +---------------------------+                                   +---------------------------+
           | StableHLO & TOSA Dialect  |                                   | Type, Tensor Shape &      |
           | Inversion (Dot, Conv, Pad)|                                   | Quantization Extraction   |
           +-------------+-------------+                                   +-------------+-------------+
                                                         |
                                                         v
                                     +---------------------------------------+
                                     | REA Evidence: Lowered IR Graph,       |
                                     | Dialect Operations & Source Provenance|
                                     +---------------------------------------+

1. MLIR Bytecode Container Ingestion

Bytecode Parser: Implement an inert parser for the MLIR bytecode format (ML\xefR magic bytes, version numbers, section headers).
String & Dialect Section: Extract dialect names, operation names, and attribute/type tables without requiring full LLVM/MLIR native libraries.

2. Multi-Dialect Operation Recovery

StableHLO Dialect: Recognize and reconstruct StableHLO tensor operations:
- stablehlo.dot_general: matrix multiplications and batch contractions.
- stablehlo.convolution: multi-dimensional convolutions, dilation, padding, and stride vectors.
- stablehlo.reduce_window, stablehlo.broadcast_in_dim, stablehlo.slice.
TOSA Dialect: Support Arm TOSA operations (tosa.conv2d, tosa.matmul, tosa.rescale).
Linalg & Affine Dialects: Recognize structured loop operations (linalg.generic, affine.for) representing lowered tensor kernels.

3. Tensor Types, Shapes & Attributes

Type System Decoding: Extract tensor ranks, element types (FP32, FP16, BF16, INT8, quantized uniform types), and dynamic dimensions.
Attributes: Recover dense resource elements, constant tensors, permutation arrays, and convolution padding geometries.

4. Graph Reconstruction & Lowering Provenance

SSA Value Flow: Track Static Single Assignment (SSA) values connecting operation results to downstream operands.
Source Location Mapping: Extract embedded loc("...") debug locations, correlating compiler IR operations back to original Python / PyTorch / JAX source lines.

Acceptance Criteria

Introduce analyze_mlir_artifact in REA CLI and MCP contracts, accepting MLIR bytecode and textual IR files.
Safely parse MLIR bytecode section tables and extract dialect names and operations.
Reconstruct StableHLO and TOSA tensor operations into structured JSON graph nodes.
Extract tensor shapes, dtypes, and operation attributes in structured evidence.
Map SSA dataflow dependencies and preserve source code location metadata.

#3637 [Feature] ML model graph reverse engineering
#3620 [Feature] Shader and GPU Bytecode Reverse Engineering
#3323 [Feature] Inspect GGUF and safetensors model files without loading them

Originally created by @matbrgz on GitHub (Oct 14, 2026). Original GitHub issue: https://github.com/morluto/rea/issues/3673 [![Type: Feature](https://img.shields.io/badge/type-feature-blue.svg?style=flat-square)](https://github.com/morluto/rea) [![Domain: Game & Engine RE](https://img.shields.io/badge/domain-game_&_engine_RE-green.svg?style=flat-square)](https://github.com/morluto/rea) [![Status: Proposed](https://img.shields.io/badge/status-proposed-blue.svg?style=flat-square)](https://github.com/morluto/rea) [![Ecosystem: REA](https://img.shields.io/badge/ecosystem-morluto%2Frea-8A2BE2.svg?style=flat-square)](https://github.com/morluto/rea) ## Summary Machine learning compilers (such as Google XLA, IREE, TVM, PyTorch Inductor, and ONNX Runtime) compile high-level neural network graphs into intermediate representations (IR) using **MLIR (Multi-Level Intermediate Representation)** and hardware-specific dialects: - **MLIR Bytecode:** Compact binary serialization of MLIR modules; - **StableHLO / MHLO:** High-Level Optimizer dialect used across JAX, TensorFlow, OpenXLA, and PyTorch; - **TOSA (Tensor Operator Set Architecture):** Arm's standard neural network operator dialect; - **XLA HLO:** Computation graph IR compiled into target CPU/GPU/TPU machine code. In compiled AI binaries, model artifacts, and accelerator kernels, the high-level framework information is lowered into these IR dialects before code generation. While issue #3637 covers runtime model formats (ONNX, TFLite), it does not cover compiler intermediate representations. This feature proposal introduces dedicated **ML Compiler IR Reverse Engineering** to REA, enabling automated parsing of MLIR bytecode, extraction of dialect operations (StableHLO, TOSA, Linalg, Affine, LLVM), and reconstruction of the tensor dataflow graph. --- ## Technical Capabilities & Architecture ``` +-----------------------------------------+ | Input: MLIR Binary Bytecode (.mlirbc) | | or Compiled XLA HLO Artifact | +--------------------+--------------------+ | +-------------------------------+-------------------------------+ | | v v +---------------------------+ +---------------------------+ | MLIR Bytecode Stream | | Dialect Registry & | | Reader (Sections, Strings)| | Operation Classifier | +-------------+-------------+ +-------------+-------------+ | | +-------------------------------+-------------------------------+ | v +---------------------------------------+ | Region, Block & Operation Traversal | | (Blocks, Values, Operands, Results) | +-------------------+-------------------+ | +-------------------------------+-------------------------------+ | | v v +---------------------------+ +---------------------------+ | StableHLO & TOSA Dialect | | Type, Tensor Shape & | | Inversion (Dot, Conv, Pad)| | Quantization Extraction | +-------------+-------------+ +-------------+-------------+ | v +---------------------------------------+ | REA Evidence: Lowered IR Graph, | | Dialect Operations & Source Provenance| +---------------------------------------+ ``` ### 1. MLIR Bytecode Container Ingestion - **Bytecode Parser:** Implement an inert parser for the MLIR bytecode format (`ML\xefR` magic bytes, version numbers, section headers). - **String & Dialect Section:** Extract dialect names, operation names, and attribute/type tables without requiring full LLVM/MLIR native libraries. ### 2. Multi-Dialect Operation Recovery - **StableHLO Dialect:** Recognize and reconstruct StableHLO tensor operations: - `stablehlo.dot_general`: matrix multiplications and batch contractions. - `stablehlo.convolution`: multi-dimensional convolutions, dilation, padding, and stride vectors. - `stablehlo.reduce_window`, `stablehlo.broadcast_in_dim`, `stablehlo.slice`. - **TOSA Dialect:** Support Arm TOSA operations (`tosa.conv2d`, `tosa.matmul`, `tosa.rescale`). - **Linalg & Affine Dialects:** Recognize structured loop operations (`linalg.generic`, `affine.for`) representing lowered tensor kernels. ### 3. Tensor Types, Shapes & Attributes - **Type System Decoding:** Extract tensor ranks, element types (FP32, FP16, BF16, INT8, quantized uniform types), and dynamic dimensions. - **Attributes:** Recover dense resource elements, constant tensors, permutation arrays, and convolution padding geometries. ### 4. Graph Reconstruction & Lowering Provenance - **SSA Value Flow:** Track Static Single Assignment (SSA) values connecting operation results to downstream operands. - **Source Location Mapping:** Extract embedded `loc("...")` debug locations, correlating compiler IR operations back to original Python / PyTorch / JAX source lines. --- ## Acceptance Criteria 1. Introduce `analyze_mlir_artifact` in REA CLI and MCP contracts, accepting MLIR bytecode and textual IR files. 2. Safely parse MLIR bytecode section tables and extract dialect names and operations. 3. Reconstruct StableHLO and TOSA tensor operations into structured JSON graph nodes. 4. Extract tensor shapes, dtypes, and operation attributes in structured evidence. 5. Map SSA dataflow dependencies and preserve source code location metadata. --- ## Related - #3637 [Feature] ML model graph reverse engineering - #3620 [Feature] Shader and GPU Bytecode Reverse Engineering - #3323 [Feature] Inspect GGUF and safetensors model files without loading them

bsakaguchi referenced this issue

2026-10-17 21:01:32 +00:00

[GH-ISSUE #2383] [Tracker] Audit analysis and debugging coverage for the TIOBE top 100 languages (October 2026) #794

bsakaguchi referenced this issue

2026-10-17 21:10:58 +00:00

[PR #1386] [MERGED] fix(javascript): handle oversized inline source maps #2845

bsakaguchi referenced this issue

2026-10-17 21:11:01 +00:00

[PR #1393] [MERGED] fix(javascript): stop reading media preload attributes as Electron preloads #2850