Add external C functions support for JIT

python/iron/__init__.py

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

python/iron/jit.py

@@ -15,6 +15,9 @@
 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
+from .config import get_current_device
+import subprocess
 
 
 # The `iron.jit` decorator below caches compiled kenrels inside the `IRON_CACHE_DIR` directory.
@@ -93,6 +96,9 @@ 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)}"
@@ -133,6 +139,22 @@ def jit(function=None, is_placed=True, use_cache=True):
 
     @functools.wraps(function)
     def decorator(*args, **kwargs):
+        # Import ExternalKernel at the top
+        from .kernel import ExternalKernel
+
+        # Clear any instances from previous runs to make sure if the user provided any broken code we don't try to recompile it
+        ExternalKernel._instances.clear()
+
+        # Find ExternalKernel instances in arguments and kwargs
+        external_kernels = []
+        for arg in args:
+            if isinstance(arg, ExternalKernel):
+                external_kernels.append(arg)
+        for value in kwargs.values():
+            if isinstance(value, ExternalKernel):
+                external_kernels.append(value)
+
+        # Execute the function to generate MLIR
         if is_placed:
             with mlir_mod_ctx() as ctx:
                 function(*args, **kwargs)
@@ -143,8 +165,15 @@ def decorator(*args, **kwargs):
         else:
             mlir_module = function(*args, **kwargs)
 
-        # Hash of the IR string
-        module_hash = hash_module(mlir_module)
+        # Compile all ExternalKernel instances that were created during this JIT compilation
+        for func in ExternalKernel._instances:
+            if (
+                not hasattr(func, "_compiled") or not func._compiled
+            ):  # Don't compile if already compiled
+                external_kernels.append(func)
+
+        # Hash of the IR string and ExternalKernel compiler options
+        module_hash = hash_module(mlir_module, external_kernels)
         kernel_dir = os.path.join(IRON_CACHE_DIR, f"{module_hash}")
         mlir_path = os.path.join(kernel_dir, "aie.mlir")
 
@@ -161,13 +190,28 @@ 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)
+
+                # Set cache directory for ExternalKernels and compile them
+                for func in external_kernels:
+                    # Compile the ExternalKernel directly in the kernel directory
+                    compile_external_kernel(func, kernel_dir)
+                # 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
+                if os.path.exists(kernel_dir):
+                    import shutil
+
+                    shutil.rmtree(kernel_dir)
+                raise e
 
         kernel_name = "MLIR_AIE"
         return NPUKernel(xclbin_path, inst_path, kernel_name=kernel_name)(
@@ -177,9 +221,112 @@ def decorator(*args, **kwargs):
     return decorator
 
 
-def hash_module(module):
+def compile_external_kernel(func, kernel_dir):
     """
-    Hash the MLIR module to create a unique identifier.
+    Compile an ExternalKernel to an object file in the kernel directory.
+
+    Args:
+        func: ExternalKernel instance to compile
+        kernel_dir: Directory to place the compiled object file
+    """
+    # 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)
+        with open(source_file, "w") as f:
+            f.write(func._source_string)
+    elif func._source_file is not None:
+        # Use source_file (copy existing file)
+        import shutil
+
+        # Check if source file exists before copying
+        if os.path.exists(func._source_file):
+            shutil.copy2(func._source_file, source_file)
+        else:
+            return
+    else:
+        raise ValueError("Neither source_string nor source_file is provided")
+
+    # Add device-specific flags based on actual device detection
+    current_device = get_current_device()
+
+    warning_flags = [
+        "-Wno-parentheses",
+        "-Wno-attributes",
+        "-Wno-macro-redefined",
+        "-Wno-empty-body",
+        "-Wno-missing-template-arg-list-after-template-kw",
+    ]
+
+    common_flags = ["-O2", "-std=c++20", "-DNDEBUG"] + warning_flags
+
+    # Build compilation command
+    cmd = [f"{os.environ.get('PEANO_INSTALL_DIR', '')}/bin/clang++"]
+
+    # Check device type and use appropriate flags
+    if isinstance(current_device, NPU2):
+        cmd.extend(common_flags + ["--target=aie2p-none-unknown-elf"])
+    elif isinstance(current_device, NPU1):
+        cmd.extend(common_flags + ["--target=aie2-none-unknown-elf"])
+    else:
+        raise RuntimeError(f"Unsupported device type: {type(current_device)}")
+
+    # Add include directories
+    for include_dir in func._include_dirs:
+        cmd.extend(["-I", include_dir])
+
+    # Add compilation flags
+    cmd.extend(func._compile_flags)
+
+    # Add source and output files
+    cmd.extend(["-c", source_file, "-o", output_file])
+
+    try:
+        result = subprocess.run(
+            cmd,
+            check=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+
+        # Mark the function as compiled
+        func._compiled = True
+    except subprocess.CalledProcessError as e:
+        error_msg = (
+            f"Compilation failed:\n{e}\n"
+            f"stdout:\n{e.stdout.decode() if e.stdout else 'No stdout'}\n"
+            f"stderr:\n{e.stderr.decode() if e.stderr else 'No stderr'}"
+        )
+        raise RuntimeError(error_msg)
+
+
+def hash_module(module, external_kernels=None):
+    """
+    Hash the MLIR module and ExternalKernel compiler options to create a unique identifier.
     """
     mlir_str = str(module)
-    return hashlib.sha256(mlir_str.encode("utf-8")).hexdigest()[:16]
+
+    # Include ExternalKernel compiler options in the hash
+    if external_kernels:
+        compiler_options = []
+        for func in external_kernels:
+            # Include include_dirs and compile_flags in the hash
+            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)
+        return hashlib.sha256(combined_str.encode("utf-8")).hexdigest()[:16]
+    else:
+        return hashlib.sha256(mlir_str.encode("utf-8")).hexdigest()[:16]

python/iron/kernel.py

@@ -7,6 +7,7 @@
 # (c) Copyright 2024 Advanced Micro Devices, Inc.
 
 import numpy as np
+import os
 
 from .. import ir  # type: ignore
 from ..extras.dialects.ext.func import FuncOp  # type: ignore
@@ -15,7 +16,36 @@
 from .resolvable import Resolvable
 
 
-class Kernel(Resolvable):
+class BaseKernel(Resolvable):
+    """Base class for kernel-like objects that resolve to FuncOp."""
+
+    def __init__(self, name: str, arg_types: list[type[np.ndarray] | np.dtype] = []):
+        """Initialize base kernel.
+
+        Args:
+            name (str): The name of the function
+            arg_types (list[type[np.ndarray] | np.dtype], optional): The type signature of the function. Defaults to [].
+        """
+        self._name = name
+        self._arg_types = arg_types
+        self._op: FuncOp | None = None
+
+    def resolve(
+        self,
+        loc: ir.Location | None = None,
+        ip: ir.InsertionPoint | None = None,
+    ) -> None:
+        """Resolve the kernel to a FuncOp. Must be implemented by subclasses."""
+        raise NotImplementedError("Subclasses must implement resolve()")
+
+    def __call__(self, *args, **kwargs):
+        """Call the kernel with the given arguments."""
+        if not self._op:
+            raise ValueError("Need to resolve kernel before it can be called")
+        call(self._op, args, **kwargs)
+
+
+class Kernel(BaseKernel):
     def __init__(
         self,
         name: str,
@@ -30,10 +60,8 @@ def __init__(
             bin_name (str): The name of the binary (used for linking to a compute core)
             arg_types (list[type[np.ndarray]  |  np.dtype], optional): The type signature of the function. Defaults to [].
         """
-        self._name = name
+        super().__init__(name, arg_types)
         self._bin_name = bin_name
-        self._arg_types = arg_types
-        self._op: FuncOp | None = None
 
     @property
     def bin_name(self) -> str:
@@ -47,7 +75,115 @@ def resolve(
         if not self._op:
             self._op = external_func(self._name, inputs=self._arg_types)
 
+
+class ExternalKernel(BaseKernel):
+    _object_files = set()
+    _instances = set()
+
+    def __init__(
+        self,
+        name: str,
+        source_file: str | None = None,
+        source_string: str | None = None,
+        arg_types: list[type[np.ndarray] | np.dtype] = [],
+        include_dirs: list[str] = [],
+        compile_flags: list[str] = [],
+    ) -> None:
+        """An ExternalKernel is a C++ source file that gets compiled to an object file and eventually resolves to a FuncOp.
+        If it is called, a CallOp will be generated.
+
+        Args:
+            name (str): The name of the function
+            source_file (str): Path to the C++ source file
+            source_string (str): C++ source code as a string
+            arg_types (list[type[np.ndarray] | np.dtype], optional): The type signature of the function. Defaults to [].
+            include_dirs (list[str], optional): Additional include directories. Defaults to [].
+            compile_flags (list[str], optional): Additional compilation flags. Defaults to [].
+        """
+        super().__init__(name, arg_types)
+        self._setup_source(source_file, source_string)
+        self._include_dirs = include_dirs
+        self._compile_flags = compile_flags
+        self._object_file_name = f"{self._name}.o"
+        self._compiled = False
+
+        # Track this instance for JIT compilation
+        ExternalKernel._instances.add(self)
+
+    def _setup_source(self, source_file: str | None, source_string: str | None) -> None:
+        """Set up the source file for compilation."""
+        if source_file is not None:
+            self._source_file = source_file
+            self._source_string = None
+        else:
+            if source_string is None:
+                raise ValueError("source_file or source_string must be provided")
+            self._source_file = None
+            self._source_string = source_string
+
+    def __enter__(self):
+        """Enter the context."""
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Exit the context."""
+        pass
+
+    @property
+    def bin_name(self) -> str:
+        return self._object_file_name
+
+    def tile_size(self, arg_index: int = 0) -> int:
+        """Get the tile size from the specified array argument type.
+
+        Args:
+            arg_index (int): Index of the argument to get tile size from. Defaults to 0.
+
+        Returns:
+            int: The tile size (first dimension) of the specified argument.
+        """
+        if not self._arg_types:
+            raise ValueError("No argument types defined")
+        if arg_index >= len(self._arg_types):
+            raise ValueError(
+                f"Argument index {arg_index} out of range (max: {len(self._arg_types) - 1})"
+            )
+
+        arg = self._arg_types[arg_index]
+
+        # Handle numpy array types like np.ndarray[(16,), np.dtype[np.int32]]
+        if hasattr(arg, "__args__") and len(arg.__args__) > 0:
+            # For types like np.ndarray[(16,), np.dtype[np.int32]], the shape is in __args__[0]
+            shape_arg = arg.__args__[0]
+            if isinstance(shape_arg, tuple) and len(shape_arg) > 0:
+                return shape_arg[0]
+
+        # Handle MLIR types like MemRefType(memref<16xi32>)
+        if (
+            hasattr(arg, "shape")
+            and hasattr(arg.shape, "__len__")
+            and len(arg.shape) > 0
+        ):
+            return arg.shape[0]
+
+        raise ValueError(
+            f"Argument {arg_index} does not have a shape or is not an array type"
+        )
+
+    def arg_types(self) -> list:
+        """Get the argument types of the ExternalKernel."""
+        return self._arg_types.copy()
+
+    def resolve(
+        self,
+        loc: ir.Location | None = None,
+        ip: ir.InsertionPoint | None = None,
+    ) -> None:
+        if not self._op:
+            # Create the external function
+            self._op = external_func(self._name, inputs=self._arg_types)
+
     def __call__(self, *args, **kwargs):
         if not self._op:
-            raise ValueError("Need to resolve Kernel before it can be called")
+            raise ValueError("Need to resolve ExternalKernel before it can be called")
         call(self._op, args, **kwargs)