Add external C functions support for JIT (#2475)

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: mawad-amd

.github/workflows/buildAndTestRyzenAI.yml

@@ -187,4 +187,4 @@ jobs:
           ninja check-reference-designs
           ninja check-programming-guide
 
-          popd
+          popd

python/iron/__init__.py

@@ -1,5 +1,5 @@
 from .globalbuffer import GlobalBuffer
-from .kernel import Kernel
+from .kernel import ExternalFunction, Kernel
 from .localbuffer import LocalBuffer
 from .program import Program
 from .worker import Worker, WorkerRuntimeBarrier

python/iron/compile/compile.py

@@ -15,6 +15,7 @@ def compile_cxx_core_function(
     source_path: str,
     target_arch: str,
     output_path: str,
+    include_dirs=None,
     compile_args=None,
     cwd=None,
     verbose=False,
@@ -28,8 +29,9 @@ def compile_cxx_core_function(
         source_path (str): Path to C++ source.
         target_arch (str): Target architecture, e.g., aie2.
         output_path (str): Output object file path.
-        compile_args (list[str]): Compile arguments to peano.
-        cwd (str): Overrides the current working directory.
+        include_dirs (list[str], optional): List of include directories to add with -I.
+        compile_args (list[str], optional): Additional compile arguments to peano.
+        cwd (str, optional): Overrides the current working directory.
         verbose (bool): Enable verbose output.
     """
     cmd = [
@@ -48,8 +50,16 @@ def compile_cxx_core_function(
         "-DNDEBUG",
         f"--target={target_arch}-none-unknown-elf",
     ]
+
+    # Add include directories
+    if include_dirs:
+        for include_dir in include_dirs:
+            cmd.extend(["-I", include_dir])
+
+    # Add additional compile arguments
     if compile_args:
         cmd = cmd + compile_args
+
     if verbose:
         print("Executing:", " ".join(cmd))
     ret = subprocess.run(

python/iron/jit.py

@@ -11,10 +11,16 @@
 import hashlib
 import numpy as np
 import pyxrt as xrt
+import shutil
+import sys
+import traceback
 
 from aie.extras.context import mlir_mod_ctx
 from ..utils.compile import compile_mlir_module_to_binary
 from ..utils.xrt import read_insts_binary
+from .device import NPU1, NPU2, NPU1Col1, NPU2Col1
+from .config import get_current_device
+from aie.dialects.aie import AIEDevice
 
 
 # The `iron.jit` decorator below caches compiled kenrels inside the `IRON_CACHE_DIR` directory.
@@ -93,11 +99,15 @@ def __call__(self, *args):
         kernel_args = []
 
         for tensor in args:
+            # Skip callable arguments since these are inlined in the kernel
+            if callable(tensor):
+                continue
             if not hasattr(tensor, "buffer_object"):
                 raise TypeError(
                     f"Expected Tensor with .buffer_object(), got {type(tensor)}"
                 )
             kernel_args.append(tensor.buffer_object())
+
         h = self.__kernel(opcode, self.__insts_buffer_bo, self.__n_insts, *kernel_args)
         r = h.wait()
         if r != xrt.ert_cmd_state.ERT_CMD_STATE_COMPLETED:
@@ -133,18 +143,61 @@ def jit(function=None, is_placed=True, use_cache=True):
 
     @functools.wraps(function)
     def decorator(*args, **kwargs):
-        if is_placed:
-            with mlir_mod_ctx() as ctx:
-                function(*args, **kwargs)
-                assert (
-                    ctx.module.operation.verify()
-                ), f"Verification failed for '{function.__name__}'"
-                mlir_module = ctx.module
-        else:
-            mlir_module = function(*args, **kwargs)
+        from .kernel import ExternalFunction
+
+        # Clear any instances from previous runs to make sure if the user provided any broken code we don't try to recompile it
+        ExternalFunction._instances.clear()
+
+        # Find ExternalFunction instances in arguments and kwargs
+        external_kernels = []
+        for arg in args:
+            if isinstance(arg, ExternalFunction):
+                external_kernels.append(arg)
+        for value in kwargs.values():
+            if isinstance(value, ExternalFunction):
+                external_kernels.append(value)
 
-        # Hash of the IR string
-        module_hash = hash_module(mlir_module)
+        # Execute the function to generate MLIR
+        try:
+            if is_placed:
+                with mlir_mod_ctx() as ctx:
+                    function(*args, **kwargs)
+                    assert (
+                        ctx.module.operation.verify()
+                    ), f"Verification failed for '{function.__name__}'"
+                    mlir_module = ctx.module
+            else:
+                mlir_module = function(*args, **kwargs)
+        except Exception as e:
+            raise
+
+        # Compile all ExternalFunction instances that were created during this JIT compilation
+        for func in ExternalFunction._instances:
+            if (
+                not hasattr(func, "_compiled") or not func._compiled
+            ):  # Don't compile if already compiled
+                external_kernels.append(func)
+
+        # Determine target architecture based on device type
+        try:
+            current_device = get_current_device()
+
+            # Determine target architecture based on device type
+            if isinstance(current_device, (NPU2, NPU2Col1)):
+                target_arch = "aie2p"
+            elif isinstance(current_device, (NPU1, NPU1Col1)):
+                target_arch = "aie2"
+            elif current_device in (AIEDevice.npu2, AIEDevice.npu2_1col):
+                target_arch = "aie2p"
+            elif current_device in (AIEDevice.npu1, AIEDevice.npu1_1col):
+                target_arch = "aie2"
+            else:
+                raise RuntimeError(f"Unsupported device type: {type(current_device)}")
+        except Exception as e:
+            raise
+
+        # Hash of the IR string, ExternalFunction compiler options, and target architecture
+        module_hash = hash_module(mlir_module, external_kernels, target_arch)
         kernel_dir = os.path.join(IRON_CACHE_DIR, f"{module_hash}")
         mlir_path = os.path.join(kernel_dir, "aie.mlir")
 
@@ -161,25 +214,120 @@ def decorator(*args, **kwargs):
         inst_exists = os.path.exists(inst_path)
 
         if not use_cache or not xclbin_exists or not inst_exists:
-            with open(mlir_path, "w", encoding="utf-8") as f:
-                print(mlir_module, file=f)
-            compile_mlir_module_to_binary(
-                mlir_module=mlir_module,
-                inst_path=inst_path,
-                xclbin_path=xclbin_path,
-            )
+            try:
+                with open(mlir_path, "w", encoding="utf-8") as f:
+                    print(mlir_module, file=f)
+
+                # Compile ExternalFunctions from inside the JIT compilation directory
+                for func in external_kernels:
+                    compile_external_kernel(func, kernel_dir, target_arch)
+
+                # Compile the MLIR module
+                compile_mlir_module_to_binary(
+                    mlir_module=mlir_module,
+                    inst_path=inst_path,
+                    xclbin_path=xclbin_path,
+                    work_dir=kernel_dir,
+                )
+            except Exception as e:
+                # Clean up cache directory on any compilation failure to avoid any corrupted objects in the cache
+                if os.path.exists(kernel_dir):
+                    shutil.rmtree(kernel_dir)
+                raise e
 
         kernel_name = "MLIR_AIE"
-        return NPUKernel(xclbin_path, inst_path, kernel_name=kernel_name)(
-            *args, **kwargs
-        )
+        try:
+            kernel = NPUKernel(xclbin_path, inst_path, kernel_name=kernel_name)
+            result = kernel(*args, **kwargs)
+            return result
+        except Exception as e:
+            raise
 
     return decorator
 
 
-def hash_module(module):
+def compile_external_kernel(func, kernel_dir, target_arch):
     """
-    Hash the MLIR module to create a unique identifier.
+    Compile an ExternalFunction to an object file in the kernel directory.
+
+    Args:
+        func: ExternalFunction instance to compile
+        kernel_dir: Directory to place the compiled object file
+        target_arch: Target architecture (e.g., "aie2" or "aie2p")
+    """
+    # Skip if already compiled
+    if hasattr(func, "_compiled") and func._compiled:
+        return
+
+    # Check if object file already exists in kernel directory
+    output_file = os.path.join(kernel_dir, func._object_file_name)
+    if os.path.exists(output_file):
+        return
+
+    # Create source file in kernel directory
+    source_file = os.path.join(kernel_dir, f"{func._name}.cc")
+
+    # Handle both source_string and source_file cases
+    if func._source_string is not None:
+        # Use source_string (write to file)
+        try:
+            with open(source_file, "w") as f:
+                f.write(func._source_string)
+        except Exception as e:
+            raise
+    elif func._source_file is not None:
+        # Use source_file (copy existing file)
+        # Check if source file exists before copying
+        if os.path.exists(func._source_file):
+            try:
+                shutil.copy2(func._source_file, source_file)
+            except Exception as e:
+                raise
+        else:
+            return
+    else:
+        raise ValueError("Neither source_string nor source_file is provided")
+
+    from .compile.compile import compile_cxx_core_function
+
+    try:
+        compile_cxx_core_function(
+            source_path=source_file,
+            target_arch=target_arch,
+            output_path=output_file,
+            include_dirs=func._include_dirs,
+            compile_args=func._compile_flags,
+            cwd=kernel_dir,
+            verbose=False,
+        )
+    except Exception as e:
+        raise
+
+    # Mark the function as compiled
+    func._compiled = True
+
+
+def hash_module(module, external_kernels=None, target_arch=None):
+    """
+    Hash the MLIR module and ExternalFunction compiler options to create a unique identifier.
     """
     mlir_str = str(module)
-    return hashlib.sha256(mlir_str.encode("utf-8")).hexdigest()[:16]
+
+    # Include ExternalFunction compiler options in the hash
+    if external_kernels:
+        compiler_options = []
+        for func in external_kernels:
+            compiler_options.extend(func._include_dirs)
+            compiler_options.extend(func._compile_flags)
+
+        # Create a combined string for hashing
+        combined_str = mlir_str + "|" + "|".join(compiler_options)
+    else:
+        combined_str = mlir_str
+
+    # Include target architecture in the hash
+    if target_arch:
+        combined_str += f"|target_arch={target_arch}"
+
+    hash_result = hashlib.sha256(combined_str.encode("utf-8")).hexdigest()[:16]
+    return hash_result