[rocjitsu] Add Python execution sessions and startup activation

emulation/rocjitsu/CMakeLists.txt

@@ -465,6 +465,22 @@ if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
         PROGRAMS scripts/rocjitsu-pretranslate.py
         DESTINATION ${CMAKE_INSTALL_BINDIR}
     )
+
+    # Python wheel inputs: lib/python3/site-packages/ and share/rocjitsu/python/.
+    # This version-independent staging directory is not on Python's import
+    # path; source installs still need `pip install ./python`.
+    install(
+        DIRECTORY python/rocjitsu
+        DESTINATION ${CMAKE_INSTALL_LIBDIR}/python3/site-packages
+        FILES_MATCHING
+        PATTERN "*.py"
+        PATTERN "__pycache__" EXCLUDE
+    )
+    # Wheel builders consume the same metadata as the standalone Python package.
+    install(
+        FILES python/pyproject.toml python/README.md
+        DESTINATION ${CMAKE_INSTALL_DATADIR}/rocjitsu/python
+    )
 endif()
 
 # Config files and schemas: share/rocjitsu/

emulation/rocjitsu/CONTRIBUTING.md

@@ -30,6 +30,7 @@
 | Shared utility (bit ops, logging, SIMD, allocators) | `lib/util/` |
 | JSON config schema change | `schemas/` + `lib/rocjitsu/src/rocjitsu/config/` |
 | CLI changes | `tools/rocjitsu/` |
+| Python execution API | `python/` (tests in `python/tests/`) |
 | New GPU topology config | `configs/` |
 | Tests | `tests/` (mirror the source structure: `tests/dbt/`, `tests/analysis/`, `tests/patch/`, etc.) |
 

emulation/rocjitsu/README.md

@@ -50,7 +50,8 @@ lib/
       config/           JSON/FlatBuffers configuration
   util/                 Shared utilities
   python/amdisa/        ISA codegen pipeline
-tools/rocjitsu/         CLI (local, daemon, attach modes)
+python/                 Python execution sessions and startup activation
+tools/rocjitsu/         CLI (local, daemon, attach, preload-only modes)
 configs/                GPU topology JSON files
 schemas/                FlatBuffers schemas
 tests/                  Test suite
@@ -146,6 +147,10 @@ rocjitsu --daemon --config configs/gfx950_mi355x_kmd.json -- \
 
 See [docs/building.md](docs/building.md) for container setup with PyTorch.
 
+For programmatic execution and notebooks, the [Python API](python/README.md)
+provides `rocjitsu.run`, persistent `Session` workers, and startup-only global
+activation with `rocjitsu.enable`.
+
 ## Documentation
 
 ### Getting started
@@ -156,6 +161,7 @@ See [docs/building.md](docs/building.md) for container setup with PyTorch.
 | [Benchmarking](docs/benchmarking.md) | Reproducible performance and memory measurement protocol |
 | [Configuration](docs/configuration.md) | JSON config format and topology |
 | [QEMU VFIO-user compute](docs/qemu-vfio.md) | Run a gfx1250 guest and qualify a GEMM workload |
+| [Python API](python/README.md) | Callable workers, sessions, notebooks, and startup activation |
 | [CLI & Transport](docs/rocjitsu-cli.md) | CLI modes, daemon RPC protocol |
 | [Race Detector](docs/race-detector.md) | Race detection tutorial and internals |
 | [Debugging with ROCgdb](docs/rocgdb-debugging.md) | Debug emulated GPU kernels with ROCgdb: breakpoints, watchpoints, faults, multi-wave |

emulation/rocjitsu/docs/rocjitsu-cli.md

@@ -3,10 +3,12 @@
 ## Overview
 
 The `rocjitsu` CLI is the primary entry point for running applications on
-the simulated GPU. It operates in four modes:
+the simulated GPU. It supports these modes:
 
 - **Local mode**: Sets `LD_PRELOAD` and `execve`s the target application.
   The simulator runs in-process via the interposer.
+- **Preload-only mode**: Preloads the interposer and defers simulation
+  configuration to `rocjitsu.enable()` before GPU discovery.
 - **Daemon mode**: Spawns a daemon child process that hosts the simulation
   engine, then `LD_PRELOAD` + `execve`s the target. Client processes
   communicate with the daemon via RPC over a Unix domain socket. GPU memory
@@ -48,12 +50,14 @@ receiving real GPU descriptors over a Unix socket.
 
 ```
 Usage: rocjitsu --config <config.json> [--daemon|--attach] -- <app> [args...]
+       rocjitsu --preload-only -- <app> [args...]
        rocjitsu --config <config.json> --vfio-socket <path>
 ```
 
 | Option | Description |
 |--------|-------------|
-| `--config <path>` | Path to simulation config JSON (required) |
+| `--config <path>` | Path to simulation config JSON (required except with `--preload-only`) |
+| `--preload-only` | Preload the interposer without a config. The application must call `rocjitsu.enable(config)` in Python before GPU discovery; the Python bootstrap `python -m rocjitsu` uses this mode. Cannot be combined with configuration or other execution modes. |
 | `--daemon` | Run in daemon mode: fork a daemon process hosting the simulation engine, then launch the application with the interposer. Without `-- <app>`, runs the daemon server only. |
 | `--attach` | Attach to a running daemon. The socket path is resolved as `$ROCJITSU_RUNTIME_DIR/daemon.sock`, then `$XDG_RUNTIME_DIR/rocjitsu/daemon.sock`, falling back to `/tmp/rocjitsu-<uid>/daemon.sock`. |
 | `--vfio-socket <path>` | Serve the configured GPU as a PCI function over VFIO-user. The VMM must share guest RAM through mmap-able file descriptors. |
@@ -335,3 +339,22 @@ kernel 6.12.59+ and 6.15+.
 - **Windows**: Named pipe transport, KMD escape driver backend
 - **RDMA**: Rack-scale transport for distributed simulation
 - **Platform-independent commands**: Abstract command enum for cross-platform
+
+### Programmatic activation environment
+
+`--preload-only` sets `ROCJITSU_PROGRAMMATIC=1`. Interposed KFD/DRM device
+and discovery paths refuse access until the application activates a config.
+A discovery attempt prevents later activation because frameworks may cache the
+failed result. Ordinary configured launches set the variable to `0`.
+
+After successful activation, Python exports `ROCJITSU_INVOCATION_DIR` pointing
+to its private config handoff and sets `ROCJITSU_PROGRAMMATIC=2` for fresh exec
+children. The interpreter that calls `enable()` owns this handoff and must
+remain alive until those children finish their GPU work.
+
+The inherited mode (`2`) permits an identical `enable` call when multiprocessing
+re-imports a module in a spawn or forkserver child. The native parser compares
+the effective configs, including any CPU thread budget override; a different
+config is refused. Configured CLI launches (`0`) still reject activation.
+These environment values and the native entry point are private implementation
+details of the Python binding.

emulation/rocjitsu/docs/sphinx/index.rst

@@ -51,6 +51,7 @@ The rocJITsu public repository is located at
       * :doc:`Status codes and compiler macros reference <reference/api-status>`
       * :doc:`rj_dbt_translate command reference <reference/rj-dbt-translate>`
       * :doc:`rocjitsu CLI reference <reference/rocjitsu-cli>`
+      * :doc:`Python execution API <reference/python-api>`
       * :doc:`Environment variable reference <reference/environment-variables>`
       * :doc:`Supported GPU architectures <reference/supported-architectures>`
 

emulation/rocjitsu/docs/sphinx/install/install.md

@@ -100,6 +100,13 @@ This installs:
 -   GPU topology JSON configuration files to
     `<prefix>/share/rocjitsu/configs/`
 -   FlatBuffers schema files to `<prefix>/share/rocjitsu/schemas/`
+-   Python package sources for wheel builders to
+    `<prefix>/lib/python3/site-packages/rocjitsu/` (Linux only)
+-   Python packaging metadata to `<prefix>/share/rocjitsu/python/` (Linux only)
+
+The staged Python sources are not on Python's import path. For a source build,
+install the package with `pip install ./python` from `emulation/rocjitsu/`;
+see the [Python execution API](../reference/python-api.md) for notebook setup.
 
 #### Third-party dependencies
 

emulation/rocjitsu/docs/sphinx/reference/environment-variables.md

@@ -91,6 +91,16 @@ Path to the AMD ROCm installation.
 | **Allowed values** | Any valid directory path pointing to a ROCm installation |
 | **Effect** | Used to locate ROCm libraries and tools for optional ROCm-dependent tests when `BUILD_TESTING=ON`. An explicitly empty CMake `ROCM_PATH` disables environment and `/opt/rocm` fallbacks. It does not select the host compiler or install destination. |
 
+### `ROCJITSU_EXECUTABLE`
+
+Override for the native executable used by the Python package.
+
+| | |
+|---|---|
+| **Default** | The bundled wheel's launcher (no fallback); source installs use the Python environment's `bin/rocjitsu`, then `rocjitsu` on `PATH` |
+| **Allowed values** | An executable name on `PATH` or a path to the native executable |
+| **Effect** | Selects the launcher for Python sessions, bootstrap and diagnostics. An explicit `executable=` argument takes priority. |
+
 ### `ROCJITSU_RUNTIME_DIR`
 
 Override for the runtime directory used by daemon socket resolution.

emulation/rocjitsu/docs/sphinx/reference/python-api.md

@@ -0,0 +1,61 @@
+# Python execution API
+
+Run upstream ROCm PyTorch and Triton on simulated AMD GPUs from a notebook,
+including Colab, using the `rocjitsu` Python package. TheRock builds containing
+the Python API and its wheel packaging support provide a `rocjitsu` wheel with
+the native runtime and JSON presets included. Install it in the notebook's
+Python environment, along with ROCm-enabled PyTorch and a Triton compiler that
+supports the simulated target:
+
+```python
+%pip install /path/to/rocjitsu-<version>-py3-none-<linux-platform>.whl
+```
+
+For a source build:
+
+```sh
+# From emulation/rocjitsu, with the notebook environment active:
+pip install ./python
+cmake --install /path/to/native/build --prefix "$VIRTUAL_ENV"
+```
+
+The native executable, worker Python and JSON presets come from that
+installation. Preset attributes match the installed JSON filename stems.
+
+```python
+import rocjitsu
+
+def add_one():
+    import torch
+    return (torch.arange(8, device="cuda") + 1).cpu()
+
+with rocjitsu.Session(rocjitsu.config.gfx1250_mi455x) as sim:
+    result = sim.run(add_one)
+result
+```
+
+`run` starts a fresh worker for one call. `Session` retains imports and worker
+state across calls; close it explicitly or use a context manager. Send CPU
+values and return CPU results. GPU tensors and CPU tensors requiring gradients
+cannot cross the process boundary. ROCm-enabled PyTorch is required in the
+worker, including when the notebook itself uses CUDA PyTorch.
+
+For a CUDA-based notebook, install `rocjitsu`, ROCm PyTorch and Triton into a
+separate environment with the same Python major/minor/patch and cloudpickle
+versions as the kernel, then select that worker interpreter:
+
+```python
+with rocjitsu.Session(
+    rocjitsu.config.gfx1250_mi455x,
+    python="/path/to/rocm-venv/bin/python",
+) as sim:
+    result = sim.run(add_one)
+```
+
+For a script, `python -m rocjitsu app.py` preloads the interposer before Python
+starts, and `app.py` must call `rocjitsu.enable(config)` before GPU discovery.
+An already running notebook uses `run` or `Session`.
+
+See the [package documentation](https://github.com/ROCm/rocm-systems/blob/develop/emulation/rocjitsu/python/README.md)
+for installation, callable transport, launch options, errors and activation
+inheritance, and the [CLI reference](rocjitsu-cli.md) for native launch modes.

emulation/rocjitsu/docs/sphinx/reference/rocjitsu-cli.md

@@ -20,7 +20,8 @@ rocjitsu --config <config.json> --vfio-socket <path>
 
 | Option | Description |
 | --- | --- |
-| `--config <path>` | Path to the simulation configuration JSON file. Required for all modes. |
+| `--config <path>` | Path to the simulation configuration JSON file. Required except with `--preload-only`. |
+| `--preload-only` | Preload the interposer without a config. The application must call `rocjitsu.enable(config)` in Python before GPU discovery; the Python bootstrap `python -m rocjitsu` uses this mode. Cannot be combined with configuration or other execution modes. |
 | `--daemon` | Run in daemon mode: fork a daemon process hosting the simulation engine, then launch the application with the interposer. Without `-- <app>`, runs the daemon server only (no application is launched). |
 | `--attach` | Attach to a running daemon. The socket path is resolved as described in [Environment variables and socket path resolution](#socket-path-resolution). |
 | `--vfio-socket <path>` | Serve the configured GPU as a PCI function over VFIO-user. Requires a build configured with `ROCJITSU_ENABLE_VFIO=ON` and a VMM that shares guest RAM through mmap-able file descriptors. |
@@ -150,3 +151,22 @@ File descriptors (`memfd` handles) are passed via `sendmsg()`/`recvmsg()` with `
 4.  The runtime calls `DESTROY_QUEUE`, `DESTROY_EVENT`, and `close(kfd_fd)`.
 5.  The client's `RemoteDriver` sends `RPC_CLOSE` to the daemon.
 6.  The daemon closes the client connection and, when all clients have disconnected, shuts down the simulation engine.
+
+### Programmatic activation environment
+
+`--preload-only` sets `ROCJITSU_PROGRAMMATIC=1`. Interposed KFD/DRM device
+and discovery paths refuse access until the application activates a config.
+A discovery attempt prevents later activation because frameworks may cache the
+failed result. Ordinary configured launches set the variable to `0`.
+
+After successful activation, Python exports `ROCJITSU_INVOCATION_DIR` pointing
+to its private config handoff and sets `ROCJITSU_PROGRAMMATIC=2` for fresh exec
+children. The interpreter that calls `enable()` owns this handoff and must
+remain alive until those children finish their GPU work.
+
+The inherited mode (`2`) permits an identical `enable` call when multiprocessing
+re-imports a module in a spawn or forkserver child. The native parser compares
+the effective configs, including any CPU thread budget override; a different
+config is refused. Configured CLI launches (`0`) still reject activation.
+These environment values and the native entry point are private implementation
+details of the Python binding.

emulation/rocjitsu/docs/sphinx/sphinx/_toc.yml

@@ -32,6 +32,7 @@ subtrees:
       - file: reference/api-status
       - file: reference/rj-dbt-translate
       - file: reference/rocjitsu-cli
+      - file: reference/python-api
       - file: reference/environment-variables
       - file: reference/supported-architectures
   - caption: About

emulation/rocjitsu/docs/sphinx/sphinx/_toc.yml.in

@@ -32,6 +32,7 @@ subtrees:
       - file: reference/api-status
       - file: reference/rj-dbt-translate
       - file: reference/rocjitsu-cli
+      - file: reference/python-api
       - file: reference/environment-variables
       - file: reference/supported-architectures
       - file: reference/changelog