From c8991db15a8258d360f69b038e338b2346aa8e2b Mon Sep 17 00:00:00 2001 From: Adam Siemieniuk Date: Mon, 10 Nov 2025 09:12:01 +0100 Subject: [PATCH 1/2] [examples][mlir] Fix shared library path handling Improves handiling around the list of shared libraries paths to avoid execution engine errors caused by path resolution. Example's documentation wording is also improved as the example kernel does not actively require any runtime utilities. However, the section is kept to showcase an example library setup as external runtime utilities are often needed for more complex kernels. --- python/examples/mlir/compile_and_run.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/python/examples/mlir/compile_and_run.py b/python/examples/mlir/compile_and_run.py index 345251a..68ff7d8 100644 --- a/python/examples/mlir/compile_and_run.py +++ b/python/examples/mlir/compile_and_run.py @@ -149,9 +149,9 @@ def main(): pm.run(kernel.operation) ### Compilation ### - # External shared libraries, containing MLIR runner utilities, are generally - # required to execute the compiled module. - # In this case, MLIR runner utils libraries are expected: + # External shared libraries, runtime utilities, might be needed to execute + # the compiled module. + # For example, MLIR runner utils libraries such as: # - libmlir_runner_utils.so # - libmlir_c_runner_utils.so # @@ -159,7 +159,10 @@ def main(): # The execution engine requires full paths to the libraries. # For example, the env variable can be set as: # LIGHTHOUSE_SHARED_LIBS=$PATH_TO_LLVM/build/lib/lib1.so:$PATH_TO_LLVM/build/lib/lib2.so - mlir_libs = os.environ.get("LIGHTHOUSE_SHARED_LIBS", default="").split(":") + mlir_libs = [] + lh_shared_libs = os.environ.get("LIGHTHOUSE_SHARED_LIBS") + if lh_shared_libs: + mlir_libs += lh_shared_libs.split(":") # JIT the kernel. eng = ExecutionEngine(kernel, opt_level=2, shared_libs=mlir_libs) From 55e3ef0a396e583eff523aac7ad2ece03030e4fe Mon Sep 17 00:00:00 2001 From: Adam Siemieniuk Date: Mon, 10 Nov 2025 13:15:49 +0100 Subject: [PATCH 2/2] Switch to input args + formatting --- python/examples/mlir/compile_and_run.py | 46 ++++++++++++++++--------- 1 file changed, 29 insertions(+), 17 deletions(-) diff --git a/python/examples/mlir/compile_and_run.py b/python/examples/mlir/compile_and_run.py index 68ff7d8..d0529f6 100644 --- a/python/examples/mlir/compile_and_run.py +++ b/python/examples/mlir/compile_and_run.py @@ -1,5 +1,5 @@ import torch -import os +import argparse from mlir import ir from mlir.dialects import transform @@ -65,7 +65,8 @@ def create_schedule(ctx: ir.Context) -> ir.Module: func = structured.MatchOp.match_op_names( named_seq.bodyTarget, ["func.func"] ) - # Use C interface wrappers - required to make function executable after jitting. + # Use C interface wrappers - required to make function executable + # after jitting. func = transform.apply_registered_pass( anytype, func, "llvm-request-c-wrappers" ) @@ -126,7 +127,7 @@ def create_pass_pipeline(ctx: ir.Context) -> PassManager: # The example's entry point. -def main(): +def main(args): ### Baseline computation ### # Create inputs. a = torch.randn(16, 32, dtype=torch.float32) @@ -149,29 +150,23 @@ def main(): pm.run(kernel.operation) ### Compilation ### + # Parse additional libraries if present. + # # External shared libraries, runtime utilities, might be needed to execute # the compiled module. - # For example, MLIR runner utils libraries such as: - # - libmlir_runner_utils.so - # - libmlir_c_runner_utils.so - # - # Get paths to MLIR runner shared libraries through an environment variable. # The execution engine requires full paths to the libraries. - # For example, the env variable can be set as: - # LIGHTHOUSE_SHARED_LIBS=$PATH_TO_LLVM/build/lib/lib1.so:$PATH_TO_LLVM/build/lib/lib2.so mlir_libs = [] - lh_shared_libs = os.environ.get("LIGHTHOUSE_SHARED_LIBS") - if lh_shared_libs: - mlir_libs += lh_shared_libs.split(":") + if args.shared_libs: + mlir_libs += args.shared_libs.split(",") # JIT the kernel. eng = ExecutionEngine(kernel, opt_level=2, shared_libs=mlir_libs) # Initialize the JIT engine. # - # The deferred initialization executes global constructors that might have been - # created by the module during engine creation (for example, when `gpu.module` - # is present) or registered afterwards. + # The deferred initialization executes global constructors that might + # have been created by the module during engine creation (for example, + # when `gpu.module` is present) or registered afterwards. # # Initialization is not strictly necessary in this case. # However, it is a good practice to perform it regardless. @@ -197,4 +192,21 @@ def main(): if __name__ == "__main__": - main() + parser = argparse.ArgumentParser() + + # External shared libraries, runtime utilities, might be needed to + # execute the compiled module. + # For example, MLIR runner utils libraries such as: + # - libmlir_runner_utils.so + # - libmlir_c_runner_utils.so + # + # Full paths to the libraries should be provided. + # For example: + # --shared-libs=$LLVM_BUILD/lib/lib1.so,$LLVM_BUILD/lib/lib2.so + parser.add_argument( + "--shared-libs", + type=str, + help="Comma-separated list of libraries to link dynamically", + ) + args = parser.parse_args() + main(args)