()
| 37 | |
| 38 | |
| 39 | def _default_globals() -> dict[str, Any]: |
| 40 | # lazy import here to avoid circular deps |
| 41 | from tvm.script import tirx as _tirx_dsl # pylint: disable=import-outside-toplevel |
| 42 | from tvm.script.parser import ( |
| 43 | ir, # pylint: disable=import-outside-toplevel |
| 44 | relax, # pylint: disable=import-outside-toplevel |
| 45 | ) |
| 46 | from tvm.script.parser import tirx as _tirx_parser # pylint: disable=import-outside-toplevel |
| 47 | from tvm.script.tirx import tile as _tirx_tile # pylint: disable=import-outside-toplevel |
| 48 | from tvm.tirx import layout as _tirx_layout # pylint: disable=import-outside-toplevel |
| 49 | |
| 50 | # Expose the layout `Axis` class so printed layout sugar like |
| 51 | # `4 @ Axis.laneid` round-trips without per-script imports. Injecting just |
| 52 | # `Axis` (one short symbol) avoids name collisions with common user shape |
| 53 | # vars like `m`, `P`, `F` that registered axes happen to share names with. |
| 54 | return { |
| 55 | "tvm": tvm, |
| 56 | "I": ir, |
| 57 | "ir": ir, |
| 58 | "T": _tirx_parser, |
| 59 | "tir": _tirx_parser, |
| 60 | "R": relax, |
| 61 | "relax": relax, |
| 62 | "Tx": _tirx_tile, |
| 63 | "tirx": _tirx_dsl, |
| 64 | "Axis": _tirx_layout.Axis, |
| 65 | } |
| 66 | |
| 67 | |
| 68 | def scan_macro(program: Any | str, extra_vars: dict[str, Any] | None = None) -> Any: |
no outgoing calls
no test coverage detected
searching dependent graphs…