# mypy: allow-untyped-defs import functools import itertools import logging from collections.abc import Iterable from typing import Any from unittest.mock import patch from torch._inductor.utils import Placeholder, unique from torch._inductor.virtualized import V from ...autotune_process import CuteDSLBenchmarkRequest, TensorMeta from ...ir import Buffer, ChoiceCaller, CuteDSLTemplateBuffer, IRNode, Layout, TensorBox from ..common import KernelTemplate from .cutedsl_kernel import CuteDSLTemplateKernel log = logging.getLogger(__name__) class CuteDSLTemplate(KernelTemplate): """Template for generating CuteDSL (CUTLASS Python DSL) kernels. Subclasses may override ``caller_type`` to attach template-specific metadata or precompile behavior while reusing the common render and benchmark request construction. """ kernel_type: type[Any] = CuteDSLTemplateKernel caller_type: type[Any] | None = None index_counter = itertools.count() all_templates: dict[str, "CuteDSLTemplate"] = {} def __init__( self, name: str, source: str, subgraph_fn: Any | None = None, mask_fn: Any | None = None, ) -> None: super().__init__(name) self.source = source self.subgraph_fn = subgraph_fn self.mask_fn = mask_fn self.template = CuteDSLTemplate._template_from_string(source) # A module that registers templates can be initialized more than once in # a single process (e.g. a double-import path). Tolerate re-registration # under an existing name as long as the template source matches, but # reject a genuine name collision between different templates. existing = self.all_templates.get(name) if existing is not None and existing.source != self.source: raise AssertionError(f"duplicate template name, {name}") CuteDSLTemplate.all_templates[name] = self @staticmethod @functools.lru_cache(None) # pyrefly: ignore [bad-override] def _template_from_string(source: str) -> Any: return KernelTemplate._template_from_string(source) def maybe_append_choice( self, choices: list[Any], **kwargs: Any ) -> NotImplementedError | None: """ Maybe generates a new ChoiceCaller and appends it into existing choices. Returns None if success, otherwise returns the error. """ try: choices.append(self.generate(**kwargs)) return None except NotImplementedError as e: log.debug("CuteDSL template choice generation failed: %s", e) return e except Exception as e: log.debug("CuteDSL template choice generation error: %s", e) return NotImplementedError(f"CuteDSL template failed: {e}") def generate(self, **kwargs: Any) -> ChoiceCaller: """Generate the CuteDSL kernel caller for template autotuning.""" input_nodes = kwargs.pop("input_nodes") layout = kwargs.pop("layout") mutated_inputs = kwargs.pop("mutated_inputs", None) subgraphs = kwargs.pop("subgraphs", None) template_kwargs = dict(kwargs) kernel_name = f"cutedsl_{self.name}_{next(self.index_counter)}" if self.template is None: raise RuntimeError("Template compilation failed (Jinja2 required)") self.output_node: Buffer = Buffer(name="buf_out", layout=layout) # Patch V.graph.get_dtype to handle the fake buf_out buffer with patch.object( V.graph, "get_dtype", KernelTemplate._fake_get_dtype(self.output_node) ): kernel = self.kernel_type( kernel_name=kernel_name, input_nodes=input_nodes, output_node=self.output_node, subgraphs=subgraphs, ) code = kernel.render(self.template, **kwargs) input_call_args = tuple(kernel.args.input_buffers.keys()) expected_input_args = tuple(unique(x.get_name() for x in input_nodes)) if input_call_args[: len(expected_input_args)] != expected_input_args: raise RuntimeError( "CuteDSL template input registration order changed while " "collecting captured subgraph buffers. Expected template " "inputs to be registered before captured buffers, got " f"{input_call_args}, expected prefix {expected_input_args}." ) extra_capture_names = input_call_args[len(expected_input_args) :] # Resolve captured nodes from the graph-level side table # (populated by realize_captures_for_cutedsl) to get view nodes. graph_captures = getattr(V.graph, "_cutedsl_capture_nodes", {}) capture_nodes_by_name: dict[str, Any] = {} extra_capture_nodes = [] for name in extra_capture_names: node = graph_captures.get(name) if node is None: node = V.graph.get_buffer(name) capture_nodes_by_name[name] = node extra_capture_nodes.append(node) input_nodes = list(input_nodes) + extra_capture_nodes kernel.set_capture_input_nodes(capture_nodes_by_name) with kernel._patch_get_dtype_for_captures(): _, call_args, _, _ = kernel.args.python_argdefs() expected_args = list(input_call_args) expected_args.append(self.output_node.get_name()) if list(call_args)[: len(expected_args)] != expected_args: raise RuntimeError( "CuteDSL template benchmark argument order changed while " "collecting dynamic scalar args. Expected prefix " f"{expected_args}, got {list(call_args)}." ) extra_args = tuple( V.graph.sizevars.optimization_hints(call_args[len(expected_args) :]) ) bmreq = CuteDSLBenchmarkRequest( kernel_name=kernel_name, input_tensor_meta=TensorMeta.from_irnodes(input_nodes), output_tensor_meta=TensorMeta.from_irnodes(self.output_node), extra_args=extra_args, source_code=code, ) def make_kernel_render(out_node, hint_override: int | None = None): """ Factory function that creates a kernel renderer for the final output. This closure captures the current template and parameters, but allows the output node to be specified later. This is used during the final kernel selection phase when the actual output buffer is available. """ render_kernel = self.kernel_type( kernel_name=str(Placeholder.KERNEL_NAME), input_nodes=input_nodes, output_node=out_node, subgraphs=subgraphs, ) render_kernel.set_capture_input_nodes(capture_nodes_by_name) def render(): return render_kernel.render(self.template, **kwargs) return render_kernel, render caller_type = self.caller_type or CuteDSLTemplateCaller return caller_type( name=kernel_name, input_nodes=input_nodes, layout=layout, make_kernel_render=make_kernel_render, bmreq=bmreq, template=self, mutated_inputs=mutated_inputs, template_kwargs=template_kwargs, ) class CuteDSLTemplateCaller(ChoiceCaller): """Caller for CuteDSL templates that integrates with the autotuning system.""" def __init__( self, name: str, input_nodes: list[Buffer], layout: Layout, make_kernel_render: Any, bmreq: CuteDSLBenchmarkRequest, template: "CuteDSLTemplate", mutated_inputs: Iterable[IRNode] | None = None, template_kwargs: dict[str, Any] | None = None, ): description = self._build_description(name, template_kwargs) super().__init__( name=name, input_nodes=input_nodes, layout=layout, description=description, ) self.make_kernel_render = make_kernel_render self.bmreq = bmreq self.template = template self.mutated_inputs = mutated_inputs def _build_description( self, name: str, template_kwargs: dict[str, Any] | None ) -> str: if not template_kwargs: return f"CuteDSL template {name}" kwargs_desc = ", ".join(f"{k}={v}" for k, v in template_kwargs.items()) return f"CuteDSL template {name} ({kwargs_desc})" def __str__(self) -> str: return f"CuteDSLTemplateCaller({self.name})" def benchmark(self, *args, out) -> float: """Benchmark the kernel execution.""" return self.bmreq.benchmark(*args, out=out) def output_node(self) -> TensorBox: """Create the output node for this template choice.""" buffer = CuteDSLTemplateBuffer( layout=self.layout, inputs=self.input_nodes, make_kernel_render=self.make_kernel_render, template=self.template, mutated_inputs=self.mutated_inputs, ) # Pass KTC annotation to the buffer for encoding if "ktc" in self.annotations: buffer.annotations["ktc"] = self.annotations["ktc"] return TensorBox.create(buffer) def call_name(self) -> str: """Return the kernel call name.""" return self.name def to_callable(self) -> Any: """Return callable that can execute this kernel.""" return self.make_kernel_render def hash_key(self) -> str: """Return unique hash key for this choice.""" return "-".join( [ self.name.rsplit("_", 1)[0], self.bmreq.module_cache_key, ] ) def info_dict(self) -> dict[str, Any]: """Return information about this kernel.""" return { "name": self.name, "backend": "CuteDSL", "template": self.template.name, }