diff --git a/src/xtc/backends/mlir/MlirBindingsExtensions.py b/src/xtc/backends/mlir/MlirBindingsExtensions.py index cfa0d6c2..a9b3cd2a 100644 --- a/src/xtc/backends/mlir/MlirBindingsExtensions.py +++ b/src/xtc/backends/mlir/MlirBindingsExtensions.py @@ -15,7 +15,10 @@ # Optional XTC extensions to the MLIR Python bindings: the module to import # mapped to the pass-pipeline entries it contributes. _EXTENSIONS: dict[str, tuple[str, ...]] = { - "mlir.xtc_transform": ("func.func(reduce-extract-slices)",), + "mlir.xtc_transform": ( + "func.func(reduce-extract-slices)", + "scf-add-vector-acc", + ), } # Reverse map from a contributed pass to its providing module, so a pass can be diff --git a/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py b/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py index 7a83f424..46c9887f 100644 --- a/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py +++ b/src/xtc/backends/mlir/MlirTarget/MlirNVGPUTarget.py @@ -34,6 +34,7 @@ from .MlirTarget import MlirTarget from ..MlirConfig import MlirConfig from ..MlirProgram import RawMlirProgram +import xtc.backends.mlir.MlirBindingsExtensions as binding_extensions __all__ = ["MlirNVGPUTarget"] @@ -437,7 +438,7 @@ def __init__( self._mlir_program = mlir_program def _lowering_pipeline(self, sm_arch: str, ptx_version: str) -> list[str]: - return [ + return binding_extensions.passes(["scf-add-vector-acc"]) + [ "canonicalize", "cse", "sccp",