[lang] Add pass-through flags for the ptx compiler

Signed-off-by: Asher Mancinelli <amancinelli@nvidia.com>

Author: ashermancinelli

experimental/cuda-lang/src/cuda/lang/_compile.py

@@ -60,11 +60,16 @@ def mlir2cubin(
     gpu_name: str,
     arch: str,
     emit_ptx: bool = False,
+    ptx_compiler_options: Sequence[str] = (),
 ) -> MLIR2CubinResult:
     executable = get_compiler_binary_path()
     argv = [executable, "-", "-o", "-", f"--gpu-name={gpu_name}", f"--arch={arch}"]
     custom_flags = os.environ.get("CUDA_LANG_MLIR2CUBIN_FLAGS", None)
 
+    argv.extend(
+        f"--ptx-compiler-option={option}" for option in ptx_compiler_options
+    )
+
     if custom_flags is not None:
         argv.extend(custom_flags.split())
 
@@ -266,10 +271,18 @@ def _dump(phase: str, contents: object) -> None:
         arch = arch or cc.arch + suffix
 
     need_ptx = log_flags.log_ptx or keep_ptx
+    ptx_compiler_options = compiler_options._ptx_compiler_options
     compiled = mlir2cubin(
-        mlir_text, gpu_name=gpu_name, arch=arch, emit_ptx=need_ptx
+        mlir_text,
+        gpu_name=gpu_name,
+        arch=arch,
+        emit_ptx=need_ptx,
+        ptx_compiler_options=ptx_compiler_options,
     )
 
+    if compiled.stderr and ptx_compiler_options:
+        _dump("PTX compiler", compiled.stderr.decode())
+
     if need_ptx:
         assert compiled.ptx is not None
 

experimental/cuda-lang/src/cuda/lang/_compiler_options.py

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
-from dataclasses import dataclass
+from dataclasses import dataclass, field, fields
 from ._exception import TypeCheckingError
 
 
@@ -14,6 +14,22 @@ class CompilerOptions:
     max_blocks_per_cluster: int | None = None
     max_registers_per_thread: int | None = None
     min_blocks_per_sm: int | None = None
+    _ptx_compiler_verbose: bool = field(
+        default=False,
+        metadata={"PTX_FLAG": "--verbose"},
+    )
+    _ptx_compiler_warn_on_local_memory_usage: bool = field(
+        default=False,
+        metadata={"PTX_FLAG": "--warn-on-local-memory-usage"},
+    )
+    _ptx_compiler_warn_on_spills: bool = field(
+        default=False,
+        metadata={"PTX_FLAG": "--warn-on-spills"},
+    )
+    _ptx_compiler_make_errors_visible_at_exit: bool = field(
+        default=False,
+        metadata={"PTX_FLAG": "--make-errors-visible-at-exit"},
+    )
 
     def __post_init__(self):
         message = (
@@ -29,15 +45,34 @@ def __post_init__(self):
             case _:
                 raise TypeCheckingError(message)
 
-        for field in (
+        for field_ in (
             "max_blocks_per_cluster",
             "max_registers_per_thread",
             "min_blocks_per_sm",
         ):
-            value = getattr(self, field)
+            value = getattr(self, field_)
             if value is not None and (not isinstance(value, int) or value < 0):
                 message = (
-                    f"Expected compiler option {field} to be a "
+                    f"Expected compiler option {field_} to be a "
                     f"positive integer but got {value}"
                 )
                 raise TypeCheckingError(message)
+
+        for field_ in fields(self):
+            if "PTX_FLAG" not in field_.metadata:
+                continue
+
+            if not isinstance(getattr(self, field_.name), bool):
+                message = (
+                    f"Expected compiler option {field_.name} to be bool "
+                    f"but got {type(getattr(self, field_.name))}"
+                )
+                raise TypeCheckingError(message)
+
+    @property
+    def _ptx_compiler_options(self) -> tuple[str, ...]:
+        return tuple(
+            field_.metadata["PTX_FLAG"]
+            for field_ in fields(self)
+            if "PTX_FLAG" in field_.metadata and getattr(self, field_.name)
+        )

experimental/cuda-lang/src/cuda/lang/_execution.py

@@ -138,6 +138,10 @@ def __init__(
         max_blocks_per_cluster: int | None = None,
         max_registers_per_thread: int | None = None,
         min_blocks_per_sm: int | None = None,
+        _ptx_compiler_verbose: bool = False,
+        _ptx_compiler_warn_on_local_memory_usage: bool = False,
+        _ptx_compiler_warn_on_spills: bool = False,
+        _ptx_compiler_make_errors_visible_at_exit: bool = False,
     ):
         """
         Args:
@@ -168,6 +172,10 @@ def __init__(
             max_blocks_per_cluster=max_blocks_per_cluster,
             max_registers_per_thread=max_registers_per_thread,
             min_blocks_per_sm=min_blocks_per_sm,
+            _ptx_compiler_verbose=_ptx_compiler_verbose,
+            _ptx_compiler_warn_on_local_memory_usage=_ptx_compiler_warn_on_local_memory_usage,
+            _ptx_compiler_warn_on_spills=_ptx_compiler_warn_on_spills,
+            _ptx_compiler_make_errors_visible_at_exit=_ptx_compiler_make_errors_visible_at_exit,
         )
         self._arch = arch
         self._gpu_name = gpu_name

experimental/cuda-lang/test/test_kernel_attributes.py

@@ -6,6 +6,7 @@
 from cuda.lang._exception import TypeCheckingError
 from test.util import compile_kernel, require_hopper_or_newer
 import pytest
+import torch
 
 
 @require_hopper_or_newer()
@@ -20,6 +21,10 @@
         dict(max_threads_per_block=1.0),
         dict(max_registers_per_thread=1.0),
         dict(min_blocks_per_sm=1.0),
+        dict(_ptx_compiler_verbose=1.0),
+        dict(_ptx_compiler_warn_on_local_memory_usage=1.0),
+        dict(_ptx_compiler_warn_on_spills=1.0),
+        dict(_ptx_compiler_make_errors_visible_at_exit=1.0),
     ),
 )
 def test_bad_kernel_attribute_types(kwarg):
@@ -119,3 +124,41 @@ def foo():
         CHECK-NEXT: .minnctapersm 2
         """,
     )
+
+
+@pytest.mark.parametrize(
+    "option",
+    (
+        "_ptx_compiler_verbose",
+        "_ptx_compiler_warn_on_local_memory_usage",
+        "_ptx_compiler_warn_on_spills",
+        "_ptx_compiler_make_errors_visible_at_exit",
+    ),
+)
+@pytest.mark.parametrize("value", (True, False))
+def test_ptx_compiler_accepts_all_options(option, value):
+    @cl.kernel(**{option: value})
+    def foo():
+        pass
+
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), foo, ())
+
+
+def test_ptx_compiler_verbose_stderr(capsys):
+    @cl.kernel(_ptx_compiler_verbose=True)
+    def foo():
+        pass
+
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), foo, ())
+    outerr = capsys.readouterr()
+    assert "ptxas info" in outerr.err
+
+
+def test_ptx_compiler_stderr_empty_when_not_verbose(capsys):
+    @cl.kernel(_ptx_compiler_verbose=False)
+    def foo():
+        pass
+
+    cl.launch(torch.cuda.current_stream(), (1,), (1,), foo, ())
+    outerr = capsys.readouterr()
+    assert "ptxas info" not in outerr.err