diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py --- a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py +++ b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py @@ -9,6 +9,7 @@ import functools import numpy as np import os +import sys # Import MLIR related modules. from mlir import all_passes_registration # Register MLIR compiler passes. @@ -16,7 +17,14 @@ from mlir import ir from mlir import runtime from mlir.dialects import sparse_tensor -from mlir.passmanager import PassManager + +_TOOL_PATH = os.path.join( + os.path.join( + os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), + "python"), "tools") +sys.path.append(_TOOL_PATH) +import sparse_compiler # Type aliases for type annotation. _SupportFunc = Callable[..., None] @@ -184,8 +192,7 @@ A JIT execution engine for the MLIR module. """ - pipeline = f"sparse-compiler" - PassManager.parse(pipeline).run(module) + sparse_compiler.SparseCompiler(options="")(module) return execution_engine.ExecutionEngine( module, opt_level=_OPT_LEVEL, shared_libs=[_get_support_lib_name()])