xtc-tools 0.3.0__py3-none-any.whl

xtc/__init__.py

@@ -0,0 +1,7 @@
+#
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024-2026 The XTC Project Authors
+#
+import importlib.metadata
+
+__version__ = importlib.metadata.version("xtc-tools")

xtc/artifacts/__init__.py

@@ -0,0 +1,25 @@
+#
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024-2026 The XTC Project Authors
+#
+from .operations import (
+    register_operation,
+    get_operation,
+    list_operations,
+    has_operation,
+)
+
+from .register_ttile_operations import _register_operations
+from .register_subgraph_operations import (
+    _register_operations as _register_subgraph_operations,
+)
+
+_register_operations()
+_register_subgraph_operations()
+
+__all__ = [
+    "register_operation",
+    "get_operation",
+    "list_operations",
+    "has_operation",
+]

xtc/artifacts/operations.py

@@ -0,0 +1,71 @@
+#
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024-2026 The XTC Project Authors
+#
+import itertools
+import logging
+
+_OPERATION_REGISTRY: dict[str, dict[str, dict]] = {}
+
+__all__ = [
+    "register_operation",
+    "get_operation",
+    "has_operation",
+    "list_operations",
+]
+
+logger = logging.getLogger(__name__)
+
+
+def register_operation(
+    operator: str, name: str, dims: dict[str, int], params: dict[str, int] | None = None
+):
+    if operator not in _OPERATION_REGISTRY:
+        _OPERATION_REGISTRY[operator] = {}
+    if params is None:
+        params = {}
+    operations = _OPERATION_REGISTRY[operator]
+    canonical = name.lower()
+    if canonical in operations:
+        logger.warning(f"operation {operator}/{canonical} is already registered")
+    _OPERATION_REGISTRY[operator][canonical] = dict(
+        name=name,
+        dims=dims,
+        params=params,
+    )
+
+
+def get_operation(operator: str, name: str) -> dict:
+    if operator not in _OPERATION_REGISTRY:
+        raise ValueError(f"operator {operator} not registered in operation registry")
+    operations = _OPERATION_REGISTRY[operator]
+    canonical = name.lower()
+    if canonical not in operations:
+        raise ValueError(
+            f"operation name {name} for operator {operator} not registered in operation registry"
+        )
+    return operations[canonical]
+
+
+def has_operation(operator: str, name: str) -> bool:
+    try:
+        get_operation(operator, name)
+    except ValueError:
+        return False
+    return True
+
+
+def list_operations(operator: str = "") -> list[tuple[str, str]]:
+    if operator == "":
+        operations = list(
+            itertools.chain(
+                *[list_operations(operator) for operator in _OPERATION_REGISTRY.keys()]
+            )
+        )
+        return operations
+    if operator not in _OPERATION_REGISTRY:
+        raise ValueError(f"operator {operator} not registered in operation registry")
+    operations = [
+        (operator, op["name"]) for op in _OPERATION_REGISTRY[operator].values()
+    ]
+    return operations

xtc/artifacts/register_subgraph_operations.py

@@ -0,0 +1,34 @@
+#
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024-2026 The XTC Project Authors
+#
+from .operations import register_operation
+from .subgraph_sizes import *
+
+__all__: list[str] = []
+
+
+def _register_conv2d_ops():
+    for group in [squeezenet_convs, alexnet_convs]:
+        for name, params in group.items():
+            register_operation(
+                "conv2d",
+                name,
+                {k: params[k] for k in ["n", "h", "w", "f", "r", "s", "c"]},
+                {"SH": params["stry"], "SW": params["strx"]},
+            )
+
+
+def _register_matmul_ops():
+    for group in [alexnet_matmuls]:
+        for name, params in group.items():
+            register_operation(
+                "matmul",
+                name,
+                {k: params[k] for k in ["i", "j", "k"]},
+            )
+
+
+def _register_operations():
+    _register_conv2d_ops()
+    _register_matmul_ops()

xtc/artifacts/register_ttile_operations.py

@@ -0,0 +1,46 @@
+#
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024-2026 The XTC Project Authors
+#
+from xtc.schedules.ttile.prob_sizes import (
+    ddsizes_Yolo,
+    ddsizes_MobilNet,
+    ddsizes_RN18,
+    ddsizes_matmul,
+)
+
+from .operations import register_operation
+
+__all__: list[str] = []
+
+
+def _register_conv2d_ops():
+    map = dict(
+        h="y",
+        w="x",
+        r="h",
+        s="w",
+    )
+    for group in [ddsizes_Yolo, ddsizes_MobilNet, ddsizes_RN18]:
+        for name, params in group.items():
+            register_operation(
+                "conv2d",
+                name,
+                {k: params[map.get(k, k)] for k in ["n", "h", "w", "f", "r", "s", "c"]},
+                {"SH": params["stry"], "SW": params["strx"]},
+            )
+
+
+def _register_matmul_ops():
+    for group in [ddsizes_matmul]:
+        for name, params in group.items():
+            register_operation(
+                "matmul",
+                name,
+                {k: params[k] for k in ["i", "j", "k"]},
+            )
+
+
+def _register_operations():
+    _register_conv2d_ops()
+    _register_matmul_ops()

xtc/artifacts/subgraph_sizes.py

@@ -0,0 +1,207 @@
+#
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024-2026 The XTC Project Authors
+#
+__all__ = [
+    "squeezenet_convs",
+    "alexnet_convs",
+    "alexnet_matmuls",
+]
+
+squeezenet_convs = {
+    "SqueezeNet_00": {
+        "n": 1,
+        "h": 224,
+        "w": 224,
+        "c": 3,
+        "r": 7,
+        "s": 7,
+        "f": 96,
+        "strx": 2,
+        "stry": 2,
+    },
+    "SqueezeNet_01": {
+        "n": 1,
+        "h": 54,
+        "w": 54,
+        "c": 96,
+        "r": 1,
+        "s": 1,
+        "f": 16,
+        "strx": 1,
+        "stry": 1,
+    },
+    "SqueezeNet_03": {
+        "n": 1,
+        "h": 56,
+        "w": 56,
+        "c": 16,
+        "r": 3,
+        "s": 3,
+        "f": 64,
+        "strx": 1,
+        "stry": 1,
+    },
+    "SqueezeNet_04": {
+        "n": 1,
+        "h": 54,
+        "w": 54,
+        "c": 128,
+        "r": 1,
+        "s": 1,
+        "f": 16,
+        "strx": 1,
+        "stry": 1,
+    },
+    "SqueezeNet_05": {
+        "n": 1,
+        "h": 54,
+        "w": 54,
+        "c": 128,
+        "r": 1,
+        "s": 1,
+        "f": 32,
+        "strx": 1,
+        "stry": 1,
+    },
+    "SqueezeNet_07": {
+        "n": 1,
+        "h": 54,
+        "w": 54,
+        "c": 32,
+        "r": 3,
+        "s": 3,
+        "f": 128,
+        "strx": 1,
+        "stry": 1,
+    },
+    "SqueezeNet_08": {
+        "n": 1,
+        "h": 27,
+        "w": 27,
+        "c": 256,
+        "r": 1,
+        "s": 1,
+        "f": 32,
+        "strx": 1,
+        "stry": 1,
+    },
+    "SqueezeNet_11": {
+        "n": 1,
+        "h": 27,
+        "w": 27,
+        "c": 256,
+        "r": 1,
+        "s": 1,
+        "f": 48,
+        "strx": 1,
+        "stry": 1,
+    },
+    "SqueezeNet_13": {
+        "n": 1,
+        "h": 27,
+        "w": 27,
+        "c": 48,
+        "r": 3,
+        "s": 3,
+        "f": 192,
+        "strx": 1,
+        "stry": 1,
+    },
+    "SqueezeNet_14": {
+        "n": 1,
+        "h": 27,
+        "w": 27,
+        "c": 384,
+        "r": 1,
+        "s": 1,
+        "f": 48,
+        "strx": 1,
+        "stry": 1,
+    },
+    "SqueezeNet_17": {
+        "n": 1,
+        "h": 27,
+        "w": 27,
+        "c": 64,
+        "r": 3,
+        "s": 3,
+        "f": 256,
+        "strx": 1,
+        "stry": 1,
+    },
+    "SqueezeNet_21": {
+        "n": 1,
+        "h": 15,
+        "w": 15,
+        "c": 512,
+        "r": 1,
+        "s": 1,
+        "f": 1000,
+        "strx": 1,
+        "stry": 1,
+    },
+}
+
+alexnet_convs = {
+    "AlexNet_00": {
+        "n": 1,
+        "h": 228,
+        "w": 228,
+        "c": 3,
+        "r": 11,
+        "s": 11,
+        "f": 64,
+        "strx": 4,
+        "stry": 4,
+    },
+    "AlexNet_01": {
+        "n": 1,
+        "h": 31,
+        "w": 31,
+        "c": 64,
+        "r": 5,
+        "s": 5,
+        "f": 192,
+        "strx": 1,
+        "stry": 1,
+    },
+    "AlexNet_02": {
+        "n": 1,
+        "h": 15,
+        "w": 15,
+        "c": 192,
+        "r": 3,
+        "s": 3,
+        "f": 384,
+        "strx": 1,
+        "stry": 1,
+    },
+    "AlexNet_03": {
+        "n": 1,
+        "h": 15,
+        "w": 15,
+        "c": 384,
+        "r": 3,
+        "s": 3,
+        "f": 256,
+        "strx": 1,
+        "stry": 1,
+    },
+    "AlexNet_04": {
+        "n": 1,
+        "h": 15,
+        "w": 15,
+        "c": 256,
+        "r": 3,
+        "s": 3,
+        "f": 256,
+        "strx": 1,
+        "stry": 1,
+    },
+}
+alexnet_matmuls = {
+    "AlexNet_FC6": {"i": 1, "j": 4096, "k": 9216},
+    "AlexNet_FC7": {"i": 1, "j": 4096, "k": 4096},
+    "AlexNet_FC8": {"i": 1, "j": 1000, "k": 4096},
+}

xtc/backends/jir/JIRBackend.py

@@ -0,0 +1,105 @@
+#
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024-2026 The XTC Project Authors
+#
+from typing_extensions import override
+from typing import Any
+
+from jir.node import JIRFunction
+from jir.backend.util.annotate_fastmath import annotate_fastmath
+from jir.parser import JIRParser
+from jir.backend.xdsl.compiler import PolygeistCompiler
+
+from xtc.utils.tools import (
+    get_geist_prefix,
+)
+
+import xtc.itf as itf
+from xtc.itf.graph import Graph
+from xtc.graphs.xtc.graph import XTCGraph
+
+from .JIROps import JIROperation
+from .JIRScheduler import JIRScheduler
+from .JIRCompiler import JIRCompiler
+
+__all__ = [
+    "JIRBackend",
+]
+
+
+class JIRBackend(itf.back.Backend):
+    def __init__(
+        self,
+        source_op: JIROperation | Graph,
+        dims: dict[str, int] | None = None,
+        parallel_dims: list[str] | None = None,
+        reduction_dims: list[str] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self._graph: Graph | None = None
+        if isinstance(source_op, XTCGraph):
+            graph = source_op
+            self._graph = graph
+            self.ops = [
+                JIROperation.from_operation(node.operation, name=node.name)
+                for node in graph.nodes.values()
+            ]
+            self.dims = self.ops[-1].operator.dims_sizes()
+            self._payload_name = self._graph.name
+        else:
+            assert isinstance(source_op, JIROperation)
+            assert dims is not None
+            self.dims = dims
+            self.ops = [source_op]
+            self._payload_name = source_op.name
+
+        self.op = self.ops[-1]
+
+        assert tuple(self.dims.keys()) == self.op.operator.dims(), (
+            f"incompatible dims names: {tuple(self.dims.keys())} != "
+            f"{self.op.operator.dims()}"
+        )
+        self.parallel_dims = self.op.operator.dims("P")
+        self.reduction_dims = self.op.operator.dims("R")
+        if parallel_dims is not None:
+            assert tuple(parallel_dims) == self.parallel_dims, (
+                f"incompatible parallel dims names: {tuple(parallel_dims)} != "
+                f"{self.parallel_dims}"
+            )
+        if reduction_dims is not None:
+            assert tuple(reduction_dims) == self.reduction_dims, (
+                f"incompatible reduction dims names: {tuple(reduction_dims)} != "
+                f"{self.reduction_dims}"
+            )
+        self._geist_install_dir = get_geist_prefix()
+        self._op_function_str, self._jir_function_str = self.op.generate(
+            self._payload_name
+        )
+        self._jir_function_op = self._parse_function(self._jir_function_str)
+        self._op_function_mlir = self._parse_primitives(self._op_function_str)
+
+    @override
+    def get_scheduler(self, **kwargs: Any) -> itf.schd.Scheduler:
+        return JIRScheduler(self, **kwargs)
+
+    @override
+    def get_compiler(self, **kwargs: Any) -> itf.comp.Compiler:
+        return JIRCompiler(self, **kwargs)
+
+    @property
+    @override
+    def payload_name(self) -> str:
+        return self._payload_name
+
+    @property
+    @override
+    def graph(self) -> itf.graph.Graph:
+        assert self._graph is not None
+        return self._graph
+
+    def _parse_function(self, jir_function: str) -> JIRFunction:
+        return JIRParser().parse_function(jir_function)
+
+    def _parse_primitives(self, op_function: str) -> str:
+        polygeist_compiler = PolygeistCompiler(f"{self._geist_install_dir}/bin/cgeist")
+        return annotate_fastmath(polygeist_compiler(op_function))