[BACKEND][NVIDIA] Update ptxas_options knobs default value (#8121)

This PR sets the default value of the ptxas_options knob to be the value
of the `PTXAS_OPTIONS` environment variable.
We need this option when we want to apply extra ptxas options across a
range of Triton kernels, without changing their call sites.

# New contributor declaration

- [x] I have written a PR description following these
  [rules](https://cbea.ms/git-commit/#why-not-how).

- [x] I have run `pre-commit run --from-ref origin/main --to-ref HEAD`.

- Select one of the following.
- [x] This PR does not need a test because this is a small change of
reading ptxas options from environment variables.

- Select one of the following.
  - [x] I have not added any `lit` tests.

Author: xuzhao9

README.md

@@ -232,6 +232,7 @@ See [`python/triton/knobs.py`](python/triton/knobs.py) for the full list of conf
 - `TRITON_F32_DEFAULT` sets the default input precision of `tl.dot` when using 32-bit floats, which can be either `ieee`, `tf32`, or `tf32x3`.
 - `TRITON_FRONT_END_DEBUGGING=1` disables exception wrapping when an error occurs in the compiler frontend, allowing the full stack trace to be seen.
 - `TRITON_DISABLE_LINE_INFO=1` removes all line information from the module.
+- `PTXAS_OPTIONS` passes additional command-line options to the PTX assembler `ptxas` (only on NVIDIA).
 
 > [!NOTE]
 > Some of these environment variables don't have a knob in `knobs.py`-- those are only relevant to the C++ layer(s), hence they don't exist in the python layer.

python/test/unit/test_knobs.py

@@ -247,30 +247,42 @@ def test_nvidia_tool(fresh_knobs, tmp_path, monkeypatch):
     default_ptxas = triton_root / "backends/nvidia/bin/ptxas"
 
     assert Path(fresh_knobs.nvidia.ptxas.path).resolve() == default_ptxas.resolve()
+    assert fresh_knobs.nvidia.ptxas_options is None
 
     tmp_ptxas = tmp_path / "ptxas-special"
     shutil.copy(default_ptxas, tmp_ptxas)
     monkeypatch.setenv("TRITON_PTXAS_PATH", str(tmp_ptxas))
+    monkeypatch.setenv("PTXAS_OPTIONS", "--verbose")
     assert Path(fresh_knobs.nvidia.ptxas.path).resolve() == tmp_ptxas.resolve()
+    assert fresh_knobs.nvidia.ptxas_options == "--verbose"
 
     # Don't prop so that the `del` is correctly tested
     fresh_knobs.propagate_env = False
     fresh_knobs.nvidia.ptxas = str(default_ptxas)
+    fresh_knobs.nvidia.ptxas_options = "--device-debug"
     fresh_knobs.propagate_env = True
     assert Path(fresh_knobs.nvidia.ptxas.path).resolve() == default_ptxas.resolve()
+    assert fresh_knobs.nvidia.ptxas_options == "--device-debug"
 
     del fresh_knobs.nvidia.ptxas
+    del fresh_knobs.nvidia.ptxas_options
     assert Path(fresh_knobs.nvidia.ptxas.path).resolve() == tmp_ptxas.resolve()
+    assert fresh_knobs.nvidia.ptxas_options == "--verbose"
 
     # Triple check scope works
     with fresh_knobs.nvidia.scope():
         fresh_knobs.nvidia.ptxas = str(default_ptxas)
+        fresh_knobs.nvidia.ptxas_options = "--device-debug"
         assert Path(fresh_knobs.nvidia.ptxas.path).resolve() == default_ptxas.resolve()
+        assert fresh_knobs.nvidia.ptxas_options == "--device-debug"
 
     assert Path(fresh_knobs.nvidia.ptxas.path).resolve() == tmp_ptxas.resolve()
+    assert fresh_knobs.nvidia.ptxas_options == "--verbose"
 
     monkeypatch.delenv("TRITON_PTXAS_PATH")
+    monkeypatch.delenv("PTXAS_OPTIONS")
     assert Path(fresh_knobs.nvidia.ptxas.path).resolve() == default_ptxas.resolve()
+    assert fresh_knobs.nvidia.ptxas_options is None
 
 
 def test_opt_bool(fresh_knobs, monkeypatch):

python/triton/knobs.py

@@ -478,6 +478,7 @@ class nvidia_knobs(base_knobs):
 
     dump_nvptx: env_bool = env_bool("NVPTX_ENABLE_DUMP")
     disable_ptxas_opt: env_bool = env_bool("DISABLE_PTXAS_OPT")
+    ptxas_options: env_opt_str = env_opt_str("PTXAS_OPTIONS")
     mock_ptx_version: env_opt_str = env_opt_str("TRITON_MOCK_PTX_VERSION")
     dump_ptxas_log: env_bool = env_bool("TRITON_DUMP_PTXAS_LOG")
 

third_party/nvidia/backend/compiler.py

@@ -113,7 +113,7 @@ class CUDAOptions:
     maxnreg: Optional[int] = None
     cluster_dims: tuple = (1, 1, 1)
     ptx_version: int = None
-    ptx_options: str = None
+    ptx_options: Optional[str] = knobs.nvidia.ptxas_options
     ir_override: Optional[str] = None  # filename of a user-defined IR (*.{ttir|ttgir|llir|ptx})
     enable_fp_fusion: bool = True
     launch_cooperative_grid: bool = False