From 87617bf510bb9828dcc5d82b7c9e01aef76c325e Mon Sep 17 00:00:00 2001 From: Sylvain Noiry Date: Fri, 22 May 2026 14:38:28 +0200 Subject: [PATCH 1/3] Fix lowering of ub dialect Lower ub dialect to llvm when using the LLVMIR target. --- src/xtc/backends/mlir/MlirTarget/MlirLLVMTarget.py | 1 + 1 file changed, 1 insertion(+) 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", From caa84d4712f8d7731c188f4d5616f9d1792ff9a1 Mon Sep 17 00:00:00 2001 From: Sylvain Noiry Date: Fri, 27 Mar 2026 11:28:49 +0100 Subject: [PATCH 2/3] Add options to relayout tensor or mark as const This patch adds a "layout" option when creating XTC tensors to relayout the storage in memory. The relayouting is considered done outside the evaluation or at compile time if the data is constat. Thus this patch also add a "const" option to specify if the tensor data is constant. --- src/xtc/backends/mlir/MlirGraphBackend.py | 10 ++- src/xtc/backends/mlir/MlirOps.py | 78 +++++++++++++++++-- src/xtc/backends/tvm/TVMOps.py | 8 ++ src/xtc/graphs/xtc/data.py | 46 ++++++++++- src/xtc/graphs/xtc/expr.py | 24 ++++-- src/xtc/itf/data/tensor.py | 20 +++++ src/xtc/runtimes/types/ndarray.py | 23 +++++- src/xtc/utils/evaluation.py | 10 ++- .../invalid/test_matmul_tvm_layout.py | 32 ++++++++ .../backends/invalid/test_relu_mlir_layout.py | 17 ++++ .../backends/test_matmul_mlir_const.py | 70 +++++++++++++++++ .../backends/test_matmul_mlir_layout.py | 78 +++++++++++++++++++ 12 files changed, 396 insertions(+), 20 deletions(-) create mode 100644 tests/filecheck/backends/invalid/test_matmul_tvm_layout.py create mode 100644 tests/filecheck/backends/invalid/test_relu_mlir_layout.py create mode 100644 tests/filecheck/backends/test_matmul_mlir_const.py create mode 100644 tests/filecheck/backends/test_matmul_mlir_layout.py 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/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/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/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 From 7294fd1648384f1aaa1db316b51d1fd1d44b6d04 Mon Sep 17 00:00:00 2001 From: Sylvain Noiry Date: Wed, 20 May 2026 15:04:15 +0200 Subject: [PATCH 3/3] Start working on Pytorch integration - Create custom ops for xtc operators in pytorch (only xtc.matmul for now) - Create a pre_grad fx pass to rewrite torch operators to xtc torch operators (e.g. linear to matmul) - Provide a eager implementation of the xtc torch operators that do one shot compile with xtc - Provide a torch inductor lowering for xtc torch operators (with cpp wrapper), and call one shot xtc compile at compile time of aot inductor run - patch runtime inductor ops lib with compiled kernel from xtc + codegen a cpp wrapper around This is still work in progress --- pyproject.toml | 4 +- src/xtc/integration/__init__.py | 4 + src/xtc/integration/pytorch/__init__.py | 20 + src/xtc/integration/pytorch/compile.py | 406 ++++++++++++++++++ src/xtc/integration/pytorch/eager.py | 39 ++ src/xtc/integration/pytorch/fx_rewrite.py | 38 ++ src/xtc/integration/pytorch/inductor_cpp.py | 357 +++++++++++++++ .../integration/pytorch/inductor_lowering.py | 206 +++++++++ src/xtc/integration/pytorch/ops.py | 77 ++++ src/xtc/integration/pytorch/torch_xtc.py | 51 +++ src/xtc/itf/comp/module.py | 12 +- src/xtc/targets/accelerator/gpu/GPUModule.py | 3 +- .../targets/accelerator/mppa/MppaModule.py | 3 +- src/xtc/targets/host/HostCppExport.py | 371 ++++++++++++++++ src/xtc/targets/host/HostModule.py | 13 +- .../pytorch/test_my_linear_compile.py | 48 +++ .../pytorch/test_my_linear_kernel.py | 45 ++ tests/pytest/mlir/test_cpp_export.py | 60 +++ .../pytorch/test_inductor_xtc_linear.py | 48 +++ 19 files changed, 1794 insertions(+), 11 deletions(-) create mode 100644 src/xtc/integration/__init__.py create mode 100644 src/xtc/integration/pytorch/__init__.py create mode 100644 src/xtc/integration/pytorch/compile.py create mode 100644 src/xtc/integration/pytorch/eager.py create mode 100644 src/xtc/integration/pytorch/fx_rewrite.py create mode 100644 src/xtc/integration/pytorch/inductor_cpp.py create mode 100644 src/xtc/integration/pytorch/inductor_lowering.py create mode 100644 src/xtc/integration/pytorch/ops.py create mode 100644 src/xtc/integration/pytorch/torch_xtc.py create mode 100644 src/xtc/targets/host/HostCppExport.py create mode 100644 tests/filecheck/pytorch/test_my_linear_compile.py create mode 100644 tests/filecheck/pytorch/test_my_linear_kernel.py create mode 100644 tests/pytest/mlir/test_cpp_export.py create mode 100644 tests/pytest/pytorch/test_inductor_xtc_linear.py 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/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/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/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()}"