Smoke test integration of XTC in Pytorch

pyproject.toml

@@ -47,7 +47,7 @@ build-backend = "setuptools.build_meta"
 
 [tool.pyright]
 pythonVersion = "3.10"
-include = ["src/xtc/itf", "src/xtc/backends", "src/xtc/graphs", "src/xtc/search", "src/xtc/runtimes", "docs/marimo"]
+include = ["src/xtc/itf", "src/xtc/backends", "src/xtc/graphs", "src/xtc/search", "src/xtc/runtimes", "src/xtc/integration", "docs/marimo"]
 reportMissingParameterType = "error"
 reportImplicitOverride = "error"
 reportUninitializedInstanceVariable = "error"
@@ -66,7 +66,7 @@ reportMissingImports = false
 
 [tool.mypy]
 python_version = "3.10"
-packages = ["xtc.itf", "xtc.backends", "xtc.graphs", "xtc.search", "xtc.runtimes", "docs.marimo"]
+packages = ["xtc.itf", "xtc.backends", "xtc.graphs", "xtc.search", "xtc.runtimes", "xtc.integration", "docs.marimo"]
 ignore_missing_imports = true
 check_untyped_defs = true
 

src/xtc/backends/mlir/MlirGraphBackend.py

@@ -152,12 +152,18 @@ def _xdsl_elt_shape_from_tensortype(self, type: XTCTensorType) -> tuple[Any, Any
 
     def _xdsl_type_from_tensortype(self, type: XTCTensorType) -> Any:
         elt_type, shape = self._xdsl_elt_shape_from_tensortype(type)
+        layout = type.layout
+        if layout is not None:
+            shape = [shape[idx] for idx in layout]
         return MemRefType(elt_type, shape)
 
     def _xdsl_attrs_from_tensortype(self, type: XTCTensorType):
+        attrs = {}
         if type.device is not None:
-            return {"memref.on_device": UnitAttr()}
-        return {}
+            attrs["memref.on_device"] = UnitAttr()
+        if type.const:
+            attrs["memref.const"] = UnitAttr()
+        return attrs
 
     def _np_types_spec(
         self, types: list[MemRefType]

src/xtc/backends/mlir/MlirOps.py

@@ -23,6 +23,7 @@
 from xdsl.builder import ImplicitBuilder
 
 from xtc.itf.graph import Operation
+from xtc.graphs.xtc.data import XTCTensorType
 from xtc.utils.math import mulall
 
 
@@ -35,6 +36,20 @@
 OpAttrs: TypeAlias = dict[str, Any]
 
 
+def _is_natural_layout(layout: list[int] | None, ndim: int) -> bool:
+    return layout is None or layout == list(range(ndim))
+
+
+def _require_natural_layouts(xtc_op: Operation) -> None:
+    for typ in xtc_op.inputs_types:
+        if isinstance(typ, XTCTensorType) and not _is_natural_layout(
+            typ.layout, typ.ndim
+        ):
+            raise NotImplementedError(
+                "tensor layout is not yet implemented in MLIR backend"
+            )
+
+
 class MlirOperation:
     def __init__(
         self,
@@ -83,6 +98,7 @@ def from_operation(cls, xtc_op: Operation, name: str | None) -> "MlirOperation":
         dtype = xtc_op.inputs_types[0].dtype  # TODO: currently get dtype from 1st arg
         args = tuple([*dims, dtype])
         attrs = xtc_op.attrs
+        args = MlirOperators.from_name(xtc_op.name).args_from_operation(xtc_op, args)
         return MlirOperation(
             MlirOperators.from_name(xtc_op.name),
             args,
@@ -125,26 +141,42 @@ def _dims(self, kind: str = "") -> tuple[str, ...]:
             return tuple(self.AXES)
         return tuple([a for a, k in zip(self.AXES, self.KINDS) if k == kind])
 
+    @classmethod
+    def args_from_operation(
+        cls, xtc_op: Operation, args: tuple[Any, ...]
+    ) -> tuple[Any, ...]:
+        _require_natural_layouts(xtc_op)
+        return args
+
 
 class MlirOperatorMatmul(MlirOperator):
     DEFAULT_NAME = "matmul"
     AXES = "ijk"
     KINDS = "PPR"
 
+    @classmethod
+    @override
+    def args_from_operation(
+        cls, xtc_op: Operation, args: tuple[Any, ...]
+    ) -> tuple[Any, ...]:
+        transpose_a = xtc_op.inputs_types[0].layout == [1, 0]
+        transpose_b = xtc_op.inputs_types[1].layout == [1, 0]
+        return (*args, transpose_a, transpose_b)
+
     @override
     def dims(self, kind: str = "") -> tuple[str, ...]:
         return self._dims(kind)
 
     @override
     def dims_sizes(self) -> dict[str, int]:
-        i, j, k, _ = self.args
+        i, j, k, _, _, _ = self.args
         return {"i": i, "j": j, "k": k}
 
     @override
     def generate_op(
         self, block: Block | None = None, args: Sequence[BlockArgument] = []
     ) -> tuple[Block, OpAttrs]:
-        Ki, Kj, Kk, dtype = self.args
+        Ki, Kj, Kk, dtype, transpose_a, transpose_b = self.args
         elt_type = {"float32": f32, "float64": f64}[dtype]
         elt_size = {"float32": 32, "float64": 64}[dtype]
         if block is None:
@@ -162,11 +194,43 @@ def generate_op(
                 inputs=(cst0.results[0],),
                 outputs=(args[2],),
             )
-            reduce = linalg.MatmulOp(
-                res=(),
-                inputs=(args[0], args[1]),
-                outputs=(args[2],),
-            )
+            if not transpose_a and not transpose_b:
+                reduce = linalg.MatmulOp(
+                    res=(),
+                    inputs=(args[0], args[1]),
+                    outputs=(args[2],),
+                )
+            else:
+                iterator_types = [
+                    StringAttr("parallel"),
+                    StringAttr("parallel"),
+                    StringAttr("reduction"),
+                ]
+                index_map_a = lambda i, j, k: (i, k)
+                index_map_b = lambda i, j, k: (k, j)
+                index_map_c = lambda i, j, k: (i, j)
+                if transpose_a:
+                    index_map_a = lambda i, j, k: (k, i)
+                if transpose_b:
+                    index_map_b = lambda i, j, k: (j, k)
+                elt_type = {"float32": f32, "float64": f64}[dtype]
+                block_in = Block(arg_types=[elt_type, elt_type, elt_type])
+                with ImplicitBuilder(block_in):
+                    mul = arith.MulfOp(block_in.args[0], block_in.args[1])
+                    add = arith.AddfOp(block_in.args[2], mul)
+                    linalg.YieldOp(add)
+                reduce = linalg.GenericOp(
+                    inputs=(args[0], args[1]),
+                    outputs=(args[2],),
+                    body=Region([block_in]),  # type: ignore # mypy issue with dataclass
+                    # ignore typing due to xdsl hints limitation
+                    indexing_maps=[
+                        AffineMapAttr(AffineMap.from_callable(index_map_a)),
+                        AffineMapAttr(AffineMap.from_callable(index_map_b)),
+                        AffineMapAttr(AffineMap.from_callable(index_map_c)),
+                    ],
+                    iterator_types=iterator_types,
+                )
         fill_node_id = f"{self.name}_0"
         reduce_node_id = f"{self.name}"
         fill.attributes[f"__xtc_id_{fill_node_id}_"] = UnitAttr()

src/xtc/backends/mlir/MlirTarget/MlirLLVMTarget.py

@@ -316,6 +316,7 @@ def _lowering_pipeline(self) -> list[str]:
             "convert-vector-to-llvm{enable-x86vector=true}",
             "convert-index-to-llvm",
             "convert-arith-to-llvm",
+            "convert-ub-to-llvm",
             "canonicalize",
             "cse",
             "sccp",

src/xtc/backends/tvm/TVMOps.py

@@ -11,6 +11,7 @@
 from xtc.utils.text import to_cname
 
 from xtc.itf.graph import Operation, Graph, Node
+from xtc.graphs.xtc.data import XTCTensorType
 
 __all__ = [
     "TVMBaseExpr",
@@ -49,6 +50,10 @@ def lower(self, tensors: Sequence[TETensor], sch: te.Schedule) -> str:
 
     @classmethod
     def from_operation(cls, xtc_op: Operation, name: str | None) -> "TVMOperation":
+        for typ in xtc_op.inputs_types:
+            if isinstance(typ, XTCTensorType):
+                if typ.layout is not None:
+                    assert False, "tensor layout is not yet implemented in TVM backend"
         dims = xtc_op.dims.values()
         dtype = xtc_op.inputs_types[0].dtype  # TODO: infer dtype form first input
         args = tuple([*dims, dtype])
@@ -145,6 +150,9 @@ def _te_shape_dtype_from_node(self, node: Node) -> Any:
         else:
             assert hasattr(node, "_expr")
             type = node._expr.type  # type: ignore
+        assert isinstance(type, XTCTensorType)
+        if type.layout is not None:
+            assert False, "tensor layout is not yet implemented in TVM backend"
         return type.shape, type.dtype
 
     def _te_tensor_from_node(self, node: Node) -> Any:

src/xtc/graphs/xtc/data.py

@@ -34,10 +34,14 @@ def __init__(
         shape: ShapeType = None,
         dtype: DataType = None,
         device: AcceleratorDevice | None = None,
+        const: bool = False,
+        layout: list[int] | None = None,
     ):
         self._shape = shape
         self._dtype = dtype
         self._device = device
+        self._const = const
+        self._layout = layout
 
     @property
     @override
@@ -54,6 +58,16 @@ def dtype(self) -> DataType:
     def device(self) -> AcceleratorDevice | None:
         return self._device
 
+    @property
+    @override
+    def const(self) -> bool:
+        return self._const
+
+    @property
+    @override
+    def layout(self) -> list[int] | None:
+        return self._layout
+
     @property
     @override
     def ndim(self) -> int:
@@ -104,6 +118,9 @@ def constant(self) -> "XTCConstantTensorType":
         return XTCConstantTensorType(
             shape=self.constant_shape,
             dtype=self.constant_dtype,
+            device=self.device,
+            const=self.const,
+            layout=self.layout,
         )
 
     @override
@@ -113,13 +130,26 @@ def __repr__(self) -> str:
         else:
             dims = "x".join([str(d if d else "?") for d in self._shape])
         dtype = self._dtype if self._dtype else "?"
-        return f"{dims}x{dtype}"
+        const = ", const" if self.const else ""
+        layout = ""
+        if self.layout is not None:
+            assert self._shape is not None
+            layout += (
+                ", <(" + ",".join([str(i) for i in range(len(self._shape))]) + ")->("
+            )
+            layout += ",".join([str(i) for i in self.layout]) + ")>"
+        return f"{dims}x{dtype}{const}{layout}"
 
     @override
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, XTCTensorType):
             return NotImplemented
-        return self.dtype == other.dtype and self.shape == other.shape
+        return (
+            self.dtype == other.dtype
+            and self.shape == other.shape
+            and self.const == other.const
+            and self.layout == other.layout
+        )
 
     @override
     def to_dict(self) -> dict[str, Any]:
@@ -140,9 +170,19 @@ def from_dict(cls, tensor_dict: dict[str, Any]) -> Self:
 
 
 class XTCConstantTensorType(XTCTensorType, ConstantTensorType):
-    def __init__(self, shape: ConstantShapeType, dtype: ConstantDataType):
+    def __init__(
+        self,
+        shape: ConstantShapeType,
+        dtype: ConstantDataType,
+        device: AcceleratorDevice | None = None,
+        const: bool = False,
+        layout: list[int] | None = None,
+    ):
         self._shape: ConstantShapeType = shape
         self._dtype: ConstantDataType = dtype
+        self._device = device
+        self._const = const
+        self._layout = layout
 
     @property
     @override

src/xtc/graphs/xtc/expr.py

@@ -123,26 +123,40 @@ def __init__(
         shape: ShapeType | DataType = None,
         dtype: DataType = None,
         device: AcceleratorDevice | None = None,
+        const: bool = False,
+        layout: list[int] | None = None,
     ) -> None:
         super().__init__()
         if tensor is None:
             assert shape is None or isinstance(shape, tuple)
             assert dtype is None or isinstance(dtype, str)
-            type = XTCTensorType(shape=shape, dtype=dtype, device=device)
+            type = XTCTensorType(
+                shape=shape, dtype=dtype, device=device, const=const, layout=layout
+            )
             value = XTCTensor(type=type)
         elif isinstance(tensor, XTCTensorType):
-            assert shape is None and dtype is None
+            assert shape is None and dtype is None and layout is None
+            assert not const, (
+                "const cannot be set when tensor type is provided directly"
+            )
             value = XTCTensor(type=tensor)
         elif isinstance(tensor, XTCTensor):
-            assert shape is None and dtype is None
+            assert shape is None and dtype is None and layout is None
+            assert not const, (
+                "const cannot be set when tensor value is provided directly"
+            )
             value = tensor
         elif isinstance(tensor, tuple):
             if shape is not None:
                 assert isinstance(shape, str)
                 assert dtype is None
-                type = XTCTensorType(shape=tensor, dtype=shape, device=device)
+                type = XTCTensorType(
+                    shape=tensor, dtype=shape, device=device, const=const, layout=layout
+                )
             else:
-                type = XTCTensorType(shape=tensor, dtype=dtype, device=device)
+                type = XTCTensorType(
+                    shape=tensor, dtype=dtype, device=device, const=const, layout=layout
+                )
             value = XTCTensor(type=type)
         self._value = value
         self._device = device

src/xtc/integration/__init__.py

@@ -0,0 +1,4 @@
+#
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024-2026 The XTC Project Authors
+#

src/xtc/integration/pytorch/__init__.py

@@ -0,0 +1,20 @@
+#
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright (c) 2024-2026 The XTC Project Authors
+#
+
+from __future__ import annotations
+
+from xtc.integration.pytorch.compile import compile_matmul_kernel
+from xtc.integration.pytorch.torch_xtc import (
+    XtcIntegration,
+    register_torch_xtc_extensions,
+)
+from xtc.integration.pytorch.eager import clear_kernel_cache
+
+__all__ = [
+    "XtcIntegration",
+    "clear_kernel_cache",
+    "compile_matmul_kernel",
+    "register_torch_xtc_extensions",
+]