kernel-lens 1.2.1__tar.gz → 2.0.0__tar.gz

{kernel_lens-1.2.1 → kernel_lens-2.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernel-lens
-Version: 1.2.1
+Version: 2.0.0
 Summary: Compile PyTorch models with custom @triton.jit kernels into ONNX Runtime and TensorRT C++ plugins
 Author-email: Justin Duc <justin.duc@galiad-research.com>
 License: MIT License
@@ -52,12 +52,12 @@ Requires-Dist: onnx>=1.15.0
 Provides-Extra: ort
 Requires-Dist: onnxruntime-gpu>=1.17.0; extra == "ort"
 Provides-Extra: trt
-Requires-Dist: tensorrt<11.0.0,>=8.6.0; extra == "trt"
-Requires-Dist: tensorrt-cu12-libs<11.0.0; extra == "trt"
+Requires-Dist: tensorrt<12.0.0,>=8.6.0; extra == "trt"
+Requires-Dist: tensorrt-cu12-libs; extra == "trt"
 Provides-Extra: all
 Requires-Dist: onnxruntime-gpu>=1.17.0; extra == "all"
-Requires-Dist: tensorrt<11.0.0,>=8.6.0; extra == "all"
-Requires-Dist: tensorrt-cu12-libs<11.0.0; extra == "all"
+Requires-Dist: tensorrt<12.0.0,>=8.6.0; extra == "all"
+Requires-Dist: tensorrt-cu12-libs; extra == "all"
 Provides-Extra: bench
 Requires-Dist: transformers>=4.40.0; extra == "bench"
 Dynamic: license-file
@@ -114,6 +114,18 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
 - **Dynamic grids in C++** — launch-grid expressions are transpiled to C++
   integer math over runtime shapes, so one plugin serves dynamic batch and
   sequence lengths.
+- **One package, several TensorRT versions** — the plugin interface is chosen
+  from the environment. TensorRT 10 deprecated the whole `IPluginV2*` hierarchy
+  and TensorRT 11 removed it, so `plugin_api="auto"` (the default) emits
+  `IPluginV3` plugins where the runtime supports them and `IPluginV2DynamicExt`
+  plugins otherwise. This spans TensorRT 8.6 through 10.x today; TensorRT 11+
+  works when you point `TENSORRT_INCLUDE_DIR` at matching headers.
+- **Shape-generic engines** — `kl.compile(..., dynamic_shapes=True)` exports the
+  leading dimension as a symbolic axis, so a single engine serves any size in a
+  declared range via a TensorRT optimization profile. Verified on Blackwell
+  (`sm_120a`) and Ampere (`sm_80`): one compile at batch 8 serving batches
+  1-16, bit-exact on both engines. Off by default, because a shape-specialised
+  engine is what the published benchmarks measure.
 - **Zero-copy VRAM binding** — PyTorch `data_ptr()` values are registered
   directly with each engine; no host staging copies.
 - **Blackwell PTX ISA clamping** — a Triton-generated `.version 9.3` header is
@@ -151,11 +163,34 @@ pip install kernel-lens[all]     # both
 Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`. TensorRT must be `< 11` (the
 generator emits `IPluginV2DynamicExt`).
 
+**TensorRT versions.** Handled automatically. The pip TensorRT wheels ship **no
+headers**, so KernelLens fetches the set matching the installed runtime from the
+TensorRT OSS repository and caches it under `~/.kernel_lens_headers/<major>.<minor>`;
+when the headers it already has match the runtime major, nothing is fetched. If
+the download is unavailable the fallback is the bundled set, and a
+header/runtime major mismatch is then reported instead of crashing inside
+TensorRT. `plugin_api="auto"` picks the interface from the runtime:
+
+| TensorRT | `IPluginV2*` | `IPluginV3` | emitted |
+| :--- | :--- | :--- | :--- |
+| < 10 (8.6, 9.x) | present | absent | `v2` |
+| 10.x | deprecated | present | `v3` |
+| 11+ | **removed** | present | `v3` |
+
+Verified on TensorRT 10.16 and 11.3, both 16/16 on the conformance suite. Force
+a choice with `plugin_api="v2"` / `"v3"`, or the `KERNEL_LENS_PLUGIN_API`
+environment variable. To supply headers yourself — offline installs, or a TensorRT
+build the OSS repository does not track — set `TENSORRT_INCLUDE_DIR`, or
+`KERNEL_LENS_FETCH_HEADERS=0` to disable fetching entirely.
+
 **Environment:** `KERNEL_LENS_CACHE_DIR` (put it on a persistent volume so the
 `.so` / `.engine` artifacts survive `kl.load`), `KERNEL_LENS_CCBIN` (host
 compiler for `nvcc` — CUDA toolkits before 12.9 reject GCC 13's exception
 specifiers, so set it to `g++-12` there), `KERNEL_LENS_DEBUG` (or `KL_VERBOSE`),
-`TENSORRT_INCLUDE_DIR`, and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
+`TENSORRT_INCLUDE_DIR` (override the TensorRT headers),
+`KERNEL_LENS_HEADER_CACHE` (where fetched TensorRT headers are kept),
+`KERNEL_LENS_FETCH_HEADERS=0` (disable fetching),
+`KERNEL_LENS_PLUGIN_API` (`auto`, `v2` or `v3`), and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
 above 2 GB from disk rather than from a memory buffer (opt-in; see the
 [changelog](CHANGELOG.md)).
 
@@ -196,6 +231,13 @@ compiled = kl.compile(RMSNorm(4096).cuda(), (x,), name="llama3_rmsnorm",
                       backends=["onnx", "tensorrt"])
 out_ort = compiled.run((x,), backend="onnx")
 out_trt = compiled.run((x,), backend="tensorrt")
+
+# One engine for any batch: mark the leading dim dynamic. dynamic_max bounds
+# the optimization profile (default 4x the traced size).
+serving = kl.compile(RMSNorm(4096).cuda(), (x,), name="llama3_rmsnorm_dyn",
+                     backends=["onnx", "tensorrt"], dynamic_shapes=True)
+for b in (1, 2, 4, 8, 16):                 # all through the engine built at b=4
+    serving.run((torch.randn(b, 512, 4096, device="cuda"),), backend="tensorrt")
 ```
 
 ## Results at a glance
@@ -212,6 +254,16 @@ and every raw JSON are on the
 it replaces is **1,137 lines across six files** for one fused RMSNorm kernel on
 both backends, measured with `tests/measure_codegen.py`.
 
+**One engine, many shapes.** With `dynamic_shapes=True` a single compiled engine
+serves any leading-dimension size in range. Measured with
+`tests/test_conformance.py`, which compiles once and then runs every shape through
+that engine — `0.00e+00` on both backends, on Blackwell and on Ampere:
+
+| Case | Compiled at | Served |
+| :--- | :--- | :--- |
+| `dynamic_batch_sweep` | batch 8 | batches 1, 2, 4, 8, 16 |
+| `dynamic_grid_sweep` | 524,288 elements | 65,536 / 131,072 / 524,288 / 1,048,576 |
+
 **Real pretrained models.** Every decoder `LayerNorm` of a real OPT checkpoint
 (125M → 6.7B) replaced by the Triton RMSNorm kernel, graded against
 `torch.compile` — the baseline that *captures the same Triton kernels* rather

{kernel_lens-1.2.1 → kernel_lens-2.0.0}/README.md

@@ -50,6 +50,18 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
 - **Dynamic grids in C++** — launch-grid expressions are transpiled to C++
   integer math over runtime shapes, so one plugin serves dynamic batch and
   sequence lengths.
+- **One package, several TensorRT versions** — the plugin interface is chosen
+  from the environment. TensorRT 10 deprecated the whole `IPluginV2*` hierarchy
+  and TensorRT 11 removed it, so `plugin_api="auto"` (the default) emits
+  `IPluginV3` plugins where the runtime supports them and `IPluginV2DynamicExt`
+  plugins otherwise. This spans TensorRT 8.6 through 10.x today; TensorRT 11+
+  works when you point `TENSORRT_INCLUDE_DIR` at matching headers.
+- **Shape-generic engines** — `kl.compile(..., dynamic_shapes=True)` exports the
+  leading dimension as a symbolic axis, so a single engine serves any size in a
+  declared range via a TensorRT optimization profile. Verified on Blackwell
+  (`sm_120a`) and Ampere (`sm_80`): one compile at batch 8 serving batches
+  1-16, bit-exact on both engines. Off by default, because a shape-specialised
+  engine is what the published benchmarks measure.
 - **Zero-copy VRAM binding** — PyTorch `data_ptr()` values are registered
   directly with each engine; no host staging copies.
 - **Blackwell PTX ISA clamping** — a Triton-generated `.version 9.3` header is
@@ -87,11 +99,34 @@ pip install kernel-lens[all]     # both
 Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`. TensorRT must be `< 11` (the
 generator emits `IPluginV2DynamicExt`).
 
+**TensorRT versions.** Handled automatically. The pip TensorRT wheels ship **no
+headers**, so KernelLens fetches the set matching the installed runtime from the
+TensorRT OSS repository and caches it under `~/.kernel_lens_headers/<major>.<minor>`;
+when the headers it already has match the runtime major, nothing is fetched. If
+the download is unavailable the fallback is the bundled set, and a
+header/runtime major mismatch is then reported instead of crashing inside
+TensorRT. `plugin_api="auto"` picks the interface from the runtime:
+
+| TensorRT | `IPluginV2*` | `IPluginV3` | emitted |
+| :--- | :--- | :--- | :--- |
+| < 10 (8.6, 9.x) | present | absent | `v2` |
+| 10.x | deprecated | present | `v3` |
+| 11+ | **removed** | present | `v3` |
+
+Verified on TensorRT 10.16 and 11.3, both 16/16 on the conformance suite. Force
+a choice with `plugin_api="v2"` / `"v3"`, or the `KERNEL_LENS_PLUGIN_API`
+environment variable. To supply headers yourself — offline installs, or a TensorRT
+build the OSS repository does not track — set `TENSORRT_INCLUDE_DIR`, or
+`KERNEL_LENS_FETCH_HEADERS=0` to disable fetching entirely.
+
 **Environment:** `KERNEL_LENS_CACHE_DIR` (put it on a persistent volume so the
 `.so` / `.engine` artifacts survive `kl.load`), `KERNEL_LENS_CCBIN` (host
 compiler for `nvcc` — CUDA toolkits before 12.9 reject GCC 13's exception
 specifiers, so set it to `g++-12` there), `KERNEL_LENS_DEBUG` (or `KL_VERBOSE`),
-`TENSORRT_INCLUDE_DIR`, and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
+`TENSORRT_INCLUDE_DIR` (override the TensorRT headers),
+`KERNEL_LENS_HEADER_CACHE` (where fetched TensorRT headers are kept),
+`KERNEL_LENS_FETCH_HEADERS=0` (disable fetching),
+`KERNEL_LENS_PLUGIN_API` (`auto`, `v2` or `v3`), and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
 above 2 GB from disk rather than from a memory buffer (opt-in; see the
 [changelog](CHANGELOG.md)).
 
@@ -132,6 +167,13 @@ compiled = kl.compile(RMSNorm(4096).cuda(), (x,), name="llama3_rmsnorm",
                       backends=["onnx", "tensorrt"])
 out_ort = compiled.run((x,), backend="onnx")
 out_trt = compiled.run((x,), backend="tensorrt")
+
+# One engine for any batch: mark the leading dim dynamic. dynamic_max bounds
+# the optimization profile (default 4x the traced size).
+serving = kl.compile(RMSNorm(4096).cuda(), (x,), name="llama3_rmsnorm_dyn",
+                     backends=["onnx", "tensorrt"], dynamic_shapes=True)
+for b in (1, 2, 4, 8, 16):                 # all through the engine built at b=4
+    serving.run((torch.randn(b, 512, 4096, device="cuda"),), backend="tensorrt")
 ```
 
 ## Results at a glance
@@ -148,6 +190,16 @@ and every raw JSON are on the
 it replaces is **1,137 lines across six files** for one fused RMSNorm kernel on
 both backends, measured with `tests/measure_codegen.py`.
 
+**One engine, many shapes.** With `dynamic_shapes=True` a single compiled engine
+serves any leading-dimension size in range. Measured with
+`tests/test_conformance.py`, which compiles once and then runs every shape through
+that engine — `0.00e+00` on both backends, on Blackwell and on Ampere:
+
+| Case | Compiled at | Served |
+| :--- | :--- | :--- |
+| `dynamic_batch_sweep` | batch 8 | batches 1, 2, 4, 8, 16 |
+| `dynamic_grid_sweep` | 524,288 elements | 65,536 / 131,072 / 524,288 / 1,048,576 |
+
 **Real pretrained models.** Every decoder `LayerNorm` of a real OPT checkpoint
 (125M → 6.7B) replaced by the Triton RMSNorm kernel, graded against
 `torch.compile` — the baseline that *captures the same Triton kernels* rather

{kernel_lens-1.2.1 → kernel_lens-2.0.0}/kernel_lens/__init__.py

@@ -7,7 +7,7 @@ from .config import set_verbose, is_verbose
 
 
 
-__version__ = "1.2.1"
+__version__ = "2.0.0"
 
 __all__ = [
     "compile",

{kernel_lens-1.2.1 → kernel_lens-2.0.0}/kernel_lens/backends/builder.py

@@ -249,6 +249,21 @@ def build_trt_plugin(trt_plugins_dir: str, cache_dir: str):
             arch_flag = "-gencode=arch=compute_75,code=sm_75"
 
         trt_inc_dirs = []
+
+        # Headers matching the installed TensorRT, fetched and cached on first
+        # use when the bundled set is a different major version. See
+        # utils/trt_headers.py for why this is necessary: the pip TensorRT
+        # wheels ship no headers at all.
+        try:
+            from ..utils.trt_headers import resolve as _resolve_headers
+            _hdr_dir, _hdr_reason = _resolve_headers()
+            from ..config import debug_print as _dp
+            _dp(f"[KernelLens] TensorRT headers: {_hdr_dir or 'bundled'} ({_hdr_reason})")
+            if _hdr_dir:
+                trt_inc_dirs.append(_hdr_dir)
+        except Exception:
+            pass
+
         if os.environ.get("TENSORRT_INCLUDE_DIR"):
             trt_inc_dirs.append(os.environ["TENSORRT_INCLUDE_DIR"])
         user_trt_inc = os.path.expanduser("~/tensorrt_headers")

{kernel_lens-1.2.1 → kernel_lens-2.0.0}/kernel_lens/backends/ort_gen.py

@@ -150,31 +150,26 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
         arg_setup_lines.append("auto ref_in = ctx.GetInput(0);")
         arg_setup_lines.append("auto info = ref_in.GetTensorTypeAndShapeInfo();")
         arg_setup_lines.append("std::vector<int64_t> dim_values = info.GetShape();")
+
+        # Dims the caller asked to keep dynamic: read them from the tensor that
+        # is actually being run, not from the traced shape.
+        _dyn_axes = tuple(getattr(manifest, "dynamic_axes", ()) or ())
         
         import re
         fn_match = re.search(r'\.entry\s+([a-zA-Z0-9_]+)', manifest.ptx)
         ptx_entry_name = fn_match.group(1) if fn_match else manifest.kernel_name
         
-        # --- ROBUST GRID EVALUATION WITH GENERIC REGEX SUBSTITUTION ---
-        grid_strs = []
-        if hasattr(manifest, '_sym_grid_asts') and manifest._sym_grid_asts:
-            for g in manifest._sym_grid_asts:
-                expr = str(g.node.expr) if hasattr(g, 'node') else str(g)
-                expr = re.sub(r'([a-zA-Z0-9_]+)\*\*([a-zA-Z0-9_]+)', r'std::pow(\1, \2)', expr)
-                expr = expr.replace("//", "/")
-                expr = re.sub(r'floor\((.*?)\)', r'(\1)', expr)
-                
-                # Replace any s(\d+) symbol dynamically using regex
-                expr = re.sub(r'\bs(\d+)\b', r"(int64_t)(dim_values.size() > \1 ? dim_values[\1] : 1)", expr)
-                grid_strs.append(expr)
-                
+        # --- GRID EVALUATION, from the manifest's own symbol translation ---
+        grid_strs = manifest.grid_exprs(
+            lambda i: f"(int64_t)(dim_values.size() > {i} ? dim_values[{i}] : 1)"
+        )
         while len(grid_strs) < 3:
             grid_strs.append("1")
             
         grid_eval_lines = [
             f"unsigned int grid_x = std::max(1u, (unsigned int)({grid_strs[0]}));",
             f"unsigned int grid_y = std::max(1u, (unsigned int)({grid_strs[1]}));",
-            f"unsigned int grid_z = std::max(1u, (unsigned int)({grid_strs[2]}));"
+            f"unsigned int grid_z = std::max(1u, (unsigned int)({grid_strs[2]}));",
         ]
 
         entry_match = re.search(r'\.entry\s+[a-zA-Z0-9_]+\s*\((.*?)\)\s*(?:\.reqntid|\{)', manifest.ptx, re.DOTALL)
@@ -215,6 +210,23 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
             
             ptx_ordered_slots.append({"type": c_type, "ident": str(param_num), "is_ptr": is_ptr, "decl": p_decl, "param_num": param_num})
 
+        # The node's inputs, in the order ONNX Runtime presents them. This is
+        # the same list GetInputTypeCount() reports, so indexing into it is
+        # authoritative. A running counter over manifest.arguments is not: any
+        # drift between the two silently points a scalar read at the wrong
+        # slot, and the kernel then masks against a garbage element count.
+        inputs_to_node = [
+            a for a in manifest.arguments
+            if (a.kind in ('input', 'inplace')
+                or (a.kind == 'scalar' and not getattr(a, 'is_constexpr', False)))
+        ]
+
+        def _node_input_index(a, fallback):
+            for k, cand in enumerate(inputs_to_node):
+                if cand is a:
+                    return k
+            return fallback
+
         arg_setup_lines.append("// Extract inputs/outputs/scalars from ORT context")
         onnx_input_counter = 0
         ort_output_counter = 0
@@ -227,7 +239,7 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
                     if _esz == 0.5
                     else f"(in_tensor_{slot_idx}.GetTensorTypeAndShapeInfo().GetElementCount() * {int(_esz)})"
                 )
-                arg_setup_lines.append(f"auto in_tensor_{slot_idx} = ctx.GetInput({onnx_input_counter});")
+                arg_setup_lines.append(f"auto in_tensor_{slot_idx} = ctx.GetInput({_node_input_index(arg, onnx_input_counter)});")
                 arg_setup_lines.append(f"static thread_local const void* arg_ptr_{slot_idx};")
                 arg_setup_lines.append(f"const void* raw_ptr_{slot_idx} = (const void*)in_tensor_{slot_idx}.GetTensorData<uint8_t>();")
                 arg_setup_lines.append(f"static thread_local void* align_buf_{slot_idx} = nullptr;")
@@ -242,6 +254,9 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
             elif arg.kind == 'output':
                 out_dim_exprs = []
                 for dim_idx, d in enumerate(arg.shape):
+                    if dim_idx in _dyn_axes:
+                        out_dim_exprs.append(f"(int64_t)(dim_values.size() > {dim_idx} ? dim_values[{dim_idx}] : 1)")
+                        continue
                     try:
                         d_int = int(d)
                         out_dim_exprs.append(str(d_int))
@@ -257,6 +272,9 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
             elif arg.kind == 'inplace':
                 out_dim_exprs = []
                 for dim_idx, d in enumerate(arg.shape):
+                    if dim_idx in _dyn_axes:
+                        out_dim_exprs.append(f"(int64_t)(dim_values.size() > {dim_idx} ? dim_values[{dim_idx}] : 1)")
+                        continue
                     try:
                         d_int = int(d)
                         out_dim_exprs.append(str(d_int))
@@ -264,7 +282,7 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
                         out_dim_exprs.append(f"(int64_t)(dim_values.size() > {dim_idx} ? dim_values[{dim_idx}] : 1)")
                 out_shape_str = "{" + ", ".join(out_dim_exprs) + "}" if out_dim_exprs else "dim_values"
                 
-                arg_setup_lines.append(f"auto in_tensor_{slot_idx} = ctx.GetInput({onnx_input_counter});")
+                arg_setup_lines.append(f"auto in_tensor_{slot_idx} = ctx.GetInput({_node_input_index(arg, onnx_input_counter)});")
                 arg_setup_lines.append(f"std::vector<int64_t> out_dims_{ort_output_counter} = {out_shape_str};")
                 arg_setup_lines.append(f"auto out_tensor_{slot_idx} = ctx.GetOutput({ort_output_counter}, out_dims_{ort_output_counter}.data(), out_dims_{ort_output_counter}.size());")
                 _esz = element_size_bytes(arg.dtype)
@@ -282,16 +300,33 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
                 onnx_input_counter += 1
                 ort_output_counter += 1
             elif arg.kind == 'scalar':
-                if not getattr(arg, 'is_constexpr', False):
-                    arg_setup_lines.append(f"auto scalar_tensor_{slot_idx} = ctx.GetInput({onnx_input_counter});")
+                expr = getattr(arg, 'cxx_expr', '') or ''
+                if 'dim_values' in expr:
+                    # Shape-derived (n = x.numel()): compute it from the tensor
+                    # being run. The exported graph folds this scalar into a
+                    # constant -- it is not a graph input -- so reading the slot
+                    # would hand the kernel the element count captured while
+                    # tracing, and it would then process a fraction of the
+                    # tensor and leave the rest untouched.
+                    arg_setup_lines.append(f"double scalar_val_{slot_idx} = (double)({expr});")
+                    onnx_input_counter += 1
+                elif not getattr(arg, 'is_constexpr', False):
+                    arg_setup_lines.append(f"auto scalar_tensor_{slot_idx} = ctx.GetInput({_node_input_index(arg, onnx_input_counter)});")
                     arg_setup_lines.append(f"double scalar_val_{slot_idx} = 0;")
                     arg_setup_lines.append(f"if (scalar_tensor_{slot_idx}.GetTensorTypeAndShapeInfo().GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) {{ scalar_val_{slot_idx} = (double)(scalar_tensor_{slot_idx}.GetTensorData<int64_t>()[0]); }}")
                     arg_setup_lines.append(f"else if (scalar_tensor_{slot_idx}.GetTensorTypeAndShapeInfo().GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32) {{ scalar_val_{slot_idx} = (double)(scalar_tensor_{slot_idx}.GetTensorData<int32_t>()[0]); }}")
                     arg_setup_lines.append(f"else {{ scalar_val_{slot_idx} = (double)(scalar_tensor_{slot_idx}.GetTensorData<float>()[0]); }}")
                     onnx_input_counter += 1
                 else:
-                    s_val = arg.value if getattr(arg, 'value', None) is not None else 0.0
-                    arg_setup_lines.append(f"double scalar_val_{slot_idx} = {s_val};")
+                    # Prefer the manifest's translated expression: for a scalar
+                    # derived from shapes (n = x.numel()) that is a runtime read
+                    # of the input dims; otherwise it is the baked constant.
+                    expr = getattr(arg, 'cxx_expr', '') or None
+                    if expr:
+                        arg_setup_lines.append(f"double scalar_val_{slot_idx} = (double)({expr});")
+                    else:
+                        s_val = arg.value if getattr(arg, 'value', None) is not None else 0.0
+                        arg_setup_lines.append(f"double scalar_val_{slot_idx} = {s_val};")
 
         arg_setup_lines.append("static thread_local const void* null_ptr_arg = nullptr;")
         arg_setup_lines.append("static thread_local std::vector<void*> kp;")

{kernel_lens-1.2.1 → kernel_lens-2.0.0}/kernel_lens/backends/trt_gen.py

@@ -28,37 +28,31 @@ class TensorRTPluginGenerator:
         self.plugin_namespace = plugin_namespace
         self.plugin_version = plugin_version
         
-    def _generate_kernel_h(self, manifest: KernelManifest) -> str:
-        plugin_name = f"{manifest.kernel_name}Plugin"
-
-        constant_decls = []
-        for arg in manifest.arguments:
-            if arg.kind == 'scalar':
-                ctype = "float" if isinstance(arg.value, float) else "int"
-                constant_decls.append(f"{ctype} m_{arg.name};")
-        dynamic_members_cpp = "\n    ".join(constant_decls)
-
-        def get_trt_datatype(dtype_obj):
-            dtype_str = str(dtype_obj).lower()
-            if 'float8' in dtype_str or 'fp8' in dtype_str or 'e4m3' in dtype_str or 'e5m2' in dtype_str:
-                return "nvinfer1::DataType::kFP8"
-            elif 'bool' in dtype_str:
-                return "nvinfer1::DataType::kBOOL"
-            elif 'uint8' in dtype_str or 'int8' in dtype_str or 'int4' in dtype_str or 'uint4' in dtype_str:
-                return "nvinfer1::DataType::kINT8"
-            elif 'bfloat16' in dtype_str or 'bf16' in dtype_str:
-                return "nvinfer1::DataType::kBF16"
-            elif 'float16' in dtype_str or 'half' in dtype_str:
-                return "nvinfer1::DataType::kHALF"
-            elif 'float' in dtype_str:
-                return "nvinfer1::DataType::kFLOAT"
-            elif 'int64' in dtype_str or 'long' in dtype_str:
-                return "nvinfer1::DataType::kINT64"
-            elif 'int' in dtype_str:
-                return "nvinfer1::DataType::kINT32"
-            else:
-                return "nvinfer1::DataType::kFLOAT"
-
+    @staticmethod
+    def trt_datatype(dtype_obj: str) -> str:
+        """Map a torch dtype onto the TensorRT DataType."""
+        dtype_str = str(dtype_obj).lower()
+        if 'float8' in dtype_str or 'fp8' in dtype_str or 'e4m3' in dtype_str or 'e5m2' in dtype_str:
+            return "nvinfer1::DataType::kFP8"
+        elif 'bool' in dtype_str:
+            return "nvinfer1::DataType::kBOOL"
+        elif 'uint8' in dtype_str or 'int8' in dtype_str or 'int4' in dtype_str or 'uint4' in dtype_str:
+            return "nvinfer1::DataType::kINT8"
+        elif 'bfloat16' in dtype_str or 'bf16' in dtype_str:
+            return "nvinfer1::DataType::kBF16"
+        elif 'float16' in dtype_str or 'half' in dtype_str:
+            return "nvinfer1::DataType::kHALF"
+        elif 'float' in dtype_str:
+            return "nvinfer1::DataType::kFLOAT"
+        elif 'int64' in dtype_str or 'long' in dtype_str:
+            return "nvinfer1::DataType::kINT64"
+        elif 'int' in dtype_str:
+            return "nvinfer1::DataType::kINT32"
+        else:
+            return "nvinfer1::DataType::kFLOAT"
+    @staticmethod
+    def supports_format_body(manifest: KernelManifest, tensor_prefix: str = "") -> str:
+        """The supportsFormatCombination body, shared by both plugin emitters."""
         in_args = [a for a in manifest.arguments if a.kind in ('input', 'inplace')]
 
         # Enforce the dtype the kernel was compiled for, position by position.
@@ -67,7 +61,7 @@ class TensorRTPluginGenerator:
         _ordered_in = [a for a in manifest.arguments if a.kind in ('input', 'inplace')]
         _ordered_sc = [a for a in manifest.arguments if a.kind == 'scalar' and not getattr(a, 'is_constexpr', False)]
         _ordered_out = [a for a in manifest.arguments if a.kind in ('output', 'inplace')]
-        _pos_types = [get_trt_datatype(a.dtype) for a in (_ordered_in + _ordered_sc + _ordered_out)]
+        _pos_types = [TensorRTPluginGenerator.trt_datatype(a.dtype) for a in (_ordered_in + _ordered_sc + _ordered_out)]
         _pos_is_tensor = ([True] * len(_ordered_in) + [False] * len(_ordered_sc) + [True] * len(_ordered_out))
         _types_cpp = ", ".join(_pos_types) if _pos_types else "nvinfer1::DataType::kFLOAT"
         _is_tensor_cpp = ", ".join("true" if t else "false" for t in _pos_is_tensor) if _pos_is_tensor else "true"
@@ -96,14 +90,37 @@ class TensorRTPluginGenerator:
             )
         else:
             supports_format_cxx = type_guard_cxx + "return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;"
-        
+        # IPluginV3 hands supportsFormatCombination a DynamicPluginTensorDesc,
+        # whose fields live under .desc; V2 hands it a PluginTensorDesc directly.
+        if tensor_prefix:
+            supports_format_cxx = supports_format_cxx.replace("inOut[pos].", f"inOut[pos]{tensor_prefix}.")
+        return supports_format_cxx
+
+    def _generate_kernel_h(self, manifest: KernelManifest) -> str:
+        plugin_name = f"{manifest.kernel_name}Plugin"
+
+        constant_decls = []
+        for arg in manifest.arguments:
+            if arg.kind == 'scalar':
+                ctype = "float" if isinstance(arg.value, float) else "int"
+                constant_decls.append(f"{ctype} m_{arg.name};")
+        dynamic_members_cpp = "\n    ".join(constant_decls)
+
+        get_trt_datatype = TensorRTPluginGenerator.trt_datatype
+        supports_format_cxx = TensorRTPluginGenerator.supports_format_body(manifest)
         output_dim_cases = []
         out_args = [a for a in manifest.arguments if a.kind in ('output', 'inplace')]
+        _dyn_axes = tuple(getattr(manifest, "dynamic_axes", ()) or ())
         for out_idx, out_arg in enumerate(out_args):
             output_dim_cases.append(f"if (outputIndex == {out_idx}) {{")
             output_dim_cases.append(f"    nvinfer1::DimsExprs res;")
             output_dim_cases.append(f"    res.nbDims = {len(out_arg.shape)};")
             for d_idx, d in enumerate(out_arg.shape):
+                if d_idx in _dyn_axes:
+                    # Declared dynamic: take the dim from the input at run time
+                    # instead of freezing the traced value.
+                    output_dim_cases.append(f"    res.d[{d_idx}] = inputs[0].d[{d_idx}];")
+                    continue
                 try:
                     d_int = int(d)
                     output_dim_cases.append(f"    res.d[{d_idx}] = exprBuilder.constant({d_int});")
@@ -211,30 +228,27 @@ private:
 '''
         return textwrap.dedent(tpl).strip()
 
-    def _generate_kernel_cu(self, manifest: KernelManifest) -> str:
-        plugin_name = f"{manifest.kernel_name}Plugin"
-        ptx_str = manifest.ptx
-        ptx_encoded = json.dumps(ptx_str)
-        block_size = (manifest.num_warps if manifest.num_warps > 0 else 4) * 32
-        
+    @staticmethod
+    def build_shared_body(manifest: KernelManifest, block_size: int):
+        """Build the pieces the enqueue implementation needs.
+
+        Shared verbatim between the IPluginV2DynamicExt and the
+        IPluginV3 emitters: the PTX entry name, the kernelParams
+        assembly and the output-type dispatch are the same; only the
+        class hierarchy and the method signatures around them differ.
+        """
         scalars = [arg for arg in manifest.arguments if arg.kind == 'scalar']
         nb_outputs = sum(1 for arg in manifest.arguments if arg.kind == 'output')
         nb_outputs = max(1, nb_outputs)
         
         # --- LOCAL TRT GRID AST EVALUATION ---
         
-        grid_strs = []
-        if hasattr(manifest, '_sym_grid_asts') and manifest._sym_grid_asts:
-            for g in manifest._sym_grid_asts:
-                expr = str(g.node.expr) if hasattr(g, 'node') else str(g)
-                expr = re.sub(r'([a-zA-Z0-9_]+)\*\*([a-zA-Z0-9_]+)', r'std::pow(\1, \2)', expr)
-                expr = expr.replace("//", "/")
-                expr = re.sub(r'floor\((.*?)\)', r'(\1)', expr)
-                
-                # Replace any s(\d+) symbol dynamically using regex
-                expr = re.sub(r'\bs(\d+)\b', r"(int64_t)(inputDesc[0].dims.nbDims > \1 ? inputDesc[0].dims.d[\1] : 1)", expr)
-                grid_strs.append(expr)
-                
+        _trt_dim = lambda i: f"(int64_t)(inputDesc[0].dims.nbDims > {i} ? inputDesc[0].dims.d[{i}] : 1)"  # noqa: E731
+        grid_strs = manifest.grid_exprs(_trt_dim)
+        # A scalar derived from the shape (n = x.numel()) must be recomputed per
+        # enqueue from the runtime input, or the kernel masks against the traced
+        # element count and silently does half the work.
+        _scalar_exprs = manifest.scalar_exprs(_trt_dim)
         while len(grid_strs) < 3:
             grid_strs.append("1")
             
@@ -356,7 +370,7 @@ private:
                     arg_setup_lines.append(f"kernelParams[{slot_idx}] = (void*)&tmp_null_{slot_idx};")
             else:
                 if arg and arg.kind == 'scalar':
-                    expr = getattr(arg, 'cxx_expr', '') or f"m_{arg.name}"
+                    expr = _scalar_exprs.get(arg.name) or getattr(arg, 'cxx_expr', '') or f"m_{arg.name}"
                     scalar_ctype = "int64_t" if "64" in str(arg.dtype) else ("float" if ("float" in str(arg.dtype) or "float" in c_type or "double" in c_type) else "int32_t")
                     arg_setup_lines.append(f"static thread_local {scalar_ctype} tmp_scalar_{slot_idx}; tmp_scalar_{slot_idx} = ({scalar_ctype})({expr});")
                     arg_setup_lines.append(f"kernelParams[{slot_idx}] = (void*)&tmp_scalar_{slot_idx};")
@@ -388,10 +402,41 @@ private:
                 
             output_type_lines.append(f"if (index == {i}) return {trt_type};")
             
-        dynamic_output_types = "\n    ".join(output_type_lines)
+        return {
+            "ptx_entry_name": ptx_entry_name,
+            "dynamic_args_cpp": "\n    ".join(arg_setup_lines),
+            "dynamic_output_types": "\n    ".join(output_type_lines),
+            "grid_x_cxx": grid_x_cxx,
+            "grid_y_cxx": grid_y_cxx,
+            "grid_z_cxx": grid_z_cxx,
+            "num_ptx_slots": num_ptx_slots,
+            "init_str": init_str,
+            "size_additions": size_additions,
+            "deserialize_cpp": deserialize_cpp,
+            "serialize_cpp": serialize_cpp,
+            "clone_cpp": clone_cpp,
+        }
 
-        dynamic_args_cpp = "\n    ".join(arg_setup_lines)
+    def _generate_kernel_cu(self, manifest: KernelManifest) -> str:
+        plugin_name = f"{manifest.kernel_name}Plugin"
+        ptx_str = manifest.ptx
+        ptx_encoded = json.dumps(ptx_str)
+        block_size = (manifest.num_warps if manifest.num_warps > 0 else 4) * 32
         
+        shared = TensorRTPluginGenerator.build_shared_body(manifest, block_size)
+        ptx_entry_name = shared["ptx_entry_name"]
+        dynamic_args_cpp = shared["dynamic_args_cpp"]
+        dynamic_output_types = shared["dynamic_output_types"]
+        grid_x_cxx = shared["grid_x_cxx"]
+        grid_y_cxx = shared["grid_y_cxx"]
+        grid_z_cxx = shared["grid_z_cxx"]
+        num_ptx_slots = shared["num_ptx_slots"]
+        init_str = shared["init_str"]
+        size_additions = shared["size_additions"]
+        deserialize_cpp = shared["deserialize_cpp"]
+        serialize_cpp = shared["serialize_cpp"]
+        clone_cpp = shared["clone_cpp"]
+
         tpl = f'''
 #include "{manifest.kernel_name}Plugin.h"
 #include <cuda.h>