diff --git a/pyproject.toml b/pyproject.toml index 9153f8024..cb72d08c6 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/xtc/backends/mlir/MlirGraphBackend.py b/src/xtc/backends/mlir/MlirGraphBackend.py index ca9650e48..2e5154f6e 100644 --- a/src/xtc/backends/mlir/MlirGraphBackend.py +++ b/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] diff --git a/src/xtc/backends/mlir/MlirOps.py b/src/xtc/backends/mlir/MlirOps.py index be687b444..4e061d1c8 100644 --- a/src/xtc/backends/mlir/MlirOps.py +++ b/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() diff --git a/src/xtc/backends/mlir/MlirTarget/MlirLLVMTarget.py b/src/xtc/backends/mlir/MlirTarget/MlirLLVMTarget.py index 28b4c003d..23a39e4ca 100644 --- a/src/xtc/backends/mlir/MlirTarget/MlirLLVMTarget.py +++ b/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", diff --git a/src/xtc/backends/tvm/TVMOps.py b/src/xtc/backends/tvm/TVMOps.py index 559957c9d..3ef40e975 100644 --- a/src/xtc/backends/tvm/TVMOps.py +++ b/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: diff --git a/src/xtc/graphs/xtc/data.py b/src/xtc/graphs/xtc/data.py index 501495e98..4c0dfdec5 100644 --- a/src/xtc/graphs/xtc/data.py +++ b/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 diff --git a/src/xtc/graphs/xtc/expr.py b/src/xtc/graphs/xtc/expr.py index 4a26c78d8..431101920 100644 --- a/src/xtc/graphs/xtc/expr.py +++ b/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 diff --git a/src/xtc/integration/__init__.py b/src/xtc/integration/__init__.py new file mode 100644 index 000000000..4a40722d1 --- /dev/null +++ b/src/xtc/integration/__init__.py @@ -0,0 +1,4 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# diff --git a/src/xtc/integration/pytorch/__init__.py b/src/xtc/integration/pytorch/__init__.py new file mode 100644 index 000000000..4bb27f92c --- /dev/null +++ b/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", +] diff --git a/src/xtc/integration/pytorch/compile.py b/src/xtc/integration/pytorch/compile.py new file mode 100644 index 000000000..cf96d0f44 --- /dev/null +++ b/src/xtc/integration/pytorch/compile.py @@ -0,0 +1,406 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import dataclass +from math import prod +from pathlib import Path + +import torch +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend +from xtc.itf.comp.module import Module +from xtc.itf.schd.scheduler import Scheduler +from xtc.utils.cfunc import CFunc +from xtc.utils.ext_tools import get_shlib_extension +from xtc.utils.loader import LibLoader + +# --------------------------------------------------------------------------- +# Shared operator compilation utilities +# --------------------------------------------------------------------------- + + +def torch_dtype_to_xtc(dtype: torch.dtype) -> str: + if dtype == torch.float32: + return "float32" + if dtype == torch.float64: + return "float64" + raise ValueError(f"unsupported dtype for XTC operator: {dtype}") + + +def default_kernel_cache_dir() -> Path: + return Path(os.environ.get("XTC_PYTORCH_CACHE", ".xtc_kernel_cache")) + + +def resolve_cache_dir(cache_dir: Path | None) -> Path: + resolved = cache_dir if cache_dir is not None else default_kernel_cache_dir() + resolved.mkdir(parents=True, exist_ok=True) + return resolved + + +def contiguous_strides_for_shape(shape: tuple[int, ...]) -> tuple[int, ...]: + """Row-major contiguous strides for ``shape`` (last dim stride 1).""" + if not shape: + return () + rev: list[int] = [1] + for dim in reversed(shape[1:]): + rev.append(rev[-1] * dim) + return tuple(reversed(rev)) + + +def validate_contiguous_strides( + shape: tuple[int, ...], + stride: tuple[int, ...], + *, + what: str, +) -> None: + """Raise ``NotImplementedError`` when ``stride`` is not contiguous for ``shape``.""" + expected = contiguous_strides_for_shape(shape) + if stride != expected: + raise NotImplementedError( + f"XTC operator requires contiguous {what} " + f"(shape={shape}, stride={stride}, expected stride={expected})" + ) + + +@dataclass(frozen=True) +class KernelArtifacts: + """Paths produced when compiling a specialized XTC operator kernel.""" + + export_name: str + cache_dir: Path + module_path: Path + export_dir: Path + header_path: Path + xtc_lib_path: Path + payload_name: str + shim_lib_path: Path + shim_lib_name: str + + +# --------------------------------------------------------------------------- +# Matmul operator compilation +# --------------------------------------------------------------------------- + + +def validate_matmul_xtc_support( + x_shape: tuple[int, ...], + w_shape: tuple[int, int], + dtype: torch.dtype, + device: torch.device, + *, + b_shape: tuple[int, ...] | None = None, +) -> None: + """Raise ``NotImplementedError`` when XTC cannot specialize this matmul.""" + if device.type != "cpu": + raise NotImplementedError( + f"xtc::matmul: XTC matmul is only supported on CPU (got {device})" + ) + try: + torch_dtype_to_xtc(dtype) + except ValueError as exc: + raise NotImplementedError( + f"xtc::matmul: XTC matmul does not support dtype {dtype}" + ) from exc + if len(x_shape) < 1: + raise NotImplementedError( + f"xtc::matmul: XTC matmul requires x with rank >= 1 (got shape {x_shape})" + ) + if len(w_shape) != 2: + raise NotImplementedError( + f"xtc::matmul: XTC matmul requires 2D weight (got shape {w_shape})" + ) + j, k_w = w_shape + *leading, k = x_shape + if k != k_w: + raise NotImplementedError( + f"xtc::matmul: in_features mismatch (x k={k}, w k={k_w})" + ) + for name, dims in (("x", x_shape), ("w", w_shape)): + if any(d <= 0 for d in dims): + raise NotImplementedError( + f"xtc::matmul: XTC matmul requires positive {name} dimensions " + f"(got {dims})" + ) + if b_shape is not None: + if len(b_shape) != 1 or b_shape[0] != j: + raise NotImplementedError( + f"xtc::matmul: bias must be 1D with size {j} (got {b_shape})" + ) + if b_shape[0] <= 0: + raise NotImplementedError( + f"xtc::matmul: bias must have positive size (got {b_shape})" + ) + _ = leading + + +def validate_matmul_inductor_no_bias(has_bias: bool) -> None: + """Raise ``NotImplementedError`` when the Inductor C++ shim receives a bias.""" + if has_bias: + raise NotImplementedError( + "xtc::matmul: XTC Inductor C++ integration does not support bias yet" + ) + + +def validate_matmul_inductor_layout( + x_shape: tuple[int, ...], + x_stride: tuple[int, ...], + w_shape: tuple[int, int], + w_stride: tuple[int, ...], +) -> None: + """Require row-major contiguous tensor layouts for the Inductor C++ shim.""" + validate_contiguous_strides(x_shape, x_stride, what="x") + validate_contiguous_strides(w_shape, w_stride, what="weight") + + +def matmul_cache_key( + x_shape: tuple[int, ...], + w_shape: tuple[int, int], + dtype: torch.dtype, +) -> tuple[int, int, int, str]: + *leading, k = x_shape + j, k_w = w_shape + if k != k_w: + raise ValueError(f"in_features mismatch: x has {k}, w has {k_w}") + i = prod(leading) if leading else 1 + return (i, j, k, torch_dtype_to_xtc(dtype)) + + +def matmul_export_name(i: int, j: int, k: int, xtc_dtype: str) -> str: + return f"matmul_{i}_{j}_{k}_{xtc_dtype}" + + +def validate_matmul_inductor_specialized_shapes( + x_shape: tuple[int, ...], + w_shape: tuple[int, int], + dtype: torch.dtype, + *, + out_shape: tuple[int, ...] | None = None, +) -> tuple[int, int, int]: + """Require 2D ``(i, k)``, ``(j, k)``, and ``(i, j)`` tensors for the AOTI C++ shim. + + The shim calls a specialized XTC matmul with fixed ``(i, j, k)``; shapes are + checked here at Inductor lowering time instead of in generated C++. + """ + i, j, k, _ = matmul_cache_key(x_shape, w_shape, dtype) + expected_x = (i, k) + expected_w = (j, k) + expected_out = (i, j) + + if x_shape != expected_x: + raise NotImplementedError( + "xtc::matmul: XTC Inductor C++ integration requires 2D activations " + f"with shape {expected_x} (got x shape {x_shape})" + ) + if w_shape != expected_w: + raise NotImplementedError( + "xtc::matmul: XTC Inductor C++ integration requires weight shape " + f"{expected_w} (got {w_shape})" + ) + if out_shape is not None and out_shape != expected_out: + raise NotImplementedError( + "xtc::matmul: XTC Inductor C++ integration requires output shape " + f"{expected_out} (got {out_shape})" + ) + return i, j, k + + +def schedule_matmul(sch: Scheduler) -> None: + sch.strip_mine("i", {"i1": 2}) + sch.strip_mine("j", {"j1": 16}) + sch.interchange(["k", "i", "j", "i1", "j1"]) + sch.vectorize(["j1"]) + sch.unroll({"i1": 2}) + + +@dataclass(frozen=True) +class MatmulKernelArtifacts(KernelArtifacts): + """Paths for a specialized matmul used by Inductor C++ integration.""" + + i: int + j: int + k: int + xtc_dtype: str + + +LinearKernelArtifacts = MatmulKernelArtifacts + + +def build_matmul_module( + x_shape: tuple[int, ...], + w_shape: tuple[int, int], + dtype: torch.dtype, + *, + cache_dir: Path | None = None, + dump_file: str | None = None, +) -> tuple[Module, MatmulKernelArtifacts]: + i, j, k, xtc_dtype = matmul_cache_key(x_shape, w_shape, dtype) + *leading, _ = x_shape + leading_shape = tuple(leading) + + a = O.tensor((i, k), xtc_dtype, name="A") + # PyTorch weights are (j, k); layout marks row-major storage for transpose_b. + b = O.tensor((k, j), xtc_dtype, name="B", layout=[1, 0]) + + with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C") + + resolved_cache = resolve_cache_dir(cache_dir) + export_name = matmul_export_name(i, j, k, xtc_dtype) + if dump_file is None: + dump_file = str(resolved_cache / export_name) + + impl = Backend(gb.graph) + sch = impl.get_scheduler() + schedule_matmul(sch) + sched = sch.schedule() + + comp = impl.get_compiler(shared_lib=True, dump_file=dump_file) + module = comp.compile(sched) + + module_path = Path(module.file_name).resolve() + if not module_path.is_file(): + raise FileNotFoundError( + f"XTC shared library not found after compile: {module_path}" + ) + + export_dir = resolved_cache / f"export_{export_name}" + header_path = export_dir / "include" / f"{export_name}.h" + ext = get_shlib_extension() + xtc_lib_path = export_dir / "lib" / f"lib{export_name}.{ext}" + shim_lib_name = f"xtc_matmul_shim_{i}_{j}_{k}_{xtc_dtype}" + shim_lib_path = resolved_cache / f"lib{shim_lib_name}.{ext}" + + artifacts = MatmulKernelArtifacts( + i=i, + j=j, + k=k, + xtc_dtype=xtc_dtype, + export_name=export_name, + cache_dir=resolved_cache, + module_path=module_path, + export_dir=export_dir, + header_path=header_path, + xtc_lib_path=xtc_lib_path, + payload_name=module.payload_name, + shim_lib_path=shim_lib_path, + shim_lib_name=shim_lib_name, + ) + _ = leading_shape # used by compile_matmul_kernel closure + return module, artifacts + + +def compile_matmul_at_lowering( + x_shape: tuple[int, ...], + w_shape: tuple[int, int], + dtype: torch.dtype, + *, + inductor_cpp: bool = False, + cache_dir: Path | None = None, + force: bool = False, + x_stride: tuple[int, ...] | None = None, + w_stride: tuple[int, ...] | None = None, + out_shape: tuple[int, ...] | None = None, + has_bias: bool = False, +) -> MatmulKernelArtifacts: + """Compile the XTC matmul during Inductor lowering (export AOTI shim if needed).""" + validate_matmul_xtc_support( + x_shape, + w_shape, + dtype, + torch.device("cpu"), + ) + if inductor_cpp: + validate_matmul_inductor_no_bias(has_bias) + validate_matmul_inductor_specialized_shapes( + x_shape, w_shape, dtype, out_shape=out_shape + ) + if x_stride is None or w_stride is None: + raise NotImplementedError( + "xtc::matmul: XTC Inductor integration requires static " + "strides at compile time" + ) + validate_matmul_inductor_layout( + x_shape, + x_stride, + w_shape, + w_stride, + ) + from xtc.integration.pytorch.inductor_cpp import ( + ensure_inductor_matmul_artifacts, + ) + + return ensure_inductor_matmul_artifacts( + x_shape, w_shape, dtype, cache_dir=cache_dir, force=force + ) + _, artifacts = build_matmul_module(x_shape, w_shape, dtype, cache_dir=cache_dir) + return artifacts + + +def export_matmul_kernel( + x_shape: tuple[int, ...], + w_shape: tuple[int, int], + dtype: torch.dtype, + *, + cache_dir: Path | None = None, + force: bool = False, +) -> MatmulKernelArtifacts: + """Compile (if needed), export C++ link tree, and build the Inductor AOTI shim.""" + module, artifacts = build_matmul_module( + x_shape, w_shape, dtype, cache_dir=cache_dir + ) + + if force or not artifacts.header_path.is_file(): + artifacts.export_dir.mkdir(parents=True, exist_ok=True) + module.export(artifacts.export_dir, name=artifacts.export_name) + + if force or not artifacts.shim_lib_path.is_file(): + from xtc.integration.pytorch.inductor_cpp import build_matmul_aoti_shim + + build_matmul_aoti_shim(artifacts, force=force) + + return artifacts + + +def compile_matmul_kernel( + x_shape: tuple[int, ...], + w_shape: tuple[int, int], + dtype: torch.dtype, + *, + cache_dir: Path | None = None, + dump_file: str | None = None, +) -> Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None], torch.Tensor]: + """Compile ``x @ w.mT`` (+ optional bias) with XTC and return a callable kernel.""" + i, j, k, xtc_dtype = matmul_cache_key(x_shape, w_shape, dtype) + *leading, _ = x_shape + + module, artifacts = build_matmul_module( + x_shape, w_shape, dtype, cache_dir=cache_dir, dump_file=dump_file + ) + + loader = LibLoader(str(artifacts.module_path)) + func = getattr(loader.lib, module.payload_name) + func.packed = not getattr(module, "_bare_ptr", True) + cfunc = CFunc(func) + + leading_shape = tuple(leading) + + def kernel( + x: torch.Tensor, w: torch.Tensor, b: torch.Tensor | None = None + ) -> torch.Tensor: + x2d = x.reshape(-1, k).contiguous() + w2d = w.contiguous() + y = torch.empty((*leading_shape, j), dtype=dtype, device=x.device) + y2d = y.reshape(i, j).contiguous() + cfunc(x2d.numpy(), w2d.numpy(), y2d.numpy()) + if b is not None: + y = y + b + return y + + kernel.loader = loader # type: ignore[attr-defined] + kernel.cfunc = cfunc # type: ignore[attr-defined] + return kernel diff --git a/src/xtc/integration/pytorch/eager.py b/src/xtc/integration/pytorch/eager.py new file mode 100644 index 000000000..0962e9484 --- /dev/null +++ b/src/xtc/integration/pytorch/eager.py @@ -0,0 +1,39 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# +from __future__ import annotations + +from collections.abc import Callable + +import torch + +from xtc.integration.pytorch.compile import compile_matmul_kernel + +_matmul_kernel_cache: dict[ + tuple[ + tuple[int, ...], + tuple[int, int], + tuple[int, ...] | None, + torch.dtype, + torch.device, + ], + Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None], torch.Tensor], +] = {} + + +def clear_kernel_cache() -> None: + """Clear eager kernel caches for all XTC PyTorch custom ops.""" + _matmul_kernel_cache.clear() + + +def matmul_cpu( + x: torch.Tensor, w: torch.Tensor, b: torch.Tensor | None = None +) -> torch.Tensor: + x_shape = tuple(x.shape) + w_shape = (int(w.shape[0]), int(w.shape[1])) + b_shape = tuple(b.shape) if b is not None else None + key = (x_shape, w_shape, b_shape, x.dtype, x.device) + if key not in _matmul_kernel_cache: + _matmul_kernel_cache[key] = compile_matmul_kernel(x_shape, w_shape, x.dtype) + return _matmul_kernel_cache[key](x, w, b) diff --git a/src/xtc/integration/pytorch/fx_rewrite.py b/src/xtc/integration/pytorch/fx_rewrite.py new file mode 100644 index 000000000..5fe4cba71 --- /dev/null +++ b/src/xtc/integration/pytorch/fx_rewrite.py @@ -0,0 +1,38 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from torch._inductor.custom_graph_pass import CustomGraphPass + +_LINEAR_TARGETS = frozenset( + { + F.linear, + torch.ops.aten.linear.default, + } +) + + +def replace_linear_pass(graph: torch.fx.Graph) -> None: + for node in graph.nodes: + if node.op == "call_function" and node.target in _LINEAR_TARGETS: + node.target = torch.ops.xtc.matmul.default + graph.eliminate_dead_code() + + +class ReplaceLinearPass(CustomGraphPass): + def __call__(self, graph: torch.fx.Graph) -> None: + replace_linear_pass(graph) + + def uuid(self) -> str: + return "replace_linear_with_xtc_matmul" + + +def register_pre_grad_pass() -> None: + import torch._inductor.config as inductor_config + + inductor_config.pre_grad_custom_pass = ReplaceLinearPass() diff --git a/src/xtc/integration/pytorch/inductor_cpp.py b/src/xtc/integration/pytorch/inductor_cpp.py new file mode 100644 index 000000000..2f2afbc47 --- /dev/null +++ b/src/xtc/integration/pytorch/inductor_cpp.py @@ -0,0 +1,357 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# +from __future__ import annotations + +import os +import subprocess +import sys +import textwrap +from functools import wraps +from pathlib import Path +from collections.abc import Callable, Sequence +from typing import Any, cast + +from xtc.integration.pytorch.compile import ( + MatmulKernelArtifacts, + export_matmul_kernel, +) + +MATMUL_C_SHIM_DECL = textwrap.dedent( + """\ + AOTITorchError + aoti_torch_cpu_matmul_cpp( + AtenTensorHandle out, + AtenTensorHandle x, + AtenTensorHandle w, + AtenTensorHandle* b)""" +) + +_active_artifacts: MatmulKernelArtifacts | None = None +_cpp_wrapper_patch_installed = False + + +def _update_aot_custom_op_libs(artifacts: MatmulKernelArtifacts) -> None: + import torch._inductor.config as inductor_config + + libs = list(inductor_config.aot_inductor.custom_op_libs or []) + if artifacts.shim_lib_name not in libs: + libs.append(artifacts.shim_lib_name) + inductor_config.aot_inductor.custom_op_libs = libs + + cache = str(artifacts.cache_dir.resolve()) + existing = os.environ.get("LIBRARY_PATH", "") + if cache not in existing.split(":"): + os.environ["LIBRARY_PATH"] = f"{cache}:{existing}" if existing else cache + + +def set_active_inductor_artifacts(artifacts: MatmulKernelArtifacts | None) -> None: + global _active_artifacts + _active_artifacts = artifacts + if artifacts is not None: + _update_aot_custom_op_libs(artifacts) + + +def get_active_inductor_link_flags() -> list[str]: + if _active_artifacts is None: + return [] + art = _active_artifacts + shim = str(art.shim_lib_path.resolve()) + lib_dir = str(art.cache_dir.resolve()) + # Link the shim archive directly (same pattern as HalideCodeCache extra_flags). + return [ + shim, + f"-Wl,-rpath,{lib_dir}", + f"-Wl,-rpath,{art.xtc_lib_path.parent.resolve()}", + ] + + +def _torch_build_paths() -> tuple[list[str], list[str], list[str]]: + from torch.utils.cpp_extension import include_paths, library_paths + + includes = include_paths() + lib_dirs = library_paths() + libs = ["torch", "torch_cpu", "c10"] + return includes, lib_dirs, libs + + +def _c_type_for_dtype(xtc_dtype: str) -> str: + if xtc_dtype == "float32": + return "float" + if xtc_dtype == "float64": + return "double" + raise ValueError(f"unsupported xtc dtype for inductor shim: {xtc_dtype}") + + +def _write_generated_shim(artifacts: MatmulKernelArtifacts, src_path: Path) -> None: + c_type = _c_type_for_dtype(artifacts.xtc_dtype) + header = artifacts.export_name + payload = artifacts.payload_name + + content = textwrap.dedent( + f"""\ + #include "{header}.h" + + #include + + namespace {{ + + using Float = {c_type}; + + }} // namespace + + extern "C" {{ + + AOTITorchError aoti_torch_cpu_matmul_cpp( + AtenTensorHandle out, + AtenTensorHandle x, + AtenTensorHandle w, + AtenTensorHandle* b) {{ + + void* x_ptr_v = nullptr; + void* w_ptr_v = nullptr; + void* out_ptr_v = nullptr; + if (aoti_torch_get_data_ptr(x, &x_ptr_v)) return 1; + if (aoti_torch_get_data_ptr(w, &w_ptr_v)) return 1; + if (aoti_torch_get_data_ptr(out, &out_ptr_v)) return 1; + + auto* x_ptr = static_cast(x_ptr_v); + auto* w_ptr = static_cast(w_ptr_v); + auto* out_ptr = static_cast(out_ptr_v); + + {payload}(x_ptr, w_ptr, out_ptr); + + return 0; + }} + + }} // extern "C" + """ + ) + src_path.write_text(content, encoding="utf-8") + + +def build_matmul_aoti_shim( + artifacts: MatmulKernelArtifacts, + *, + force: bool = False, +) -> Path: + if not force and artifacts.shim_lib_path.is_file(): + return artifacts.shim_lib_path + + gen_dir = artifacts.cache_dir / f"shim_build_{artifacts.export_name}" + gen_dir.mkdir(parents=True, exist_ok=True) + src_path = gen_dir / "matmul_aoti_shim.cpp" + _write_generated_shim(artifacts, src_path) + + includes, lib_dirs, libs = _torch_build_paths() + ext = artifacts.shim_lib_path.suffix.lstrip(".") + if sys.platform == "darwin": + shared_flag = "-dynamiclib" + rpath_flag = f"-Wl,-rpath,{artifacts.xtc_lib_path.parent}" + else: + shared_flag = "-shared" + rpath_flag = f"-Wl,-rpath,{artifacts.xtc_lib_path.parent}" + + cmd: list[str] = [ + os.environ.get("CXX", "g++"), + "-O2", + "-std=c++17", + "-fPIC", + shared_flag, + "-o", + str(artifacts.shim_lib_path), + str(src_path), + f"-I{artifacts.export_dir / 'include'}", + f"-L{artifacts.xtc_lib_path.parent}", + f"-l{artifacts.export_name}", + rpath_flag, + ] + for inc in includes: + cmd.append(f"-I{inc}") + for lib_dir in lib_dirs: + cmd.extend(["-L", lib_dir]) + for lib in libs: + cmd.append(f"-l{lib}") + + result = subprocess.run(cmd, capture_output=True, text=True, check=False) + if result.returncode != 0: + raise RuntimeError( + "failed to build XTC matmul AOTI shim\n" + f"command: {' '.join(cmd)}\n" + f"stdout:\n{result.stdout}\n" + f"stderr:\n{result.stderr}" + ) + return artifacts.shim_lib_path + + +def ensure_inductor_matmul_artifacts( + x_shape: tuple[int, ...], + w_shape: tuple[int, int], + dtype: object, + *, + cache_dir: Path | None = None, + force: bool = False, +) -> MatmulKernelArtifacts: + import torch + + if not isinstance(dtype, torch.dtype): + raise TypeError(f"expected torch.dtype, got {type(dtype)!r}") + artifacts = export_matmul_kernel( + x_shape, w_shape, dtype, cache_dir=cache_dir, force=force + ) + set_active_inductor_artifacts(artifacts) + return artifacts + + +ensure_inductor_linear_artifacts = ensure_inductor_matmul_artifacts + + +def _partition_extra_flags( + flags: tuple[str, ...], +) -> tuple[tuple[str, ...], tuple[str, ...]]: + """Split compile flags from link-only flags (-L/-l/-Wl).""" + compile_flags: list[str] = [] + link_flags: list[str] = [] + i = 0 + while i < len(flags): + flag = flags[i] + if flag == "-L" and i + 1 < len(flags): + link_flags.extend([flag, flags[i + 1]]) + i += 2 + continue + if ( + flag.startswith("-L") + or flag.startswith("-l") + or flag.startswith("-Wl,") + or flag.endswith(".so") + or flag.endswith(".a") + ): + link_flags.append(flag) + else: + compile_flags.append(flag) + i += 1 + return tuple(compile_flags), tuple(link_flags) + + +def _patch_cpp_wrapper_code_cache() -> None: + global _cpp_wrapper_patch_installed + if _cpp_wrapper_patch_installed: + return + + from torch._inductor.codecache import CppCodeCache + + _orig_load_async = cast( + Callable[[type[Any], str, str, Any, Sequence[str], str | None], Any], + getattr(CppCodeCache.load_async, "__func__", CppCodeCache.load_async), + ) + + @wraps(_orig_load_async) + def _load_async_impl( + cls: type[Any], + main_code: str, + device_type: str = "cpu", + submit_fn: Any = None, + extra_flags: Sequence[str] = (), + optimized_code: str | None = None, + ) -> Any: + merged = tuple(extra_flags) + tuple(get_active_inductor_link_flags()) + compile_flags, link_flags = _partition_extra_flags(merged) + flags = compile_flags + link_flags + if not link_flags: + return _orig_load_async( + cls, + main_code, + device_type, + submit_fn, + flags, + optimized_code, + ) + + from torch._inductor import codecache as codecache_mod + + orig_precompile = codecache_mod._precompile_header + + def precompile_without_link( + header: str, + hashable_cmd_line: str, + min_optimize: bool = False, + **compile_command: Any, + ) -> str: + cmd = dict(compile_command) + pch_flags, _ = _partition_extra_flags(tuple(cmd.get("extra_flags", ()))) + cmd["extra_flags"] = pch_flags + return orig_precompile( + header, hashable_cmd_line, min_optimize=min_optimize, **cmd + ) + + codecache_mod._precompile_header = precompile_without_link # type: ignore[assignment] + try: + return _orig_load_async( + cls, + main_code, + device_type, + submit_fn, + flags, + optimized_code, + ) + finally: + codecache_mod._precompile_header = orig_precompile + + CppCodeCache.load_async = classmethod(_load_async_impl) # type: ignore[method-assign,assignment] + _cpp_wrapper_patch_installed = True + + +def _declare_matmul_shim_on_wrapper(wrapper: object) -> None: + if getattr(wrapper, "_xtc_matmul_shim_declared", False): + return + wrapper._xtc_matmul_shim_declared = True # type: ignore[attr-defined] + decl = f""" +#include +extern "C" {{ +{MATMUL_C_SHIM_DECL}; +}} +""" + wrapper.prefix._lines.insert(0, decl) # type: ignore[attr-defined] + + +def _patch_cpp_wrapper_extern_call() -> None: + from torch._inductor.codegen.cpp_wrapper_cpu import CppWrapperCpu + + if getattr(CppWrapperCpu, "_xtc_extern_patch_installed", False): + return + CppWrapperCpu._xtc_extern_patch_installed = True # type: ignore[attr-defined] + + orig = CppWrapperCpu.generate_c_shim_extern_kernel_call + + def patched( + self: object, + kernel: str, + args: list[str], + device: str, + **kwargs: object, + ) -> None: + shim = self.get_c_shim_func_name(kernel, device) # type: ignore[attr-defined] + if shim == "aoti_torch_cpu_matmul_cpp": + _declare_matmul_shim_on_wrapper(self) + return orig(self, kernel, args, device, **kwargs) # type: ignore[arg-type] + + CppWrapperCpu.generate_c_shim_extern_kernel_call = patched # type: ignore[method-assign] + + +def register_aot_inductor_cpp_shims() -> None: + """Register C shim declarations for AOT Inductor (use inside config.patch).""" + import torch + import torch._inductor.config as inductor_config + + op = torch.ops.xtc.matmul.default + shims = dict(inductor_config.aot_inductor.custom_ops_to_c_shims) + shims[op] = [MATMUL_C_SHIM_DECL] + inductor_config.aot_inductor.custom_ops_to_c_shims = shims + if _active_artifacts is not None: + _update_aot_custom_op_libs(_active_artifacts) + + +def register_inductor_cpp_hooks() -> None: + _patch_cpp_wrapper_code_cache() + _patch_cpp_wrapper_extern_call() diff --git a/src/xtc/integration/pytorch/inductor_lowering.py b/src/xtc/integration/pytorch/inductor_lowering.py new file mode 100644 index 000000000..06cd6ae24 --- /dev/null +++ b/src/xtc/integration/pytorch/inductor_lowering.py @@ -0,0 +1,206 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# +# pyright: reportFunctionMemberAccess=false +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any, cast + +import torch +from sympy import Expr +from torch._inductor import ir +from torch._inductor.ir import ExternKernelAlloc, ExternKernelOut, TensorBox +from torch._inductor.lowering import lowerings, register_lowering +from torch._inductor.virtualized import V + +from xtc.integration.pytorch.compile import ( + compile_matmul_at_lowering, + validate_matmul_xtc_support, +) + +# --------------------------------------------------------------------------- +# Shared Inductor lowering helpers (reused by future XTC custom ops) +# --------------------------------------------------------------------------- + + +def _static_dim_or_nyi(expr: Expr | int, *, what: str) -> int: + try: + return V.graph.sizevars.size_hint_or_throw(expr) + except Exception as exc: + raise NotImplementedError( + f"XTC operator requires static {what} at compile time" + ) from exc + + +def _static_strides_or_nyi( + tensor: Any, shape: tuple[int, ...], *, what: str +) -> tuple[int, ...]: + stride = tensor.maybe_get_stride() + if stride is None: + raise NotImplementedError( + f"XTC operator requires static {what} strides at compile time" + ) + if len(stride) != len(shape): + raise NotImplementedError( + f"XTC operator {what} stride rank mismatch " + f"(shape={shape}, stride len={len(stride)})" + ) + return tuple(_static_dim_or_nyi(s, what=f"{what} stride") for s in stride) + + +def _register_extern_lowering( + op_overload: Any, + inputs: list[Any], + kwargs: dict[str, Any], + *, + output_size: list[Any], + device: torch.device, + dtype: torch.dtype, + cpp_kernel_name: str | None = None, +) -> Any: + """Build an Inductor ExternKernel node for a registered XTC custom op.""" + output_layout = ir.FixedLayout( + device=device, + dtype=dtype, + size=output_size, + stride=ir.FlexibleLayout.contiguous_strides(cast(Sequence[int], output_size)), + ) + + extern_cls = ExternKernelOut if V.graph.cpp_wrapper else ExternKernelAlloc + extern_kwargs: dict[str, Any] = { + "layout": output_layout, + "inputs": inputs, + "kwargs": kwargs, + "op_overload": op_overload, + } + if extern_cls is ExternKernelOut: + extern_kwargs["output_view"] = None + extern_kwargs["cpp_kernel_name"] = cpp_kernel_name + + extern = extern_cls(**extern_kwargs) + return TensorBox.create(extern) + + +# --------------------------------------------------------------------------- +# Matmul Inductor lowering (xtc::matmul) +# --------------------------------------------------------------------------- + + +def _register_matmul_inductor_lowering() -> None: + @register_lowering(torch.ops.xtc.matmul.default, type_promotion_kind=None) + def matmul_lowering(x: Any, w: Any, b: Any | None = None) -> Any: + x.realize() + w.realize() + if b is not None: + b.realize() + + if V.graph.cpp_wrapper and b is not None: + raise NotImplementedError( + "xtc::matmul: XTC Inductor C++ integration does not support bias yet" + ) + + x_sizes = list(x.get_size()) + k = x_sizes[-1] + x_prefix = x_sizes[:-1] + j, k_w = w.get_size() + V.graph.sizevars.check_equals(k, k_w) + + device = x.get_device() + if device is None: + raise NotImplementedError( + "xtc::matmul: XTC matmul requires a concrete device at compile time" + ) + if device.type != "cpu": + raise NotImplementedError( + f"xtc::matmul: XTC matmul is only supported on CPU (got {device})" + ) + + dtype = x.get_dtype() + try: + x_static = tuple(_static_dim_or_nyi(s, what="x shape") for s in x_sizes) + j_static = _static_dim_or_nyi(j, what="weight out_features") + k_static = _static_dim_or_nyi(k_w, what="weight in_features") + except NotImplementedError: + raise + except Exception as exc: + raise NotImplementedError( + "xtc::matmul: XTC matmul requires fully static shapes at compile time" + ) from exc + + w_static = (j_static, k_static) + b_static: tuple[int, ...] | None = None + if b is not None: + b_sizes = list(b.get_size()) + try: + b_static = tuple( + _static_dim_or_nyi(s, what="bias shape") for s in b_sizes + ) + except NotImplementedError: + raise + b_device = b.get_device() + if b_device is not None and b_device.type != "cpu": + raise NotImplementedError( + f"xtc::matmul: XTC matmul bias must be on CPU (got {b_device})" + ) + if b.get_dtype() != dtype: + raise NotImplementedError( + "xtc::matmul: XTC matmul requires bias dtype to match activations" + ) + + validate_matmul_xtc_support(x_static, w_static, dtype, device, b_shape=b_static) + + x_stride = _static_strides_or_nyi(x, x_static, what="x") + w_stride = _static_strides_or_nyi(w, w_static, what="weight") + + output_size = x_prefix + [j] + out_static: tuple[int, ...] | None = None + if V.graph.cpp_wrapper: + try: + out_static = tuple( + _static_dim_or_nyi(s, what="output shape") for s in output_size + ) + except NotImplementedError: + raise + except Exception as exc: + raise NotImplementedError( + "xtc::matmul: XTC matmul requires fully static output shape " + "at compile time" + ) from exc + + compile_matmul_at_lowering( + x_static, + w_static, + dtype, + inductor_cpp=V.graph.cpp_wrapper, + x_stride=x_stride if V.graph.cpp_wrapper else None, + w_stride=w_stride if V.graph.cpp_wrapper else None, + out_shape=out_static, + has_bias=b is not None, + ) + + inputs: list[Any] = [x, w] + kwargs: dict[str, Any] = {} + if b is not None: + inputs.append(b) + else: + kwargs["b"] = None + + return _register_extern_lowering( + torch.ops.xtc.matmul.default, + inputs, + kwargs, + output_size=output_size, + device=device, + dtype=dtype, + cpp_kernel_name="xtc::matmul_cpp" if V.graph.cpp_wrapper else None, + ) + + +def register_inductor_lowerings() -> None: + _register_matmul_inductor_lowering() + + assert torch.ops.xtc.matmul.default in lowerings, ( + "matmul Inductor lowering was not registered" + ) diff --git a/src/xtc/integration/pytorch/ops.py b/src/xtc/integration/pytorch/ops.py new file mode 100644 index 000000000..d39600833 --- /dev/null +++ b/src/xtc/integration/pytorch/ops.py @@ -0,0 +1,77 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# +# pyright: reportFunctionMemberAccess=false +from __future__ import annotations + +from math import prod +from typing import Any + +import torch +from torch.library import custom_op + +from xtc.integration.pytorch.eager import matmul_cpu + +# --------------------------------------------------------------------------- +# Matmul custom op (xtc::matmul) +# --------------------------------------------------------------------------- + + +@custom_op("xtc::matmul", mutates_args=()) +def matmul( + x: torch.Tensor, + w: torch.Tensor, + b: torch.Tensor | None = None, +) -> torch.Tensor: + raise RuntimeError("xtc::matmul: no kernel registered for this device") + + +@matmul.register_fake +def matmul_meta( + x: torch.Tensor, w: torch.Tensor, b: torch.Tensor | None = None +) -> torch.Tensor: + torch._check(x.dim() >= 1) + torch._check(w.dim() == 2) + torch._check(x.shape[-1] == w.shape[1]) + if b is not None: + torch._check(b.dim() == 1) + torch._check(w.shape[0] == b.shape[0]) + torch._check(x.device == b.device) + torch._check(x.device == w.device) + return x.new_empty(*x.shape[:-1], w.shape[0]) + + +def _matmul_setup_context( + ctx: Any, inputs: tuple[Any, ...], output: torch.Tensor +) -> None: + x, w, b = inputs + ctx.save_for_backward(x, w) + ctx.has_bias = b is not None + + +def _matmul_backward( + ctx: Any, grad_output: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + x, w = ctx.saved_tensors + k = w.shape[1] + leading = x.shape[:-1] + batch = prod(leading) if leading else 1 + x2d = x.reshape(batch, k) + go2d = grad_output.reshape(batch, w.shape[0]) + grad_x = (go2d @ w).reshape(x.shape) + grad_w = go2d.mT @ x2d + grad_b = go2d.sum(dim=0) if ctx.has_bias else None + return grad_x, grad_w, grad_b + + +# --------------------------------------------------------------------------- +# Register ops +# --------------------------------------------------------------------------- + + +def register_ops() -> None: + # Register eager kernel + matmul.register_kernel("cpu")(matmul_cpu) + # Register autograd + matmul.register_autograd(_matmul_backward, setup_context=_matmul_setup_context) diff --git a/src/xtc/integration/pytorch/torch_xtc.py b/src/xtc/integration/pytorch/torch_xtc.py new file mode 100644 index 000000000..7d5daaac1 --- /dev/null +++ b/src/xtc/integration/pytorch/torch_xtc.py @@ -0,0 +1,51 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# +"""PyTorch integration: XTC-backed ``xtc::matmul`` and ``torch.compile`` hooks.""" + +from __future__ import annotations + +from enum import Enum + +from xtc.integration.pytorch.fx_rewrite import register_pre_grad_pass +from xtc.integration.pytorch.ops import register_ops + + +class XtcIntegration(Enum): + """How XTC-backed ops are invoked under ``torch.compile``.""" + + EAGER = "eager" + INDUCTOR_CPP = "inductor-cpp" + + +def register_torch_xtc_extensions( + *, + xtc_integration: XtcIntegration | str = XtcIntegration.EAGER, +) -> None: + """Register XTC custom ops, Inductor pass, and integration mode.""" + import torch._inductor.config as inductor_config + + # IMPORT CUSTOM OPS AND REWRITE PASS + + register_ops() + + register_pre_grad_pass() + + # REGISTER INDUCTOR CPP LOWERING IF NEEDED + + integration = ( + xtc_integration + if isinstance(xtc_integration, XtcIntegration) + else XtcIntegration(xtc_integration) + ) + + if integration == XtcIntegration.INDUCTOR_CPP: + inductor_config.cpp_wrapper = True + from xtc.integration.pytorch.inductor_cpp import register_inductor_cpp_hooks + from xtc.integration.pytorch.inductor_lowering import ( + register_inductor_lowerings, + ) + + register_inductor_cpp_hooks() + register_inductor_lowerings() diff --git a/src/xtc/itf/comp/module.py b/src/xtc/itf/comp/module.py index e91f07939..3fb10506f 100644 --- a/src/xtc/itf/comp/module.py +++ b/src/xtc/itf/comp/module.py @@ -3,6 +3,7 @@ # Copyright (c) 2024-2026 The XTC Project Authors # from abc import ABC, abstractmethod +from pathlib import Path from typing import Any from ..exec import Evaluator, Executor @@ -72,12 +73,13 @@ def file_name(self) -> str: ... @abstractmethod - def export(self) -> None: - """Exports the module to a format suitable for execution. + def export(self, out_dir: str | Path, **kwargs: Any) -> None: + """Export the module for use from external C/C++ code. - This method handles the final step of making the compiled code - available for execution, typically by writing it to a shared - object file or similar executable format. + Args: + out_dir: Destination directory for exported artifacts (header, + shared library, test harness, etc.). + kwargs: Target-specific options (e.g. export name, random seed). """ ... diff --git a/src/xtc/itf/data/tensor.py b/src/xtc/itf/data/tensor.py index 7e5832afc..9d9a61a15 100644 --- a/src/xtc/itf/data/tensor.py +++ b/src/xtc/itf/data/tensor.py @@ -54,6 +54,26 @@ def device(self) -> AcceleratorDevice | None: """ ... + @property + @abstractmethod + def const(self) -> bool: + """Returns if the tensor contains constant data. + + Returns: + If the tensor contains constant data. + """ + ... + + @property + @abstractmethod + def layout(self) -> list[int] | None: + """Returns the layout of the tensor. + + Returns: + The layout of the tensor + """ + ... + @property @abstractmethod def ndim(self) -> int: diff --git a/src/xtc/runtimes/types/ndarray.py b/src/xtc/runtimes/types/ndarray.py index 351f4ecd1..197c223cf 100644 --- a/src/xtc/runtimes/types/ndarray.py +++ b/src/xtc/runtimes/types/ndarray.py @@ -40,7 +40,10 @@ class NDArray: rev_np_dtype_map: dict[tuple[int, int], str] = {} def __init__( - self, array: Any, runtime: CommonRuntimeInterface | None = None + self, + array: Any, + runtime: CommonRuntimeInterface | None = None, + layout: list[int] | None = None, ) -> None: if not self.rev_np_dtype_map: self.rev_np_dtype_map.update( @@ -53,6 +56,7 @@ def __init__( if self.runtime is None: self.runtime = HostRuntime() self.location = NDArrayLocation.HOST + self.layout = layout if isinstance(array, NDArray): raise RuntimeError("TODO: copy from CNDArray not supported yet") elif isinstance(array, np.ndarray): @@ -63,6 +67,11 @@ def __init__( self._to_device() def _from_numpy(self, nparray: np.ndarray) -> None: + if self.is_transposed(): + transposed_shape = [nparray.shape[1], nparray.shape[0]] + transposed_array = np.empty(transposed_shape, dtype=nparray.dtype) + np.copyto(transposed_array, np.transpose(nparray)) + nparray = transposed_array assert nparray.flags["C_CONTIGUOUS"] self.handle = self._new(nparray.shape, str(nparray.dtype)) self._copy_from(self.handle, nparray.ctypes.data_as(ctypes.c_voidp)) @@ -72,6 +81,11 @@ def _to_numpy(self) -> np.ndarray: np_dtype = self.dtype_str nparray = np.empty(shape=shape, dtype=np_dtype) self._copy_to(self.handle, nparray.ctypes.data_as(ctypes.c_voidp)) + if self.is_transposed(): + transposed_shape = [nparray.shape[1], nparray.shape[0]] + transposed_array = np.empty(transposed_shape, dtype=nparray.dtype) + np.copyto(transposed_array, np.transpose(nparray)) + nparray = transposed_array return nparray def _copy_to_numpy(self, out: np.ndarray) -> np.ndarray: @@ -81,6 +95,12 @@ def _copy_to_numpy(self, out: np.ndarray) -> np.ndarray: self._copy_to(self.handle, out.ctypes.data_as(ctypes.c_voidp)) return out + def is_transposed(self) -> bool: + if self.layout is None: + return False + assert len(self.layout) == 2, "Only 2D reshape is implemented" + return self.layout == [1, 0] + def numpy(self, out: np.ndarray | None = None) -> np.ndarray: if self.is_on_device(): assert isinstance(self.runtime, AcceleratorDevice) @@ -92,6 +112,7 @@ def numpy(self, out: np.ndarray | None = None) -> np.ndarray: if out is None: return self._to_numpy() else: + assert not self.is_transposed() return self._copy_to_numpy(out) def _to_device(self) -> None: diff --git a/src/xtc/targets/accelerator/gpu/GPUModule.py b/src/xtc/targets/accelerator/gpu/GPUModule.py index 9d66611bb..29caeb705 100644 --- a/src/xtc/targets/accelerator/gpu/GPUModule.py +++ b/src/xtc/targets/accelerator/gpu/GPUModule.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: BSD-3-Clause # Copyright (c) 2024-2026 The XTC Project Authors # +from pathlib import Path from typing import Any from typing_extensions import override @@ -71,7 +72,7 @@ def file_name(self) -> str: return self._file_name @override - def export(self) -> None: + def export(self, out_dir: str | Path, **kwargs: Any) -> None: raise NotImplementedError("GPUModule.export is not implemented") @override diff --git a/src/xtc/targets/accelerator/mppa/MppaModule.py b/src/xtc/targets/accelerator/mppa/MppaModule.py index 746329e16..f45af4577 100644 --- a/src/xtc/targets/accelerator/mppa/MppaModule.py +++ b/src/xtc/targets/accelerator/mppa/MppaModule.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: BSD-3-Clause # Copyright (c) 2024-2026 The XTC Project Authors # +from pathlib import Path from typing import Any, cast from typing_extensions import override @@ -76,7 +77,7 @@ def file_name(self) -> str: return self._file_name @override - def export(self) -> None: + def export(self, out_dir: str | Path, **kwargs: Any) -> None: raise NotImplementedError("AcceleratorModule.export is not implemented") @override diff --git a/src/xtc/targets/host/HostCppExport.py b/src/xtc/targets/host/HostCppExport.py new file mode 100644 index 000000000..359a2fb9e --- /dev/null +++ b/src/xtc/targets/host/HostCppExport.py @@ -0,0 +1,371 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# +from __future__ import annotations + +import re +import shutil +from dataclasses import dataclass +from functools import reduce +import operator +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable + +import numpy as np + +from xtc.utils.ext_tools import get_shlib_extension +from xtc.utils.numpy import np_init + +if TYPE_CHECKING: + from .HostModule import HostModule + +__all__ = ["HostCppExporter"] + +_DTYPE_MAP: dict[str, tuple[str, str]] = { + "float32": ("float", "f"), + "float64": ("double", "d"), +} + + +@dataclass(frozen=True) +class _TensorArg: + name: str + c_name: str + shape: tuple[int, ...] + dtype: str + c_type: str + numel: int + is_input: bool + + +def _c_ident(name: str) -> str: + ident = re.sub(r"[^a-zA-Z0-9_]", "_", name) + if ident and ident[0].isdigit(): + ident = f"_{ident}" + return ident or "tensor" + + +def _macro_prefix(c_name: str) -> str: + return f"XTC_{c_name.upper()}" + + +class HostCppExporter: + """Export a CPU shared-library module for linking from external C/C++ code.""" + + def __init__( + self, + module: HostModule, + out_dir: Path, + *, + name: str | None = None, + seed: int = 0, + ) -> None: + self._module = module + self._out_dir = out_dir + self._export_name = name or module.name + self._seed = seed + + def export(self) -> None: + inputs_spec_fn = self._module._np_inputs_spec + outputs_spec_fn = self._module._np_outputs_spec + reference_impl = self._module._reference_impl + if inputs_spec_fn is None or outputs_spec_fn is None: + raise ValueError("module is missing np_inputs_spec / np_outputs_spec") + if reference_impl is None: + raise ValueError("module is missing reference_impl (graph required)") + + inputs_spec = inputs_spec_fn() + outputs_spec = outputs_spec_fn() + input_names, output_names = self._resolve_arg_names( + len(inputs_spec), len(outputs_spec) + ) + + input_args = [ + self._tensor_arg(name, spec, is_input=True) + for name, spec in zip(input_names, inputs_spec) + ] + output_args = [ + self._tensor_arg(name, spec, is_input=False) + for name, spec in zip(output_names, outputs_spec) + ] + + self._out_dir.mkdir(parents=True, exist_ok=True) + (self._out_dir / "include").mkdir(exist_ok=True) + (self._out_dir / "lib").mkdir(exist_ok=True) + (self._out_dir / "data" / "inputs").mkdir(parents=True, exist_ok=True) + (self._out_dir / "data" / "outputs").mkdir(parents=True, exist_ok=True) + + self._copy_shared_lib() + self._write_golden_data(input_args, output_args, reference_impl) + self._write_header(input_args, output_args) + self._write_test_cpp(input_args, output_args) + self._write_makefile() + self._write_readme() + + def _resolve_arg_names( + self, num_inputs: int, num_outputs: int + ) -> tuple[list[str], list[str]]: + graph = self._module._graph + if graph is not None: + return ( + [node.name for node in graph.inputs_nodes], + [node.name for node in graph.outputs_nodes], + ) + return ( + [f"in{i}" for i in range(num_inputs)], + [f"out{i}" for i in range(num_outputs)], + ) + + def _tensor_arg( + self, name: str, spec: dict[str, Any], *, is_input: bool + ) -> _TensorArg: + dtype = spec["dtype"] + if dtype not in _DTYPE_MAP: + raise NotImplementedError( + f"export does not support dtype {dtype!r} (supported: {list(_DTYPE_MAP)})" + ) + shape = tuple(spec["shape"]) + c_type = _DTYPE_MAP[dtype][0] + c_name = _c_ident(name) + numel = reduce(operator.mul, shape, 1) + return _TensorArg( + name=name, + c_name=c_name, + shape=shape, + dtype=dtype, + c_type=c_type, + numel=numel, + is_input=is_input, + ) + + def _copy_shared_lib(self) -> None: + ext = get_shlib_extension() + dest = self._out_dir / "lib" / f"lib{self._export_name}.{ext}" + shutil.copy2(self._module.file_name, dest) + + def _write_golden_data( + self, + input_args: list[_TensorArg], + output_args: list[_TensorArg], + reference_impl: Callable[..., None], + ) -> None: + np.random.seed(self._seed) + inputs: list[np.ndarray[Any, Any]] = [] + for arg in input_args: + arr = np_init(shape=arg.shape, dtype=arg.dtype) + path = self._out_dir / "data" / "inputs" / f"{arg.c_name}.bin" + arr.tofile(path) + inputs.append(arr) + + outputs: list[np.ndarray[Any, Any]] = [] + for arg in output_args: + arr = np.zeros(arg.shape, dtype=arg.dtype) + outputs.append(arr) + + reference_impl(*inputs, *outputs) + + for arg, arr in zip(output_args, outputs): + path = self._out_dir / "data" / "outputs" / f"{arg.c_name}.bin" + arr.tofile(path) + + def _write_header( + self, input_args: list[_TensorArg], output_args: list[_TensorArg] + ) -> None: + payload = self._module.payload_name + all_args = input_args + output_args + params = ", ".join(f"{a.c_type}* {a.c_name}" for a in all_args) + meta_lines: list[str] = [] + for arg in all_args: + prefix = _macro_prefix(arg.c_name) + meta_lines.append(f"#define {prefix}_NDIM {len(arg.shape)}") + for i, dim in enumerate(arg.shape): + meta_lines.append(f"#define {prefix}_SHAPE_{i} {dim}") + meta_lines.append(f'#define {prefix}_DTYPE "{arg.dtype}"') + meta_lines.append(f"#define {prefix}_NUMEL {arg.numel}") + + meta = "\n".join(meta_lines) + content = f"""\ +#pragma once + +#include + +#ifdef __cplusplus +extern "C" {{ +#endif + +void {payload}({params}); + +/* Tensor metadata (shapes, dtypes, element counts). */ +{meta} + +#ifdef __cplusplus +}} +#endif +""" + path = self._out_dir / "include" / f"{self._export_name}.h" + path.write_text(content, encoding="utf-8") + + def _write_test_cpp( + self, input_args: list[_TensorArg], output_args: list[_TensorArg] + ) -> None: + payload = self._module.payload_name + header = self._export_name + all_args = input_args + output_args + c_types = {a.c_type for a in all_args} + if len(c_types) != 1: + raise NotImplementedError( + "export test.cpp requires a single floating-point dtype across all tensors" + ) + c_type = c_types.pop() + + load_blocks: list[str] = [] + ptr_names: list[str] = [] + check_blocks: list[str] = [] + + for arg in input_args: + load_blocks.append(self._load_vector_block(arg, c_type, "inputs")) + ptr_names.append(f"{arg.c_name}.data()") + + for arg in output_args: + load_blocks.append( + self._load_vector_block(arg, c_type, "outputs", zero_init=True) + ) + ptr_names.append(f"{arg.c_name}.data()") + + for arg in output_args: + check_blocks.append( + f"""\ + {{ + std::vector<{c_type}> ref = load_binary<{c_type}>( + "data/outputs/{arg.c_name}.bin", {arg.numel}); + if (!allclose({arg.c_name}, ref)) {{ + std::cerr << "FAIL {arg.c_name}\\n"; + return 1; + }} + std::cout << "OK {arg.c_name}\\n"; + }}""" + ) + + load_src = "\n\n".join(load_blocks) + checks = "\n".join(check_blocks) + call_args = ", ".join(ptr_names) + + content = f"""\ +#include "{header}.h" + +#include +#include +#include +#include +#include + +template +std::vector load_binary(const std::string& path, size_t expected_numel) {{ + std::ifstream in(path, std::ios::binary); + if (!in) {{ + throw std::runtime_error("cannot open " + path); + }} + in.seekg(0, std::ios::end); + const auto nbytes = static_cast(in.tellg()); + in.seekg(0, std::ios::beg); + const size_t numel = nbytes / sizeof(T); + if (numel != expected_numel) {{ + throw std::runtime_error("size mismatch for " + path); + }} + std::vector data(numel); + in.read(reinterpret_cast(data.data()), + static_cast(nbytes)); + return data; +}} + +template +bool allclose(const std::vector& got, const std::vector& ref, + T rtol = static_cast(1e-5), T atol = static_cast(1e-6)) {{ + if (got.size() != ref.size()) {{ + return false; + }} + for (size_t i = 0; i < got.size(); ++i) {{ + const T diff = std::abs(got[i] - ref[i]); + const T tol = atol + rtol * std::abs(ref[i]); + if (diff > tol) {{ + std::cerr << " mismatch at index " << i << ": got=" << got[i] + << " ref=" << ref[i] << "\\n"; + return false; + }} + }} + return true; +}} + +int main() {{ + try {{ +{load_src} + + {payload}({call_args}); + +{checks} + std::cout << "All checks passed.\\n"; + return 0; + }} catch (const std::exception& ex) {{ + std::cerr << "Error: " << ex.what() << "\\n"; + return 1; + }} +}} +""" + (self._out_dir / "test.cpp").write_text(content, encoding="utf-8") + + def _load_vector_block( + self, arg: _TensorArg, c_type: str, subdir: str, *, zero_init: bool = False + ) -> str: + if zero_init: + init = f"std::vector<{c_type}>({arg.numel})" + else: + init = ( + f'load_binary<{c_type}>("data/{subdir}/{arg.c_name}.bin", {arg.numel})' + ) + return f" std::vector<{c_type}> {arg.c_name} = {init};" + + def _write_makefile(self) -> None: + ext = get_shlib_extension() + if ext == "dylib": + rpath = "-Wl,-rpath,@loader_path/lib" + else: + rpath = "-Wl,-rpath,'$$ORIGIN/lib'" + content = f"""\ +NAME := {self._export_name} +CXX ?= c++ +CXXFLAGS ?= -O2 -std=c++17 -Wall -Wextra +INCLUDES := -Iinclude +LDFLAGS := -Llib -l$(NAME) {rpath} + +test: test.cpp +\t$(CXX) $(CXXFLAGS) $(INCLUDES) $< -o $@ $(LDFLAGS) + +.PHONY: run clean +run: test +\t./test + +clean: +\trm -f test +""" + (self._out_dir / "Makefile").write_text(content, encoding="utf-8") + + def _write_readme(self) -> None: + content = f"""\ +# XTC exported kernel: {self._export_name} + +Artifacts produced by `module.export()` for CPU linking from C/C++. + +## Build and test + +```bash +make +make run +``` + +- `include/{self._export_name}.h` — kernel declaration and tensor metadata macros +- `lib/lib{self._export_name}.{get_shlib_extension()}` — compiled shared library +- `data/inputs/*.bin` / `data/outputs/*.bin` — golden inputs and reference outputs +- `test.cpp` — loads golden data, calls the kernel, compares against reference +""" + (self._out_dir / "README.md").write_text(content, encoding="utf-8") diff --git a/src/xtc/targets/host/HostModule.py b/src/xtc/targets/host/HostModule.py index 681c934ee..933e9aa7b 100644 --- a/src/xtc/targets/host/HostModule.py +++ b/src/xtc/targets/host/HostModule.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: BSD-3-Clause # Copyright (c) 2024-2026 The XTC Project Authors # +from pathlib import Path from typing_extensions import override from typing import Any @@ -92,8 +93,16 @@ def file_name(self) -> str: return self._file_name @override - def export(self) -> None: - pass + def export(self, out_dir: str | Path, **kwargs: Any) -> None: + if self._file_type != "shlib": + raise NotImplementedError( + f"export only supports shlib modules (got {self._file_type})" + ) + from .HostCppExport import HostCppExporter + + name = kwargs.get("name") + seed = kwargs.get("seed", 0) + HostCppExporter(self, Path(out_dir), name=name, seed=seed).export() @override def get_evaluator(self, **kwargs: Any) -> itf.exec.Evaluator: diff --git a/src/xtc/utils/evaluation.py b/src/xtc/utils/evaluation.py index 1e372015e..d5d37c5c7 100644 --- a/src/xtc/utils/evaluation.py +++ b/src/xtc/utils/evaluation.py @@ -36,6 +36,7 @@ def _graph_np_inputs_spec() -> list[dict[str, Any]]: "shape": type.constant_shape, "dtype": type.constant_dtype, "device": type.device, + "layout": type.layout, } for type in inputs_types ] @@ -88,13 +89,18 @@ def ensure_ndarray_parameters( out_init = np.zeros if init_zero else np.empty inputs = [ ( - np_init(**{k: v for k, v in spec.items() if k != "device"}), + np_init( + **{k: v for k, v in spec.items() if k != "device" and k != "layout"} + ), spec["device"] if "device" in spec else HostRuntime.get(), + spec["layout"] if "layout" in spec else None, ) for spec in inputs_spec ] outputs = [ - out_init(**{k: v for k, v in spec.items() if k != "device"}) + out_init( + **{k: v for k, v in spec.items() if k != "device" and k != "layout"} + ) for spec in outputs_spec ] parameters = ( diff --git a/tests/filecheck/backends/invalid/test_matmul_tvm_layout.py b/tests/filecheck/backends/invalid/test_matmul_tvm_layout.py new file mode 100644 index 000000000..a2ebc6315 --- /dev/null +++ b/tests/filecheck/backends/invalid/test_matmul_tvm_layout.py @@ -0,0 +1,32 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: module_tvm + +import xtc.graphs.xtc.op as O +from xtc.backends.tvm import Backend + +I, J, K, dtype = 4, 32, 512, "float32" +a = O.tensor((I, K), dtype, name="A", layout=[1, 0]) +b = O.tensor((K, J), dtype, name="B") + +with O.graph(name="matmul_layout") as gb: + O.matmul(a, b, name="C") + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="matmul_tvm_layout", + print_source_ir=True, + print_transformed_ir=False, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") +# XFAIL: * diff --git a/tests/filecheck/backends/invalid/test_relu_mlir_layout.py b/tests/filecheck/backends/invalid/test_relu_mlir_layout.py new file mode 100644 index 000000000..59d587384 --- /dev/null +++ b/tests/filecheck/backends/invalid/test_relu_mlir_layout.py @@ -0,0 +1,17 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: module_mlir + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend + +I, J, dtype = 4, 32, "float32" +a = O.tensor((I, J), dtype, name="A", layout=[1, 0]) + +with O.graph(name="relu_layout") as gb: + O.relu(a, name="relu") + +graph = gb.graph +Backend(graph) + +# CHECK: NotImplementedError: tensor layout is not yet implemented in MLIR backend +# XFAIL: * diff --git a/tests/filecheck/backends/test_matmul_mlir_const.py b/tests/filecheck/backends/test_matmul_mlir_const.py new file mode 100644 index 000000000..cf6377b1f --- /dev/null +++ b/tests/filecheck/backends/test_matmul_mlir_const.py @@ -0,0 +1,70 @@ +# RUN: python %s 2>&1 | filecheck %s + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend + +I, J, K, dtype = 4, 32, 512, "float32" +a = O.tensor((I, K), dtype, name="A", const=True) +b = O.tensor((K, J), dtype, name="B") + +with O.graph(name="matmul_const") as gb: + O.matmul(a, b, name="C") + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="matmul_mlir_const", + print_source_ir=True, + print_transformed_ir=False, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul_const(%arg0: memref<4x512xf32> {llvm.noalias, memref.const}, %arg1: memref<512x32xf32> {llvm.noalias}, %arg2: memref<4x32xf32> {llvm.noalias}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<4x32xf32>) +# CHECK-NEXT: linalg.matmul {__xtc_id_C_} ins(%arg0, %arg1 : memref<4x512xf32>, memref<512x32xf32>) outs(%arg2 : memref<4x32xf32>) +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %1 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "./j" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "./k" : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: matmul_const +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 4x512xfloat32, const +# CHECK-NEXT: - %1 : 512x32xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 4x32xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x512xfloat32, const, 512x32xfloat32] -> [4x32xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/backends/test_matmul_mlir_layout.py b/tests/filecheck/backends/test_matmul_mlir_layout.py new file mode 100644 index 000000000..5fb211dc3 --- /dev/null +++ b/tests/filecheck/backends/test_matmul_mlir_layout.py @@ -0,0 +1,78 @@ +# RUN: python %s 2>&1 | filecheck %s + +import xtc.graphs.xtc.op as O +from xtc.backends.mlir import Backend + +I, J, K, dtype = 4, 32, 512, "float32" +a = O.tensor((I, K), dtype, name="A", layout=[1, 0]) +b = O.tensor((K, J), dtype, name="B") + +with O.graph(name="matmul_layout") as gb: + O.matmul(a, b, name="C") + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="matmul_mlir_layout", + print_source_ir=True, + print_transformed_ir=False, +) +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") +# CHECK: // -----// IR Dump Before transform //----- // +# CHECK-NEXT: #map = affine_map<(d0, d1, d2) -> (d2, d0)> +# CHECK-NEXT: #map1 = affine_map<(d0, d1, d2) -> (d2, d1)> +# CHECK-NEXT: #map2 = affine_map<(d0, d1, d2) -> (d0, d1)> +# CHECK-NEXT: module attributes {transform.with_named_sequence} { +# CHECK-NEXT: func.func @matmul_layout(%arg0: memref<512x4xf32> {llvm.noalias}, %arg1: memref<512x32xf32> {llvm.noalias}, %arg2: memref<4x32xf32> {llvm.noalias}) { +# CHECK-NEXT: %cst = arith.constant 0.000000e+00 : f32 +# CHECK-NEXT: linalg.fill {__xtc_id_C_0_} ins(%cst : f32) outs(%arg2 : memref<4x32xf32>) +# CHECK-NEXT: linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : memref<512x4xf32>, memref<512x32xf32>) outs(%arg2 : memref<4x32xf32>) attrs = {__xtc_id_C_} { +# CHECK-NEXT: ^bb0(%in: f32, %in_0: f32, %out: f32): +# CHECK-NEXT: %0 = arith.mulf %in, %in_0 : f32 +# CHECK-NEXT: %1 = arith.addf %out, %0 : f32 +# CHECK-NEXT: linalg.yield %1 : f32 +# CHECK-NEXT: } +# CHECK-NEXT: return +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @_vecto(%arg0: !transform.any_op {transform.consumed}) { +# CHECK-NEXT: transform.structured.vectorize %arg0 : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { +# CHECK-NEXT: %0 = transform.structured.match attributes {__xtc_id_C_0_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op, %loops = transform.structured.tile_using_for %0 tile_sizes [1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_0, %loops_1 = transform.structured.tile_using_for %tiled_linalg_op tile_sizes [0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_1 "./j" : !transform.any_op +# CHECK-NEXT: %1 = transform.structured.match attributes {__xtc_id_C_} in %arg0 : (!transform.any_op) -> !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_2, %loops_3 = transform.structured.tile_using_for %1 tile_sizes [1, 0, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_3 "./i" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_4, %loops_5 = transform.structured.tile_using_for %tiled_linalg_op_2 tile_sizes [0, 1, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_5 "./j" : !transform.any_op +# CHECK-NEXT: %tiled_linalg_op_6, %loops_7 = transform.structured.tile_using_for %tiled_linalg_op_4 tile_sizes [0, 0, 1] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) +# CHECK-NEXT: transform.annotate %loops_7 "./k" : !transform.any_op +# CHECK-NEXT: transform.yield +# CHECK-NEXT: } +# CHECK-NEXT: } +# CHECK-NEXT: +# CHECK-NEXT: graph: +# CHECK-NEXT: name: matmul_layout +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 4x512xfloat32, <(0,1)->(1,0)> +# CHECK-NEXT: - %1 : 512x32xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 4x32xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x512xfloat32, <(0,1)->(1,0)>, 512x32xfloat32] -> [4x32xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/pytorch/test_my_linear_compile.py b/tests/filecheck/pytorch/test_my_linear_compile.py new file mode 100644 index 000000000..2f7122099 --- /dev/null +++ b/tests/filecheck/pytorch/test_my_linear_compile.py @@ -0,0 +1,48 @@ +# RUN: python %s 2>&1 | filecheck %s +# UNSUPPORTED: mlir-target=nvgpu + +from __future__ import annotations + +import sys + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from xtc.integration.pytorch import ( + XtcIntegration, + clear_kernel_cache, + register_torch_xtc_extensions, +) + + +def main() -> None: + clear_kernel_cache() + register_torch_xtc_extensions(xtc_integration=XtcIntegration.EAGER) + + class M(nn.Module): + def __init__(self) -> None: + super().__init__() + self.w = nn.Parameter(torch.randn(16, 8)) + self.b = nn.Parameter(torch.randn(16)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.linear(x, self.w, self.b) + + m = M() + x = torch.randn(4, 8) + y_ref = m(x) + compiled = torch.compile(m) + y = compiled(x) + max_err = (y - y_ref).abs().max().item() + print(f"MAX_ERR: {max_err}") + if not torch.allclose(y, y_ref): + sys.exit(1) + print("OK: torch.compile + XTC linear matches eager reference") + + +if __name__ == "__main__": + main() + +# CHECK: MAX_ERR: 0.0 +# CHECK-NEXT: OK: torch.compile + XTC linear matches eager reference diff --git a/tests/filecheck/pytorch/test_my_linear_kernel.py b/tests/filecheck/pytorch/test_my_linear_kernel.py new file mode 100644 index 000000000..0e768cc79 --- /dev/null +++ b/tests/filecheck/pytorch/test_my_linear_kernel.py @@ -0,0 +1,45 @@ +# RUN: python %s 2>&1 | filecheck %s +# UNSUPPORTED: mlir-target=nvgpu + +from __future__ import annotations + +import sys +import tempfile +from pathlib import Path + +import torch +import torch.nn.functional as F + +from xtc.integration.pytorch import clear_kernel_cache, compile_matmul_kernel + + +def main() -> None: + clear_kernel_cache() + with tempfile.TemporaryDirectory(prefix="xtc_pytorch_") as tmp: + cache_dir = Path(tmp) + kernel = compile_matmul_kernel( + (4, 8), + (16, 8), + torch.float32, + cache_dir=cache_dir, + ) + print("KERNEL: compiled") + + x = torch.randn(4, 8) + w = torch.randn(16, 8) + b = torch.randn(16) + y = kernel(x, w, b) + y_ref = F.linear(x, w, b) + max_err = (y - y_ref).abs().max().item() + print(f"MAX_ERR: {max_err}") + if not torch.allclose(y, y_ref): + sys.exit(1) + print("OK: XTC linear kernel matches F.linear") + + +if __name__ == "__main__": + main() + +# CHECK: KERNEL: compiled +# CHECK-NEXT: MAX_ERR: 0 +# CHECK-NEXT: OK: XTC linear kernel matches F.linear diff --git a/tests/pytest/mlir/test_cpp_export.py b/tests/pytest/mlir/test_cpp_export.py new file mode 100644 index 000000000..232497ffd --- /dev/null +++ b/tests/pytest/mlir/test_cpp_export.py @@ -0,0 +1,60 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +import pytest + +from mlir_utils import matmul_impl, requires_mlir +from xtc.utils.ext_tools import get_shlib_extension + +I, J, K, DTYPE = 4, 16, 8, "float32" + + +@requires_mlir +@pytest.mark.skipif(sys.platform != "linux", reason="cpp export integration test (linux)") +def test_cpp_export_matmul(tmp_path: Path) -> None: + impl = matmul_impl(I, J, K, DTYPE, "matmul_export") + sch = impl.get_scheduler() + sch.set_dims(["i", "j", "k"]) + sched = sch.schedule() + + comp = impl.get_compiler(shared_lib=True, dump_file=str(tmp_path / "matmul_export")) + module = comp.compile(sched) + + export_dir = tmp_path / "export" + module.export(export_dir) + + ext = get_shlib_extension() + export_name = "matmul_export" + assert (export_dir / "include" / f"{export_name}.h").is_file() + assert (export_dir / "lib" / f"lib{export_name}.{ext}").is_file() + assert (export_dir / "test.cpp").is_file() + assert (export_dir / "Makefile").is_file() + assert (export_dir / "README.md").is_file() + assert list((export_dir / "data" / "inputs").glob("*.bin")) + assert list((export_dir / "data" / "outputs").glob("*.bin")) + + build = subprocess.run( + ["make", "test"], + cwd=export_dir, + capture_output=True, + text=True, + check=False, + ) + assert build.returncode == 0, f"make failed:\n{build.stdout}\n{build.stderr}" + + run = subprocess.run( + ["./test"], + cwd=export_dir, + capture_output=True, + text=True, + check=False, + ) + assert run.returncode == 0, f"test binary failed:\n{run.stdout}\n{run.stderr}" + assert "All checks passed." in run.stdout diff --git a/tests/pytest/pytorch/test_inductor_xtc_linear.py b/tests/pytest/pytorch/test_inductor_xtc_linear.py new file mode 100644 index 000000000..f3e705b54 --- /dev/null +++ b/tests/pytest/pytorch/test_inductor_xtc_linear.py @@ -0,0 +1,48 @@ +# +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2024-2026 The XTC Project Authors +# +from __future__ import annotations + +import sys + +import pytest + +pytest.importorskip("torch") + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch._inductor.utils import fresh_cache + +from xtc.integration.pytorch import ( + XtcIntegration, + clear_kernel_cache, + register_torch_xtc_extensions, +) + + +@pytest.mark.skipif(sys.platform != "linux", reason="inductor-cpp XTC integration (linux)") +def test_inductor_cpp_matmul_matches_eager() -> None: + try: + import mlir # noqa: F401 + except ImportError as exc: + pytest.skip(f"MLIR backend required: {exc}") + + clear_kernel_cache() + register_torch_xtc_extensions(xtc_integration=XtcIntegration.INDUCTOR_CPP) + + class M(nn.Module): + def __init__(self) -> None: + super().__init__() + self.w = nn.Parameter(torch.randn(16, 8)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.linear(x, self.w) + + m = M() + x = torch.randn(4, 8) + y_ref = m(x) + with fresh_cache(): + y = torch.compile(m, backend="inductor")(x) + assert torch.allclose(y, y_ref), f"max err={(y - y_ref).abs().max().item()}"