Complete definition of a computational workload. A Definition provides a formal, machine-readable specification for a computational workload. It defines the tensor formats, dimension semantics, and computational logic through a reference implementation. This serves as the single source
| 134 | """Union type representing all possible axis specifications.""" |
| 135 | |
| 136 | class Definition(BaseModelWithDocstrings): |
| 137 | """Complete definition of a computational workload. |
| 138 | |
| 139 | A Definition provides a formal, machine-readable specification for a computational |
| 140 | workload. It defines the tensor formats, dimension semantics, and computational |
| 141 | logic through a reference implementation. This serves as the single source of |
| 142 | truth for kernel development and optimization. |
| 143 | """ |
| 144 | |
| 145 | name: NonEmptyString |
| 146 | """A unique, human-readable name for the kernel definition.""" |
| 147 | op_type: Optional[NonEmptyString] = Field(default=None) |
| 148 | """The general compute category.""" |
| 149 | axes: dict[NonEmptyString, AxisSpec] |
| 150 | """Dictionary of symbolic dimensions used in tensor shapes. The axes will be bound to the |
| 151 | input tensor dimensions at runtime.""" |
| 152 | custom_inputs_entrypoint: Optional[NonEmptyString] = Field(default=None) |
| 153 | """The entrypoint function to generate the inputs. The signature should follow entrypoint(axes_and_scalars: dict[str, int], device: torch.device) -> dict[str, torch.Tensor]""" |
| 154 | inputs: dict[NonEmptyString, TensorSpec] |
| 155 | """Named input tensors required by this kernel. The order of inputs is preserved as the |
| 156 | kernel's interface.""" |
| 157 | outputs: dict[NonEmptyString, TensorSpec] |
| 158 | """Named output tensors produced by this kernel. The names of the output must not overlap |
| 159 | with the names of the inputs. The order of outputs is preserved as the kernel's interface.""" |
| 160 | reference: NonEmptyString |
| 161 | """Reference implementation code. It defines the compute logic of the kernel. Must be a valid |
| 162 | Python code with a 'run' function that takes the input tensors and returns the output tensors. |
| 163 | """ |
| 164 | description: Optional[str] = Field(default=None) |
| 165 | """Optional human-readable description of the kernel's purpose.""" |
| 166 | hf_id: Optional[NonEmptyString] = Field(default=None) |
| 167 | """Optional HuggingFace model ID that the definition was sourced from.""" |
| 168 | |
| 169 | @model_validator(mode="after") |
| 170 | def _validate_reference_code(self) -> Definition: |
| 171 | """Validate that reference contains valid Python code with a 'run' function. |
| 172 | |
| 173 | Raises |
| 174 | ------ |
| 175 | ValueError |
| 176 | If the reference code is not valid Python syntax or doesn't contain |
| 177 | a top-level 'run' function. |
| 178 | """ |
| 179 | try: |
| 180 | mod = ast.parse(self.reference, mode="exec") |
| 181 | except SyntaxError as e: |
| 182 | raise ValueError(f"Reference must be valid Python code: {e}") from e |
| 183 | |
| 184 | # Check for 'run' function |
| 185 | has_run_func = any( |
| 186 | isinstance(node, ast.FunctionDef) and node.name == "run" |
| 187 | for node in mod.body |
| 188 | ) |
| 189 | if not has_run_func: |
| 190 | raise ValueError("Reference must define a top-level function named 'run'") |
| 191 | return self |
| 192 | |
| 193 | @model_validator(mode="after") |
no outgoing calls