[FEATURE] NKI interpreter

[latentCall145]

sorry for big diff :(

Adds the NKI frontend to Triton-Viz:

• added a NKI interpreter to apply language-level patching (e.g. allow NKI kernels to be executed line-by-line with NumPy; modifies nki.py)
• AST rewriting to convert NKI loads/stores to masked loads/stores (modifies nki_extract_slice.py)
• Masked load/store implementation (modifies nki_masked_load.py)
• added backend (one of "nki" or "triton") as arguments for patching (modifies patch.py, client.py)
• added tracer callbacks for masked loads/stores (modifies data.py, tracer.py)
• separate Triton and NKI Trace objects (modifies trace.py)
• a lot of visualization stuff (credit: @gujialiang123)

Added Tests

• nki-examples/: demo programs to try out the visualizer (currently, only nki-examples/matmul.py is supported)
• tests/test_nki.py: basic NDArray slicing
• tests/test_masked_load.py: make sure masked load/store works correctly

[Jokeren]

Hi @mark14wu , can you also review the PR?

[mark14wu]

Hi @mark14wu , can you also review the PR?

Yes. I have dropped some comments, but haven't finished reviewing all files.

[mark14wu]

Some tests in CI is not working.

Traceback:
/opt/hostedtoolcache/Python/3.10.19/x64/lib/python3.10/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
tests/test_wrapper.py:8: in <module>
    from triton_viz.wrapper import SANITIZER_COMMAND, PROFILER_COMMAND
triton_viz/__init__.py:1: in <module>
    from .core import trace, clear, config
triton_viz/core/__init__.py:1: in <module>
    from .trace import trace, clear
triton_viz/core/trace.py:7: in <module>
    from ..clients import Sanitizer, Profiler, Tracer
triton_viz/clients/__init__.py:1: in <module>
    from .profiler.profiler import Profiler
triton_viz/clients/profiler/profiler.py:1: in <module>
    from ...core.client import Client
triton_viz/core/client.py:8: in <module>
    from .patch import (
triton_viz/core/patch.py:61: in <module>
    from triton_viz.core.nki import nki_builder
triton_viz/core/nki.py:3: in <module>
    import neuronxcc.nki.language as nl
E   ModuleNotFoundError: No module named 'neuronxcc'

Please fix.