[FRONTEND] Expand MLIR bindings for out-of-tree walk by aakhundov · Pull Request #3191 · triton-lang/triton

aakhundov · 2024-02-24T03:31:50Z

Summary: This PR adds a few simple MLIR bindings to ir.cc to allow walking the MLIR structure of the TTIR module out of tree. This will help making the Triton kernel analysis performed in PyTorch 2 more robust and reliable (related PR in PT2: pytorch/pytorch#120476).

Test Plan:

$ python3 -m pytest -s python/test/unit/runtime/test_bindings.py
===================== test session starts =====================
platform linux -- Python 3.11.3, pytest-8.0.1, pluggy-1.4.0
rootdir: /data/users/aakhundov/triton/python
plugins: hypothesis-6.29.3, hydra-core-1.3.2
collected 1 item

python/test/unit/runtime/test_bindings.py .

===================== 1 passed in 3.02s =======================

Summary: This PR adds a few simple MLIR bindings to `ir.cc` to allow walking the MLIR structure of the TTIR module out of tree. This will help making the Triton kernel analysis performed in PyTorch 2 more robust and reliable. Related PR in PT2: pytorch/pytorch#120476 Test Plan: ``` $ python3 -m pytest -s python/test/unit/runtime/test_bindings.py ===================== test session starts ===================== platform linux -- Python 3.11.3, pytest-8.0.1, pluggy-1.4.0 rootdir: /data/users/aakhundov/triton/python plugins: hypothesis-6.29.3, hydra-core-1.3.2 collected 1 item python/test/unit/runtime/test_bindings.py . ===================== 1 passed in 3.02s ======================= ```

oulgen · 2024-02-24T03:39:35Z

@ptillet This PR is regarding the conversation we had over Slack about exposing more IR bindings for read-only introspection: https://triton-lang.slack.com/archives/C01L894JD98/p1708105403856079

@jansel just FYI

…riton kernel mutation analysis" Summary: Previously, we relied on the `lark`-based parsing of the string TTIR representation dumped by the Triton compiler. However, this has proven to be brittle in the face of changes both in the user-written Triton kernel code and in the Triton compiler code. In this PR, we add an alternative way of mining the function information from the TTIR based on walking the tree of structured MLIR entities. To this end, we rely on the MLIR bindings exposed by `libtriton` (related PR in Triton: triton-lang/triton#3191). For now, we introduce gating based on whether `ttir_module.hasattr("walk")`. This will allow switching to the newly introduced TTIR analysis approach only when the new MLIR bindings (including that of `ModuleOp::walk`) become available in the Triton pin. Before then, we'll keep using the old string TTIR parsing-based approach. Test Plan: The new functionality was tested locally with the latest Triton version compiled with the added new MLIR bindings: all Triton kernel mutation tests in `test_triton_kernels.py` are passing. Here we rely on the CI for regression testing, but it won't cover the new functionality due to gating. Reviewers: Subscribers: Tasks: Tags: cc zou3519 ydwu4 [ghstack-poisoned]

…d MLIR walk in Triton kernel mutation analysis" Summary: Previously, we relied on the `lark`-based parsing of the string TTIR representation dumped by the Triton compiler. However, this has proven to be brittle in the face of changes both in the user-written Triton kernel code and in the Triton compiler code. In this PR, we add an alternative way of mining the function information from the TTIR based on walking the tree of structured MLIR entities. To this end, we rely on the MLIR bindings exposed by `libtriton` (related PR in Triton: triton-lang/triton#3191). For now, we introduce gating based on whether `ttir_module.hasattr("walk")`. This will allow switching to the newly introduced TTIR analysis approach only when the new MLIR bindings (including that of `ModuleOp::walk`) become available in the Triton pin. Before then, we'll keep using the old string TTIR parsing-based approach. Test Plan: The new functionality was tested locally with the latest Triton version compiled with the added new MLIR bindings: all Triton kernel mutation tests in `test_triton_kernels.py` are passing. Here we rely on the CI for regression testing, but it won't cover the new functionality due to gating. Reviewers: Subscribers: Tasks: Tags: cc zou3519 ydwu4 [ghstack-poisoned]