From 9ca1889d1c10a269b604865828f4cb5a5b7e7100 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 03:12:03 -0700
Subject: [PATCH 01/37] Add `ExternalKernel` support for JIT

---
 python/iron/__init__.py |   2 +-
 python/iron/jit.py      | 179 +++++++++++++++++++++++++++++++++++++---
 python/iron/kernel.py   | 146 ++++++++++++++++++++++++++++++--
 python/iron/tensor.py   |  13 +++
 python/iron/worker.py   |   4 +-
 python/utils/compile.py |  10 ++-
 6 files changed, 333 insertions(+), 21 deletions(-)

diff --git a/python/iron/__init__.py b/python/iron/__init__.py
index c99407ccac3..3c8be7f8a0d 100644
--- a/python/iron/__init__.py
+++ b/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
diff --git a/python/iron/jit.py b/python/iron/jit.py
index f4567bde404..810c6b36bb1 100644
--- a/python/iron/jit.py
+++ b/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,120 @@ 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")
+
+    # Build compilation command
+    cmd = [
+        f"{os.environ.get('PEANO_INSTALL_DIR', '')}/bin/clang++",
+        "-O2",
+        "-std=c++20",
+        "--target=aie2-none-unknown-elf",
+        "-Wno-parentheses",
+        "-Wno-attributes",
+        "-Wno-macro-redefined",
+        "-Wno-empty-body",
+        "-DNDEBUG",
+    ]
+
+    # Add AIEOPT include directory
+    try:
+        aieopt_path = subprocess.check_output(["which", "aie-opt"], text=True).strip()
+        aieopt_dir = os.path.dirname(os.path.dirname(os.path.realpath(aieopt_path)))
+        cmd.extend(["-I", f"{aieopt_dir}/include"])
+    except subprocess.CalledProcessError:
+        pass
+
+    # Add device-specific flags based on actual device detection
+    current_device = get_current_device()
+
+    # Check device type and use appropriate flags
+    if isinstance(current_device, NPU2):
+        cmd.extend(os.environ.get("PEANOWRAP2P_FLAGS", "").split())
+    elif isinstance(current_device, NPU1):
+        cmd.extend(os.environ.get("PEANOWRAP2_FLAGS", "").split())
+    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]
diff --git a/python/iron/kernel.py b/python/iron/kernel.py
index a46304d84c4..52c9d37a65c 100644
--- a/python/iron/kernel.py
+++ b/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)
diff --git a/python/iron/tensor.py b/python/iron/tensor.py
index 804d742e1c5..358cdd13ec3 100644
--- a/python/iron/tensor.py
+++ b/python/iron/tensor.py
@@ -196,6 +196,19 @@ def numpy(self):
             self.__sync_from_device()
         return self.data
 
+    def fill_(self, value):
+        """
+        Fills the tensor with a scalar value (in-place operation).
+
+        Parameters:
+            value: The scalar value to fill the tensor with.
+
+        Note: For NPU tensors, this method syncs the filled data to device after modification.
+        """
+        self.data.fill(value)
+        if self.device == "npu":
+            self.__sync_to_device()
+
     @staticmethod
     def _ctype_from_dtype(dtype):
         """
diff --git a/python/iron/worker.py b/python/iron/worker.py
index 78b6f837f61..31835866795 100644
--- a/python/iron/worker.py
+++ b/python/iron/worker.py
@@ -16,7 +16,7 @@
 from .device import PlacementTile, AnyComputeTile, Tile
 from .dataflow.objectfifo import ObjectFifoHandle, ObjectFifo
 from .dataflow.endpoint import ObjectFifoEndpoint
-from .kernel import Kernel
+from .kernel import Kernel, ExternalKernel
 from .globalbuffer import GlobalBuffer
 from .resolvable import Resolvable
 
@@ -86,7 +86,7 @@ def do_nothing_core_fun(*args) -> None:
 
         # Check arguments to the core. Some information is saved for resolution.
         for arg in self.fn_args:
-            if isinstance(arg, Kernel):
+            if isinstance(arg, (Kernel, ExternalKernel)):
                 bin_names.add(arg.bin_name)
             elif isinstance(arg, ObjectFifoHandle):
                 arg.endpoint = self
diff --git a/python/utils/compile.py b/python/utils/compile.py
index ec1406b6f90..504d00e50a4 100644
--- a/python/utils/compile.py
+++ b/python/utils/compile.py
@@ -10,7 +10,9 @@
 import aie.compiler.aiecc.main as aiecc
 
 
-def compile_mlir_module_to_binary(mlir_module: str, inst_path: str, xclbin_path: str):
+def compile_mlir_module_to_binary(
+    mlir_module: str, inst_path: str, xclbin_path: str, work_dir: str = None
+):
     """
     Compile an MLIR module to instruction and xclbin files using the aiecc module.
 
@@ -18,6 +20,7 @@ def compile_mlir_module_to_binary(mlir_module: str, inst_path: str, xclbin_path:
         mlir_module (str): MLIR module to compile.
         inst_path (str): Path to the instruction binary file.
         xclbin_path (str): Path to the xclbin file.
+        work_dir (str, optional): Working directory for compilation. Defaults to None.
     """
 
     args = [
@@ -29,6 +32,11 @@ def compile_mlir_module_to_binary(mlir_module: str, inst_path: str, xclbin_path:
         f"--xclbin-name={xclbin_path}",
         f"--npu-insts-name={inst_path}",
     ]
+
+    # Add working directory if specified
+    if work_dir:
+        args.append(f"--tmpdir={work_dir}")
+
     try:
         aiecc.run(mlir_module, args)
     except Exception as e:

From 058ad42e6069f5e9b3c86bf8d5d21fbb818ffa9b Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 03:12:12 -0700
Subject: [PATCH 02/37] Add tests

---
 test/python/jit_extern.py            | 844 +++++++++++++++++++++++++++
 test/python/jit_extern_inside_jit.py | 307 ++++++++++
 2 files changed, 1151 insertions(+)
 create mode 100644 test/python/jit_extern.py
 create mode 100644 test/python/jit_extern_inside_jit.py

diff --git a/test/python/jit_extern.py b/test/python/jit_extern.py
new file mode 100644
index 00000000000..527c573ca8a
--- /dev/null
+++ b/test/python/jit_extern.py
@@ -0,0 +1,844 @@
+# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#
+# (c) Copyright 2025 AMD Inc.
+
+# RUN: %run_on_npu1% %pytest %s
+# RUN: %run_on_npu2% %pytest %s
+
+import numpy as np
+import os
+import tempfile
+import shutil
+import pytest
+
+import aie.iron as iron
+from aie.iron import ExternalKernel, jit
+from aie.iron import ObjectFifo, Worker, Runtime, Program
+from aie.iron.placers import SequentialPlacer
+from aie.iron.controlflow import range_
+
+
+@jit(is_placed=False)
+def transform(input, output, func):
+    """Transform kernel that applies a function to input tensor and stores result in output tensor."""
+    if input.shape != output.shape:
+        raise ValueError(
+            f"Input shapes are not the equal ({input.shape} != {output.shape})."
+        )
+    num_elements = np.size(input)
+
+    # Extract tile size from ExternalKernel (using first argument)
+    tile_size = func.tile_size(0)
+
+    # Assert that input and output arrays have the same tile size
+    assert func.tile_size(0) == func.tile_size(
+        1
+    ), f"Input and output tile sizes must match: {func.tile_size(0)} != {func.tile_size(1)}"
+
+    if num_elements % tile_size != 0:
+        raise ValueError(
+            f"Number of elements ({num_elements}) must be a multiple of {tile_size}."
+        )
+    num_tiles = num_elements // tile_size
+
+    if input.dtype != output.dtype:
+        raise ValueError(
+            f"Input data types are not the same ({input.dtype} != {output.dtype})."
+        )
+
+    dtype = input.dtype
+
+    # Define tensor types
+    tensor_ty = np.ndarray[(num_elements,), np.dtype[dtype]]
+    tile_ty = np.ndarray[(tile_size,), np.dtype[dtype]]
+
+    # AIE-array data movement with object fifos
+    of_in = ObjectFifo(tile_ty, name="in")
+    of_out = ObjectFifo(tile_ty, name="out")
+
+    # Define a task that will run on a compute tile
+    def core_body(of_in, of_out, func_to_apply):
+        # Extract tile size from ExternalKernel (using first argument)
+        tile_size = func_to_apply.tile_size(0)
+
+        # Number of sub-vector "tile" iterations
+        for i in range_(num_tiles):
+            elem_in = of_in.acquire(1)
+            elem_out = of_out.acquire(1)
+            func_to_apply(elem_in, elem_out, tile_size)
+            of_in.release(1)
+            of_out.release(1)
+
+    # Create a worker to run the task on a compute tile
+    worker = Worker(core_body, fn_args=[of_in.cons(), of_out.prod(), func])
+
+    # Runtime operations to move data to/from the AIE-array
+    rt = Runtime()
+    with rt.sequence(tensor_ty, tensor_ty) as (A, B):
+        rt.start(worker)
+        rt.fill(of_in.prod(), A)
+        rt.drain(of_out.cons(), B, wait=True)
+
+    # Place program components (assign them resources on the device) and generate an MLIR module
+    return Program(iron.get_current_device(), rt).resolve_program(SequentialPlacer())
+
+
+def test_simple_add_one():
+    """Test basic ExternalKernel with simple add_one operation."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create ExternalKernel for adding one
+    add_one = ExternalKernel(
+        "add_one",
+        source_string="""extern "C" {
+            void add_one(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + 1;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+    )
+
+    # Apply the transform
+    transform(input_tensor, output_tensor, add_one)
+
+    # Verify results
+    expected = initial_tensor + 1
+    actual = output_tensor.numpy()
+    np.testing.assert_array_equal(actual, expected)
+
+
+@pytest.mark.parametrize("tile_size", [8, 16, 32, 64])
+def test_different_tile_sizes(tile_size):
+    """Test ExternalKernel with different tile sizes."""
+    # Create input and output tensors
+    num_elements = 1024
+    input_tensor = iron.randint(0, 100, (num_elements,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((num_elements,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create ExternalKernel with specific tile size
+    add_one = ExternalKernel(
+        "add_one",
+        source_string="""extern "C" {
+            void add_one(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + 1;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(tile_size,), np.dtype[np.int32]],
+            np.ndarray[(tile_size,), np.dtype[np.int32]],
+            np.int32,
+        ],
+    )
+
+    # Apply the transform
+    transform(input_tensor, output_tensor, add_one)
+
+    # Verify results
+    expected = initial_tensor + 1
+    actual = output_tensor.numpy()
+    np.testing.assert_array_equal(actual, expected)
+
+
+@pytest.mark.parametrize(
+    "dtype,c_type",
+    [
+        (np.int32, "int"),
+        (np.float32, "float"),
+    ],
+)
+def test_different_data_types(dtype, c_type):
+    """Test ExternalKernel with different data types."""
+    # Create input and output tensors
+    input_tensor = iron.rand((1024,), dtype=dtype, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=dtype, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create ExternalKernel with specific data type
+    add_one = ExternalKernel(
+        "add_one",
+        source_string=f"""extern "C" {{
+            void add_one({c_type}* input, {c_type}* output, int tile_size) {{
+                for (int i = 0; i < tile_size; i++) {{
+                    output[i] = input[i] + 1.0f;
+                }}
+            }}
+        }}""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[dtype]],
+            np.ndarray[(16,), np.dtype[dtype]],
+            np.int32,
+        ],
+    )
+
+    # Apply the transform
+    transform(input_tensor, output_tensor, add_one)
+
+    # Verify results
+    expected = initial_tensor + 1.0
+    actual = output_tensor.numpy()
+    np.testing.assert_array_almost_equal(actual, expected, decimal=5)
+
+
+@pytest.mark.parametrize("value", [5, 42])
+def test_define_values(value):
+    """Test ExternalKernel with different define values."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    add_value = ExternalKernel(
+        "add_value",
+        source_string="""extern "C" {
+            void add_value(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + ADD_VALUE;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+        compile_flags=[f"-DADD_VALUE={value}"],
+    )
+
+    # Apply the transform
+    transform(input_tensor, output_tensor, add_value)
+
+    # Verify results
+    expected = initial_tensor + value
+    actual = output_tensor.numpy()
+    np.testing.assert_array_equal(actual, expected)
+
+
+def test_multiple_defines():
+    """Test ExternalKernel with multiple defines."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create ExternalKernel with multiple defines
+    complex_op = ExternalKernel(
+        "complex_op",
+        source_string="""extern "C" {
+            void complex_op(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    #ifdef DEBUG
+                    output[i] = input[i] + ADD_VALUE + DEBUG_OFFSET;
+                    #else
+                    output[i] = input[i] + ADD_VALUE;
+                    #endif
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+        compile_flags=["-DADD_VALUE=5", "-DDEBUG", "-DDEBUG_OFFSET=10"],
+    )
+
+    # Apply the transform
+    transform(input_tensor, output_tensor, complex_op)
+
+    # Verify results (should add 15: 5 + 10 due to DEBUG define)
+    expected = initial_tensor + 15
+    actual = output_tensor.numpy()
+    np.testing.assert_array_equal(actual, expected)
+
+
+def test_include_directories():
+    """Test ExternalKernel with include directories."""
+    # Create a temporary directory with a header file
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Create a header file
+        header_file = os.path.join(temp_dir, "math_ops.h")
+        with open(header_file, "w") as f:
+            f.write(
+                """
+#ifndef MATH_OPS_H
+#define MATH_OPS_H
+
+#define ADD_VALUE 42
+
+#endif
+"""
+            )
+
+        # Create input and output tensors
+        input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+        output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+        initial_tensor = input_tensor.numpy().copy()
+
+        # Create ExternalKernel that includes the header
+        add_value = ExternalKernel(
+            "add_value",
+            source_string="""extern "C" {
+                #include "math_ops.h"
+                void add_value(int* input, int* output, int tile_size) {
+                    for (int i = 0; i < tile_size; i++) {
+                        output[i] = input[i] + ADD_VALUE;
+                    }
+                }
+            }""",
+            arg_types=[
+                np.ndarray[(16,), np.dtype[np.int32]],
+                np.ndarray[(16,), np.dtype[np.int32]],
+                np.int32,
+            ],
+            include_dirs=[temp_dir],
+        )
+
+        # Apply the transform
+        transform(input_tensor, output_tensor, add_value)
+
+        # Verify results
+        expected = initial_tensor + 42
+        actual = output_tensor.numpy()
+        np.testing.assert_array_equal(actual, expected)
+
+
+def test_multiple_include_directories():
+    """Test ExternalKernel with multiple include directories."""
+    # Create temporary directories with header files
+    with tempfile.TemporaryDirectory() as temp_dir1, tempfile.TemporaryDirectory() as temp_dir2:
+        # Create header files
+        header1 = os.path.join(temp_dir1, "ops1.h")
+        with open(header1, "w") as f:
+            f.write("#define VALUE1 10\n")
+
+        header2 = os.path.join(temp_dir2, "ops2.h")
+        with open(header2, "w") as f:
+            f.write("#define VALUE2 20\n")
+
+        # Create input and output tensors
+        input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+        output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+        initial_tensor = input_tensor.numpy().copy()
+
+        # Create ExternalKernel that includes both headers
+        add_values = ExternalKernel(
+            "add_values",
+            source_string="""extern "C" {
+                #include "ops1.h"
+                #include "ops2.h"
+                void add_values(int* input, int* output, int tile_size) {
+                    for (int i = 0; i < tile_size; i++) {
+                        output[i] = input[i] + VALUE1 + VALUE2;
+                    }
+                }
+            }""",
+            arg_types=[
+                np.ndarray[(16,), np.dtype[np.int32]],
+                np.ndarray[(16,), np.dtype[np.int32]],
+                np.int32,
+            ],
+            include_dirs=[temp_dir1, temp_dir2],
+        )
+
+        # Apply the transform
+        transform(input_tensor, output_tensor, add_values)
+
+        # Verify results
+        expected = initial_tensor + 30  # 10 + 20
+        actual = output_tensor.numpy()
+        np.testing.assert_array_equal(actual, expected)
+
+
+def test_caching_same_source():
+    """Test that same source code produces same cached result."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create two ExternalKernels with identical source
+    add_one_1 = ExternalKernel(
+        "add_one_1",
+        source_string="""extern "C" {
+            void add_one_1(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + 1;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+    )
+
+    add_one_2 = ExternalKernel(
+        "add_one_2",
+        source_string="""extern "C" {
+            void add_one_2(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + 1;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+    )
+
+    # Apply both transforms
+    transform(input_tensor, output_tensor, add_one_1)
+    result1 = output_tensor.numpy().copy()
+
+    output_tensor.fill_(0)
+    transform(input_tensor, output_tensor, add_one_2)
+    result2 = output_tensor.numpy()
+
+    # Verify both produce same results
+    np.testing.assert_array_equal(result1, result2)
+
+
+def test_context_manager():
+    """Test ExternalKernel with context manager syntax."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create ExternalKernel and use it with context manager
+    with ExternalKernel(
+        "add_one_context",
+        source_string="""extern "C" {
+            void add_one_context(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + 1;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+    ) as add_one:
+        # Apply the transform
+        transform(input_tensor, output_tensor, add_one)
+
+    # Verify results
+    expected = initial_tensor + 1
+    actual = output_tensor.numpy()
+    np.testing.assert_array_equal(actual, expected)
+
+
+def test_context_manager_with_compiler_options():
+    """Test ExternalKernel with context manager and compiler options."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create ExternalKernel with compiler options using context manager
+    with ExternalKernel(
+        "add_value_context",
+        source_string="""extern "C" {
+            void add_value_context(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + ADD_VALUE;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+        compile_flags=["-DADD_VALUE=42"],
+    ) as add_value:
+        # Apply the transform
+        transform(input_tensor, output_tensor, add_value)
+
+    # Verify results
+    expected = initial_tensor + 42
+    actual = output_tensor.numpy()
+    np.testing.assert_array_equal(actual, expected)
+
+
+def test_source_file():
+    """Test ExternalKernel with source_file instead of source_string."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create a temporary source file
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".cc", delete=False) as f:
+        source_content = """extern "C" {
+            void add_one_from_file(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + 1;
+                }
+            }
+        }"""
+        f.write(source_content)
+        source_file_path = f.name
+
+    try:
+        # Create ExternalKernel using source_file
+        add_one_from_file = ExternalKernel(
+            "add_one_from_file",
+            source_file=source_file_path,
+            arg_types=[
+                np.ndarray[(16,), np.dtype[np.int32]],
+                np.ndarray[(16,), np.dtype[np.int32]],
+                np.int32,
+            ],
+        )
+
+        # Apply the transform
+        transform(input_tensor, output_tensor, add_one_from_file)
+
+        # Verify results
+        expected = initial_tensor + 1
+        actual = output_tensor.numpy()
+        np.testing.assert_array_equal(actual, expected)
+
+    finally:
+        # Clean up the temporary file
+        os.unlink(source_file_path)
+
+
+def test_source_file_with_compiler_options():
+    """Test ExternalKernel with source_file and compiler options."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create a temporary source file with defines
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".cc", delete=False) as f:
+        source_content = """extern "C" {
+            void add_value_from_file(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + ADD_VALUE;
+                }
+            }
+        }"""
+        f.write(source_content)
+        source_file_path = f.name
+
+    try:
+        # Create ExternalKernel using source_file with compiler options
+        add_value_from_file = ExternalKernel(
+            "add_value_from_file",
+            source_file=source_file_path,
+            arg_types=[
+                np.ndarray[(16,), np.dtype[np.int32]],
+                np.ndarray[(16,), np.dtype[np.int32]],
+                np.int32,
+            ],
+            compile_flags=["-DADD_VALUE=25"],
+        )
+
+        # Apply the transform
+        transform(input_tensor, output_tensor, add_value_from_file)
+
+        # Verify results
+        expected = initial_tensor + 25
+        actual = output_tensor.numpy()
+        np.testing.assert_array_equal(actual, expected)
+
+    finally:
+        # Clean up the temporary file
+        os.unlink(source_file_path)
+
+
+def test_transform_with_internal_func():
+    """Test transform function that creates ExternalKernel internally."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create ExternalKernel dynamically but pass it as argument
+    internal_func = ExternalKernel(
+        "internal_add_one",
+        source_string="""extern "C" {
+            void internal_add_one(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + 1;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+    )
+
+    # Apply the transform (ExternalKernel is passed as argument)
+    transform(input_tensor, output_tensor, internal_func)
+
+    # Verify results
+    expected = initial_tensor + 1
+    actual = output_tensor.numpy()
+    np.testing.assert_array_equal(actual, expected)
+
+
+def test_caching_different_flags():
+    """Test that different compile flags produce different cached results."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create ExternalKernels with same source but different flags
+    add_value_5 = ExternalKernel(
+        "add_value",
+        source_string="""extern "C" {
+            void add_value(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + ADD_VALUE;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+        compile_flags=["-DADD_VALUE=5"],
+    )
+
+    add_value_10 = ExternalKernel(
+        "add_value",
+        source_string="""extern "C" {
+            void add_value(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + ADD_VALUE;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+        compile_flags=["-DADD_VALUE=10"],
+    )
+
+    # Apply transforms
+    transform(input_tensor, output_tensor, add_value_5)
+    result_5 = output_tensor.numpy().copy()
+
+    output_tensor.fill_(0)
+    transform(input_tensor, output_tensor, add_value_10)
+    result_10 = output_tensor.numpy()
+
+    # Verify different results
+    expected_5 = initial_tensor + 5
+    expected_10 = initial_tensor + 10
+
+    np.testing.assert_array_equal(result_5, expected_5)
+    np.testing.assert_array_equal(result_10, expected_10)
+    np.testing.assert_raises(
+        AssertionError, np.testing.assert_array_equal, result_5, result_10
+    )
+
+
+@pytest.mark.parametrize(
+    "invalid_source",
+    [
+        # Missing semicolon
+        """extern "C" {
+        void invalid_func(int* input, int* output, int tile_size) {
+            for (int i = 0; i < tile_size; i++) {
+                output[i] = input[i] + 1  // Missing semicolon
+            }
+        }
+    }""",
+        # Undefined variable
+        """extern "C" {
+        void invalid_func(int* input, int* output, int tile_size) {
+            for (int i = 0; i < tile_size; i++) {
+                output[i] = input[i] + undefined_var;
+            }
+        }
+    }""",
+        # Syntax error
+        """extern "C" {
+        void invalid_func(int* input, int* output, int tile_size) {
+            for (int i = 0; i < tile_size; i++) {
+                output[i] = input[i] + 1;
+            }  // Missing closing brace
+    }""",
+    ],
+)
+def test_invalid_source(invalid_source):
+    """Test error handling for invalid C++ source."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+
+    # Create ExternalKernel with invalid C++ source
+    invalid_func = ExternalKernel(
+        "invalid_func",
+        source_string=invalid_source,
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+    )
+
+    # Should raise an error during compilation
+    with pytest.raises(Exception):
+        transform(input_tensor, output_tensor, invalid_func)
+
+
+@pytest.mark.parametrize(
+    "input_tile_size,output_tile_size",
+    [
+        (16, 32),  # Different tile sizes
+        (8, 16),  # Different tile sizes
+        (64, 32),  # Different tile sizes
+    ],
+)
+def test_mismatched_tile_sizes(input_tile_size, output_tile_size):
+    """Test error handling for mismatched tile sizes."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+
+    # Create ExternalKernel with mismatched tile sizes
+    mismatched_func = ExternalKernel(
+        "mismatched_func",
+        source_string="""extern "C" {
+            void mismatched_func(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + 1;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(input_tile_size,), np.dtype[np.int32]],
+            np.ndarray[(output_tile_size,), np.dtype[np.int32]],
+            np.int32,
+        ],
+    )
+
+    # Should raise an assertion error
+    with pytest.raises(AssertionError, match="Input and output tile sizes must match"):
+        transform(input_tensor, output_tensor, mismatched_func)
+
+
+@pytest.mark.parametrize(
+    "invalid_include",
+    [
+        "/nonexistent/directory",
+        "/tmp/nonexistent_header_dir",
+        "/usr/local/include/nonexistent",
+    ],
+)
+def test_invalid_include_directory(invalid_include):
+    """Test error handling for invalid include directory."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+
+    # Create ExternalKernel with invalid include directory
+    invalid_include_func = ExternalKernel(
+        "invalid_include_func",
+        source_string="""extern "C" {
+            #include "nonexistent.h"
+            void invalid_include_func(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + 1;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+        include_dirs=[invalid_include],
+    )
+
+    # Should raise an error during compilation
+    with pytest.raises(Exception):
+        transform(input_tensor, output_tensor, invalid_include_func)
+
+
+@pytest.mark.parametrize(
+    "compile_flags,expected_value",
+    [
+        (["-DADD_VALUE=5"], 5),
+        (["-DADD_VALUE=10", "-DMULTIPLIER=2"], 20),  # 10 * 2
+        (["-DADD_VALUE=3", "-DOFFSET=7"], 10),  # 3 + 7
+        (["-DADD_VALUE=1", "-DDEBUG", "-DDEBUG_OFFSET=9"], 10),  # 1 + 9 (DEBUG enabled)
+    ],
+)
+def test_compiler_flag_combinations(compile_flags, expected_value):
+    """Test ExternalKernel with different combinations of compiler flags."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Create source that uses the defines
+    source_template = """extern "C" {
+        void complex_op(int* input, int* output, int tile_size) {
+            for (int i = 0; i < tile_size; i++) {
+                #ifdef MULTIPLIER
+                output[i] = input[i] + ADD_VALUE * MULTIPLIER;
+                #elif defined(DEBUG)
+                output[i] = input[i] + ADD_VALUE + DEBUG_OFFSET;
+                #elif defined(OFFSET)
+                output[i] = input[i] + ADD_VALUE + OFFSET;
+                #else
+                output[i] = input[i] + ADD_VALUE;
+                #endif
+            }
+        }
+    }"""
+
+    complex_op = ExternalKernel(
+        "complex_op",
+        source_string=source_template,
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+        compile_flags=compile_flags,
+    )
+
+    # Apply the transform
+    transform(input_tensor, output_tensor, complex_op)
+
+    # Verify results
+    expected = initial_tensor + expected_value
+    actual = output_tensor.numpy()
+    np.testing.assert_array_equal(actual, expected)
+
+
+if __name__ == "__main__":
+    test_simple_add_one()
diff --git a/test/python/jit_extern_inside_jit.py b/test/python/jit_extern_inside_jit.py
new file mode 100644
index 00000000000..709569de2dd
--- /dev/null
+++ b/test/python/jit_extern_inside_jit.py
@@ -0,0 +1,307 @@
+# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#
+# (c) Copyright 2025 AMD Inc.
+
+# RUN: %run_on_npu1% %pytest %s
+# RUN: %run_on_npu2% %pytest %s
+
+import numpy as np
+import os
+import tempfile
+import shutil
+import pytest
+
+import aie.iron as iron
+from aie.iron import ExternalKernel, jit
+from aie.iron import ObjectFifo, Worker, Runtime, Program
+from aie.iron.placers import SequentialPlacer
+from aie.iron.controlflow import range_
+
+
+@jit(is_placed=False)
+def transform_with_internal_func_with_options(input, output):
+    """Transform kernel that creates ExternalKernel internally with compiler options."""
+    if input.shape != output.shape:
+        raise ValueError(
+            f"Input shapes are not the equal ({input.shape} != {output.shape})."
+        )
+    num_elements = np.size(input)
+
+    # Create ExternalKernel inside the transform with compiler options
+    internal_func = ExternalKernel(
+        "internal_add_value",
+        source_string="""extern "C" {
+            void internal_add_value(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + ADD_VALUE;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+        compile_flags=["-DADD_VALUE=1"],
+    )
+
+    # Extract tile size from ExternalKernel
+    tile_size = internal_func.tile_size(0)
+
+    if num_elements % tile_size != 0:
+        raise ValueError(
+            f"Number of elements ({num_elements}) must be a multiple of {tile_size}."
+        )
+    num_tiles = num_elements // tile_size
+
+    if input.dtype != output.dtype:
+        raise ValueError(
+            f"Input data types are not the same ({input.dtype} != {output.dtype})."
+        )
+
+    dtype = input.dtype
+
+    # Define tensor types
+    tensor_ty = np.ndarray[(num_elements,), np.dtype[dtype]]
+    tile_ty = np.ndarray[(tile_size,), np.dtype[dtype]]
+
+    # AIE-array data movement with object fifos
+    of_in = ObjectFifo(tile_ty, name="in")
+    of_out = ObjectFifo(tile_ty, name="out")
+
+    # Define a task that will run on a compute tile
+    def core_body(of_in, of_out, func_to_apply):
+        # Extract tile size from ExternalKernel
+        tile_size = func_to_apply.tile_size(0)
+
+        # Number of sub-vector "tile" iterations
+        for i in range_(num_tiles):
+            elem_in = of_in.acquire(1)
+            elem_out = of_out.acquire(1)
+            func_to_apply(elem_in, elem_out, tile_size)
+            of_in.release(1)
+            of_out.release(1)
+
+    # Create a worker to run the task on a compute tile
+    worker = Worker(core_body, fn_args=[of_in.cons(), of_out.prod(), internal_func])
+
+    # Runtime operations to move data to/from the AIE-array
+    rt = Runtime()
+    with rt.sequence(tensor_ty, tensor_ty) as (A, B):
+        rt.start(worker)
+        rt.fill(of_in.prod(), A)
+        rt.drain(of_out.cons(), B, wait=True)
+
+    # Place program components and generate an MLIR module
+    return Program(iron.get_current_device(), rt).resolve_program(SequentialPlacer())
+
+
+@jit(is_placed=False)
+def transform_with_internal_func_from_file(input, output):
+    """Transform kernel that creates ExternalKernel internally from a file."""
+    if input.shape != output.shape:
+        raise ValueError(
+            f"Input shapes are not the equal ({input.shape} != {output.shape})."
+        )
+    num_elements = np.size(input)
+
+    # Create a temporary file with the source code inside the function
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".cc", delete=False) as f:
+        f.write(
+            """extern "C" {
+            void internal_add_from_file(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + 42;
+                }
+            }
+        }"""
+        )
+        temp_file_path = f.name
+
+    # Create ExternalKernel inside the transform from a file
+    internal_func = ExternalKernel(
+        "internal_add_from_file",
+        source_file=temp_file_path,
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+    )
+
+    # Extract tile size from ExternalKernel
+    tile_size = internal_func.tile_size(0)
+
+    if num_elements % tile_size != 0:
+        raise ValueError(
+            f"Number of elements ({num_elements}) must be a multiple of {tile_size}."
+        )
+    num_tiles = num_elements // tile_size
+
+    if input.dtype != output.dtype:
+        raise ValueError(
+            f"Input data types are not the same ({input.dtype} != {output.dtype})."
+        )
+
+    dtype = input.dtype
+
+    # Define tensor types
+    tensor_ty = np.ndarray[(num_elements,), np.dtype[dtype]]
+    tile_ty = np.ndarray[(tile_size,), np.dtype[dtype]]
+
+    # AIE-array data movement with object fifos
+    of_in = ObjectFifo(tile_ty, name="in")
+    of_out = ObjectFifo(tile_ty, name="out")
+
+    # Define a task that will run on a compute tile
+    def core_body(of_in, of_out, func_to_apply):
+        # Extract tile size from ExternalKernel
+        tile_size = func_to_apply.tile_size(0)
+
+        # Number of sub-vector "tile" iterations
+        for i in range_(num_tiles):
+            elem_in = of_in.acquire(1)
+            elem_out = of_out.acquire(1)
+            func_to_apply(elem_in, elem_out, tile_size)
+            of_in.release(1)
+            of_out.release(1)
+
+    # Create a worker to run the task on a compute tile
+    worker = Worker(core_body, fn_args=[of_in.cons(), of_out.prod(), internal_func])
+
+    # Runtime operations to move data to/from the AIE-array
+    rt = Runtime()
+    with rt.sequence(tensor_ty, tensor_ty) as (A, B):
+        rt.start(worker)
+        rt.fill(of_in.prod(), A)
+        rt.drain(of_out.cons(), B, wait=True)
+
+    # Place program components and generate an MLIR module
+    return Program(iron.get_current_device(), rt).resolve_program(SequentialPlacer())
+
+
+@jit(is_placed=False)
+def transform_with_internal_func(input, output):
+    """Transform kernel that creates ExternalKernel internally."""
+    if input.shape != output.shape:
+        raise ValueError(
+            f"Input shapes are not the equal ({input.shape} != {output.shape})."
+        )
+    num_elements = np.size(input)
+
+    # Create ExternalKernel inside the transform
+    internal_func = ExternalKernel(
+        "internal_add_one",
+        source_string="""extern "C" {
+            void internal_add_one(int* input, int* output, int tile_size) {
+                for (int i = 0; i < tile_size; i++) {
+                    output[i] = input[i] + 1;
+                }
+            }
+        }""",
+        arg_types=[
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.ndarray[(16,), np.dtype[np.int32]],
+            np.int32,
+        ],
+    )
+
+    # Extract tile size from ExternalKernel
+    tile_size = internal_func.tile_size(0)
+
+    if num_elements % tile_size != 0:
+        raise ValueError(
+            f"Number of elements ({num_elements}) must be a multiple of {tile_size}."
+        )
+    num_tiles = num_elements // tile_size
+
+    if input.dtype != output.dtype:
+        raise ValueError(
+            f"Input data types are not the same ({input.dtype} != {output.dtype})."
+        )
+
+    dtype = input.dtype
+
+    # Define tensor types
+    tensor_ty = np.ndarray[(num_elements,), np.dtype[dtype]]
+    tile_ty = np.ndarray[(tile_size,), np.dtype[dtype]]
+
+    # AIE-array data movement with object fifos
+    of_in = ObjectFifo(tile_ty, name="in")
+    of_out = ObjectFifo(tile_ty, name="out")
+
+    # Define a task that will run on a compute tile
+    def core_body(of_in, of_out, func_to_apply):
+        # Extract tile size from ExternalKernel
+        tile_size = func_to_apply.tile_size(0)
+
+        # Number of sub-vector "tile" iterations
+        for i in range_(num_tiles):
+            elem_in = of_in.acquire(1)
+            elem_out = of_out.acquire(1)
+            func_to_apply(elem_in, elem_out, tile_size)
+            of_in.release(1)
+            of_out.release(1)
+
+    # Create a worker to run the task on a compute tile
+    worker = Worker(core_body, fn_args=[of_in.cons(), of_out.prod(), internal_func])
+
+    # Runtime operations to move data to/from the AIE-array
+    rt = Runtime()
+    with rt.sequence(tensor_ty, tensor_ty) as (A, B):
+        rt.start(worker)
+        rt.fill(of_in.prod(), A)
+        rt.drain(of_out.cons(), B, wait=True)
+
+    # Place program components and generate an MLIR module
+    return Program(iron.get_current_device(), rt).resolve_program(SequentialPlacer())
+
+
+def test_transform_with_internal_func_with_options_inside():
+    """Test transform function that creates ExternalKernel internally with compiler options."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Apply the transform (ExternalKernel is created inside with hardcoded compiler options)
+    transform_with_internal_func_with_options(input_tensor, output_tensor)
+
+    # Verify results
+    expected = initial_tensor + 1
+    actual = output_tensor.numpy()
+    np.testing.assert_array_equal(actual, expected)
+
+
+def test_transform_with_internal_func_inside():
+    """Test transform function that creates ExternalKernel internally."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Apply the transform (ExternalKernel is created inside)
+    transform_with_internal_func(input_tensor, output_tensor)
+
+    # Verify results
+    expected = initial_tensor + 1
+    actual = output_tensor.numpy()
+    np.testing.assert_array_equal(actual, expected)
+
+
+def test_transform_with_internal_func_from_file():
+    """Test transform function that creates ExternalKernel from a file."""
+    # Create input and output tensors
+    input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
+    output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
+    initial_tensor = input_tensor.numpy().copy()
+
+    # Apply the transform (ExternalKernel is created inside from file)
+    transform_with_internal_func_from_file(input_tensor, output_tensor)
+
+    # Verify results
+    expected = initial_tensor + 42
+    actual = output_tensor.numpy()
+    np.testing.assert_array_equal(actual, expected)

From f5c7fe4b2a30e21f92e73570b80915227e180a86 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 03:25:25 -0700
Subject: [PATCH 03/37] Add tesnor `fill_` test

---
 test/python/tensor.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/test/python/tensor.py b/test/python/tensor.py
index f1c169bd238..e78cb2f35e9 100644
--- a/test/python/tensor.py
+++ b/test/python/tensor.py
@@ -62,6 +62,34 @@ def test_arange_floats():
     assert np.allclose(iron.arange(1.0, 5.0, 1.5), np.arange(1.0, 5.0, 1.5))
 
 
+@pytest.mark.parametrize("dtype", [np.int32, np.float32])
+def test_fill_(dtype):
+    """Test the fill_ method for in-place tensor filling."""
+    t = iron.zeros((2, 3), dtype=dtype, device="npu")
+
+    # Fill with a specific value
+    fill_value = 42 if dtype == np.int32 else 42.5
+    t.fill_(fill_value)
+
+    # Verify the tensor is filled with the correct value
+    expected = np.full((2, 3), fill_value, dtype=dtype)
+    assert np.allclose(t.numpy(), expected)
+
+    # Test with different value
+    new_fill_value = 99 if dtype == np.int32 else 99.9
+    t.fill_(new_fill_value)
+    expected = np.full((2, 3), new_fill_value, dtype=dtype)
+    assert np.allclose(t.numpy(), expected)
+
+
+def test_fill_cpu_tensor():
+    """Test fill_ method on CPU tensors."""
+    t = iron.zeros((2, 2), dtype=np.int32, device="cpu")
+    t.fill_(123)
+    expected = np.full((2, 2), 123, dtype=np.int32)
+    assert np.array_equal(t.numpy(), expected)
+
+
 @pytest.mark.parametrize("dtype", [np.int32, np.float32])
 def test_zeros_like(dtype):
     t = iron.tensor([[1, 2], [3, 4]], dtype=dtype)

From 6d36175e16a03f0518bea3df0826d38cee29c3af Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 03:25:41 -0700
Subject: [PATCH 04/37] Update test name

---
 test/python/tensor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/python/tensor.py b/test/python/tensor.py
index e78cb2f35e9..ddabcc90911 100644
--- a/test/python/tensor.py
+++ b/test/python/tensor.py
@@ -63,7 +63,7 @@ def test_arange_floats():
 
 
 @pytest.mark.parametrize("dtype", [np.int32, np.float32])
-def test_fill_(dtype):
+def test_fill(dtype):
     """Test the fill_ method for in-place tensor filling."""
     t = iron.zeros((2, 3), dtype=dtype, device="npu")
 

From 136a5a7d5c4217d36c39752762542243b0028f82 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 03:32:54 -0700
Subject: [PATCH 05/37] Rename jit test files

---
 test/python/{jit_extern.py => extern_functions.py} | 0
 test/python/{jit.py => jit_compilation.py}         | 0
 2 files changed, 0 insertions(+), 0 deletions(-)
 rename test/python/{jit_extern.py => extern_functions.py} (100%)
 rename test/python/{jit.py => jit_compilation.py} (100%)

diff --git a/test/python/jit_extern.py b/test/python/extern_functions.py
similarity index 100%
rename from test/python/jit_extern.py
rename to test/python/extern_functions.py
diff --git a/test/python/jit.py b/test/python/jit_compilation.py
similarity index 100%
rename from test/python/jit.py
rename to test/python/jit_compilation.py

From e29af7838f3e10508138d935d69cd8ea0ded5fee Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 03:54:43 -0700
Subject: [PATCH 06/37] Fix compiler flags for NPU2

---
 python/iron/jit.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/python/iron/jit.py b/python/iron/jit.py
index 810c6b36bb1..c12e4148f59 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -282,11 +282,21 @@ def compile_external_kernel(func, kernel_dir):
     # 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
+
     # Check device type and use appropriate flags
     if isinstance(current_device, NPU2):
-        cmd.extend(os.environ.get("PEANOWRAP2P_FLAGS", "").split())
+        cmd.extend(common_flags + ["--target=aie2p-none-unknown-elf"])
     elif isinstance(current_device, NPU1):
-        cmd.extend(os.environ.get("PEANOWRAP2_FLAGS", "").split())
+        cmd.extend(common_flags + ["--target=aie2-none-unknown-elf"])
     else:
         raise RuntimeError(f"Unsupported device type: {type(current_device)}")
 

From def58f0c17b45e0c747e3e09fce8b56f54fc34a3 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 04:15:30 -0700
Subject: [PATCH 07/37] Use define flags instead of debug

---
 ...extern_functions.py => jit_extern_functions.py} | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)
 rename test/python/{extern_functions.py => jit_extern_functions.py} (98%)

diff --git a/test/python/extern_functions.py b/test/python/jit_extern_functions.py
similarity index 98%
rename from test/python/extern_functions.py
rename to test/python/jit_extern_functions.py
index 527c573ca8a..4a1eeb23a83 100644
--- a/test/python/extern_functions.py
+++ b/test/python/jit_extern_functions.py
@@ -240,8 +240,8 @@ def test_multiple_defines():
         source_string="""extern "C" {
             void complex_op(int* input, int* output, int tile_size) {
                 for (int i = 0; i < tile_size; i++) {
-                    #ifdef DEBUG
-                    output[i] = input[i] + ADD_VALUE + DEBUG_OFFSET;
+                    #ifdef FLAG2
+                    output[i] = input[i] + ADD_VALUE + FLAG2_OFFSET;
                     #else
                     output[i] = input[i] + ADD_VALUE;
                     #endif
@@ -253,13 +253,13 @@ def test_multiple_defines():
             np.ndarray[(16,), np.dtype[np.int32]],
             np.int32,
         ],
-        compile_flags=["-DADD_VALUE=5", "-DDEBUG", "-DDEBUG_OFFSET=10"],
+        compile_flags=["-DADD_VALUE=5", "-DFLAG2", "-DFLAG2_OFFSET=10"],
     )
 
     # Apply the transform
     transform(input_tensor, output_tensor, complex_op)
 
-    # Verify results (should add 15: 5 + 10 due to DEBUG define)
+    # Verify results (should add 15: 5 + 10 due to FLAG2 define)
     expected = initial_tensor + 15
     actual = output_tensor.numpy()
     np.testing.assert_array_equal(actual, expected)
@@ -793,7 +793,7 @@ def test_invalid_include_directory(invalid_include):
         (["-DADD_VALUE=5"], 5),
         (["-DADD_VALUE=10", "-DMULTIPLIER=2"], 20),  # 10 * 2
         (["-DADD_VALUE=3", "-DOFFSET=7"], 10),  # 3 + 7
-        (["-DADD_VALUE=1", "-DDEBUG", "-DDEBUG_OFFSET=9"], 10),  # 1 + 9 (DEBUG enabled)
+        (["-DADD_VALUE=1", "-DFLAG2", "-DFLAG2_OFFSET=9"], 10),  # 1 + 9 (FLAG2 enabled)
     ],
 )
 def test_compiler_flag_combinations(compile_flags, expected_value):
@@ -809,8 +809,8 @@ def test_compiler_flag_combinations(compile_flags, expected_value):
             for (int i = 0; i < tile_size; i++) {
                 #ifdef MULTIPLIER
                 output[i] = input[i] + ADD_VALUE * MULTIPLIER;
-                #elif defined(DEBUG)
-                output[i] = input[i] + ADD_VALUE + DEBUG_OFFSET;
+                #elif defined(FLAG2)
+                output[i] = input[i] + ADD_VALUE + FLAG2_OFFSET;
                 #elif defined(OFFSET)
                 output[i] = input[i] + ADD_VALUE + OFFSET;
                 #else

From de7e52305a70287972d997f8e7c764ed2d497d5b Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 04:26:35 -0700
Subject: [PATCH 08/37] Remove stray compiler flags

---
 python/iron/jit.py | 32 +++++++++++---------------------
 1 file changed, 11 insertions(+), 21 deletions(-)

diff --git a/python/iron/jit.py b/python/iron/jit.py
index c12e4148f59..358919832e6 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -258,27 +258,6 @@ def compile_external_kernel(func, kernel_dir):
     else:
         raise ValueError("Neither source_string nor source_file is provided")
 
-    # Build compilation command
-    cmd = [
-        f"{os.environ.get('PEANO_INSTALL_DIR', '')}/bin/clang++",
-        "-O2",
-        "-std=c++20",
-        "--target=aie2-none-unknown-elf",
-        "-Wno-parentheses",
-        "-Wno-attributes",
-        "-Wno-macro-redefined",
-        "-Wno-empty-body",
-        "-DNDEBUG",
-    ]
-
-    # Add AIEOPT include directory
-    try:
-        aieopt_path = subprocess.check_output(["which", "aie-opt"], text=True).strip()
-        aieopt_dir = os.path.dirname(os.path.dirname(os.path.realpath(aieopt_path)))
-        cmd.extend(["-I", f"{aieopt_dir}/include"])
-    except subprocess.CalledProcessError:
-        pass
-
     # Add device-specific flags based on actual device detection
     current_device = get_current_device()
 
@@ -292,6 +271,9 @@ def compile_external_kernel(func, kernel_dir):
 
     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"])
@@ -300,6 +282,14 @@ def compile_external_kernel(func, kernel_dir):
     else:
         raise RuntimeError(f"Unsupported device type: {type(current_device)}")
 
+    # Add AIEOPT include directory
+    try:
+        aieopt_path = subprocess.check_output(["which", "aie-opt"], text=True).strip()
+        aieopt_dir = os.path.dirname(os.path.dirname(os.path.realpath(aieopt_path)))
+        cmd.extend(["-I", f"{aieopt_dir}/include"])
+    except subprocess.CalledProcessError:
+        pass
+
     # Add include directories
     for include_dir in func._include_dirs:
         cmd.extend(["-I", include_dir])

From 1c63fe5dfe556057d03297a9e27849037324e453 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 04:56:10 -0700
Subject: [PATCH 09/37] Remove aie opt path addition

---
 python/iron/jit.py | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/python/iron/jit.py b/python/iron/jit.py
index 358919832e6..06389cc0676 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -282,14 +282,6 @@ def compile_external_kernel(func, kernel_dir):
     else:
         raise RuntimeError(f"Unsupported device type: {type(current_device)}")
 
-    # Add AIEOPT include directory
-    try:
-        aieopt_path = subprocess.check_output(["which", "aie-opt"], text=True).strip()
-        aieopt_dir = os.path.dirname(os.path.dirname(os.path.realpath(aieopt_path)))
-        cmd.extend(["-I", f"{aieopt_dir}/include"])
-    except subprocess.CalledProcessError:
-        pass
-
     # Add include directories
     for include_dir in func._include_dirs:
         cmd.extend(["-I", include_dir])

From 4feb272ab54d5a40d8e061ddfb1d1674a7fbc282 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 14:39:35 -0700
Subject: [PATCH 10/37] Use existing `compile_cxx_core_function`

---
 python/iron/compile/compile.py | 14 +++++++--
 python/iron/jit.py             | 57 +++++++++-------------------------
 2 files changed, 27 insertions(+), 44 deletions(-)

diff --git a/python/iron/compile/compile.py b/python/iron/compile/compile.py
index d6c883a6dbb..a08757c493f 100644
--- a/python/iron/compile/compile.py
+++ b/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(
diff --git a/python/iron/jit.py b/python/iron/jit.py
index 06389cc0676..08636db6d9a 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -17,7 +17,6 @@
 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.
@@ -261,54 +260,28 @@ def compile_external_kernel(func, kernel_dir):
     # 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
+    # Determine target architecture based on device type
     if isinstance(current_device, NPU2):
-        cmd.extend(common_flags + ["--target=aie2p-none-unknown-elf"])
+        target_arch = "aie2p"
     elif isinstance(current_device, NPU1):
-        cmd.extend(common_flags + ["--target=aie2-none-unknown-elf"])
+        target_arch = "aie2"
     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)
+    from .compile.compile import compile_cxx_core_function
 
-    # Add source and output files
-    cmd.extend(["-c", source_file, "-o", output_file])
+    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
+    )
 
-    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)
+    # Mark the function as compiled
+    func._compiled = True
 
 
 def hash_module(module, external_kernels=None):

From 8926f75a6af79b729b359173cdb6006e0691c864 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <112003944+mawad-amd@users.noreply.github.com>
Date: Tue, 5 Aug 2025 14:42:51 -0700
Subject: [PATCH 11/37] Update python/iron/jit.py

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
---
 python/iron/jit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/iron/jit.py b/python/iron/jit.py
index 08636db6d9a..1c6ffb86497 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -277,7 +277,7 @@ def compile_external_kernel(func, kernel_dir):
         include_dirs=func._include_dirs,
         compile_args=func._compile_flags,
         cwd=kernel_dir,
-        verbose=False
+        verbose=False,
     )
 
     # Mark the function as compiled

From 37f61d3e52fe584287e0103ecfb241e9cbeabe97 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 14:52:26 -0700
Subject: [PATCH 12/37] Move import to the top

---
 python/iron/jit.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/python/iron/jit.py b/python/iron/jit.py
index 1c6ffb86497..79215c79920 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -11,6 +11,7 @@
 import hashlib
 import numpy as np
 import pyxrt as xrt
+import shutil
 
 from aie.extras.context import mlir_mod_ctx
 from ..utils.compile import compile_mlir_module_to_binary
@@ -207,8 +208,6 @@ def decorator(*args, **kwargs):
             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
 
@@ -247,7 +246,6 @@ def compile_external_kernel(func, kernel_dir):
             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):

From 433aa64e334402313e3702681d2395c8acb4d99b Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:04:14 -0700
Subject: [PATCH 13/37] Try debugging the workflow. Can't repro locally.

---
 .github/workflows/buildAndTestRyzenAI.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index d97c2906480..2476fbb50a5 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -22,7 +22,7 @@ on:
         type: boolean
         description: 'Run the build with tmate debugging enabled (https://github.com/marketplace/actions/debugging-with-tmate)'
         required: false
-        default: false
+        default: true
 
 defaults:
   run:

From 7877e21801fa62758dcba2d8eec9a276b6368793 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:08:47 -0700
Subject: [PATCH 14/37] Always trigger debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index 2476fbb50a5..a3ae1921771 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -72,8 +72,7 @@ jobs:
         uses: mxschmitt/action-tmate@v3
         # To run this, launch it manually on the default branch and
         # click on "launch_tmate_terminal_for_debug"
-        if: github.event_name == 'workflow_dispatch'
-                && inputs.launch_tmate_terminal_for_debug
+        if: inputs.launch_tmate_terminal_for_debug
 
       - name: Run commands
         run: |
@@ -150,7 +149,6 @@ jobs:
         # To run this, launch it manually on the default branch and
         # click on "launch_tmate_terminal_for_debug"
         if: github.event_name == 'workflow_dispatch'
-                && inputs.launch_tmate_terminal_for_debug
 
       - name: Run commands
         run: |

From f32061f1a69985319073fbb3fc62cbc7c6dcf9c0 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:09:04 -0700
Subject: [PATCH 15/37] always trigger debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index a3ae1921771..c7f70428bf7 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -148,7 +148,7 @@ jobs:
         uses: mxschmitt/action-tmate@v3
         # To run this, launch it manually on the default branch and
         # click on "launch_tmate_terminal_for_debug"
-        if: github.event_name == 'workflow_dispatch'
+        if: inputs.launch_tmate_terminal_for_debug
 
       - name: Run commands
         run: |

From e31859cb387825ab8f829be272e80c160f4069e2 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:10:17 -0700
Subject: [PATCH 16/37] Run only broken test

---
 .github/workflows/buildAndTestRyzenAI.yml | 116 +++++++++++-----------
 1 file changed, 58 insertions(+), 58 deletions(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index c7f70428bf7..14b41310fcf 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -59,7 +59,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        runner_type: [ amd7940hs, amdhx370 ]
+        runner_type: [ amdhx370 ]
     steps:
       - uses: actions/checkout@v4
         with:
@@ -129,60 +129,60 @@ jobs:
           ninja check-aie
           popd
 
-  build-quick-setup:
-    name: Run Examples on Ryzen AI
-    runs-on: ${{ matrix.runner_type }}
-    strategy:
-      fail-fast: false
-      matrix:
-        runner_type: [ amd7940hs, amdhx370 ]
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          submodules: "true"
-
-      # Launch an ssh session via a proxy server if there is a need
-      # for debug. This seems to live for 35 min max
-      # https://github.com/mxschmitt/action-tmate
-      - name: Setup tmate session
-        uses: mxschmitt/action-tmate@v3
-        # To run this, launch it manually on the default branch and
-        # click on "launch_tmate_terminal_for_debug"
-        if: inputs.launch_tmate_terminal_for_debug
-
-      - name: Run commands
-        run: |
-          sudo prlimit -lunlimited --pid $$
-          pip cache purge
-          source /opt/xilinx/xrt/setup.sh
-          export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
-          source utils/quick_setup.sh
-          # quick_setup changes directory to programming_examples, so we need to return to mlir-aie
-          cd ..
-          pip install -r python/requirements_test.txt
-
-          ./utils/build-mlir-aie-from-wheels.sh
-
-          # I have no clue why but the system clock on GHA containers is like 12 hours ahead.
-          # That means wheels have file with time stamps in the future which makes ninja loop
-          # forever when configuring. Set the time to some arbitrary stamp in the past just to be safe.
-          pushd my_install
-          find mlir -exec touch -a -m -t 201108231405.14 {} \;
-          popd
-
-          # build is created by the build-mlir-aie-from-wheels.sh script
-          pushd build
-
-          # -j here to reduce the number of parallel chess jobs.
-          # -j4 for 32GB RAM, -j12 for 64GB RAM
-          if [ x"${{ matrix.runner_type }}" == x"amdhx370" ]; then
-            LIT_OPTS="-j4 $LIT_OPTS"
-          else
-            LIT_OPTS="-j12 $LIT_OPTS"
-          fi
-
-          ninja install
-          ninja check-reference-designs
-          ninja check-programming-guide
-
-          popd
+  # build-quick-setup:
+  #   name: Run Examples on Ryzen AI
+  #   runs-on: ${{ matrix.runner_type }}
+  #   strategy:
+  #     fail-fast: false
+  #     matrix:
+  #       runner_type: [ amd7940hs, amdhx370 ]
+  #   steps:
+  #     - uses: actions/checkout@v4
+  #       with:
+  #         submodules: "true"
+
+  #     # Launch an ssh session via a proxy server if there is a need
+  #     for debug. This seems to live for 35 min max
+  #     # https://github.com/mxschmitt/action-tmate
+  #     - name: Setup tmate session
+  #       uses: mxschmitt/action-tmate@v3
+  #       # To run this, launch it manually on the default branch and
+  #       # click on "launch_tmate_terminal_for_debug"
+  #       if: inputs.launch_tmate_terminal_for_debug
+
+  #     - name: Run commands
+  #       run: |
+  #         sudo prlimit -lunlimited --pid $$
+  #         pip cache purge
+  #         source /opt/xilinx/xrt/setup.sh
+  #         export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
+  #         source utils/quick_setup.sh
+  #         # quick_setup changes directory to programming_examples, so we need to return to mlir-aie
+  #         cd ..
+  #         pip install -r python/requirements_test.txt
+
+  #         ./utils/build-mlir-aie-from-wheels.sh
+
+  #         # I have no clue why but the system clock on GHA containers is like 12 hours ahead.
+  #         # That means wheels have file with time stamps in the future which makes ninja loop
+  #         forever when configuring. Set the time to some arbitrary stamp in the past just to be safe.
+  #         pushd my_install
+  #         find mlir -exec touch -a -m -t 201108231405.14 {} \;
+  #         popd
+
+  #         # build is created by the build-mlir-aie-from-wheels.sh script
+  #         pushd build
+
+  #         # -j here to reduce the number of parallel chess jobs.
+  #         # -j4 for 32GB RAM, -j12 for 64GB RAM
+  #         if [ x"${{ matrix.runner_type }}" == x"amdhx370" ]; then
+  #           LIT_OPTS="-j4 $LIT_OPTS"
+  #           else
+  #           LIT_OPTS="-j12 $LIT_OPTS"
+  #         fi
+
+  #         ninja install
+  #         ninja check-reference-designs
+  #         ninja check-programming-guide
+
+  #         popd

From 3db23a95d3afd4e8bee0ed3c613ce766d1aa5ba5 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:13:25 -0700
Subject: [PATCH 17/37] Remove if

---
 .github/workflows/buildAndTestRyzenAI.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index 14b41310fcf..bf1aa7b59c9 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -72,7 +72,7 @@ jobs:
         uses: mxschmitt/action-tmate@v3
         # To run this, launch it manually on the default branch and
         # click on "launch_tmate_terminal_for_debug"
-        if: inputs.launch_tmate_terminal_for_debug
+        #if: inputs.launch_tmate_terminal_for_debug
 
       - name: Run commands
         run: |

From 00193fdeff316e5377e79592bfca6f846f98c975 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:15:06 -0700
Subject: [PATCH 18/37] Remove debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index bf1aa7b59c9..14b41310fcf 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -72,7 +72,7 @@ jobs:
         uses: mxschmitt/action-tmate@v3
         # To run this, launch it manually on the default branch and
         # click on "launch_tmate_terminal_for_debug"
-        #if: inputs.launch_tmate_terminal_for_debug
+        if: inputs.launch_tmate_terminal_for_debug
 
       - name: Run commands
         run: |

From 4b7571411ab7104f78395aafb21f456097d8ea24 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:20:59 -0700
Subject: [PATCH 19/37] Manual debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index 14b41310fcf..dc24a4dbe32 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -126,7 +126,16 @@ jobs:
             $CMAKE_ARGS
 
           ninja install
-          ninja check-aie
+          
+          cd ..
+          xrt-smi examine
+
+          python programming_examples/basic/vector_vector_add/vector_vector_add.py
+          pytest -vvv -s test/python/jit_compilation.py
+          pytest -vvv -s test/python/jit_extern_functions.py
+          pytest -vvv -s test/python/jit_extern_inside_jit.py
+    
+          #ninja check-aie
           popd
 
   # build-quick-setup:

From a1be75f12798e3ff3b35a8044f72a3fd87b51666 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:25:52 -0700
Subject: [PATCH 20/37] Debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index dc24a4dbe32..ac0cae8c9d8 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -128,6 +128,11 @@ jobs:
           ninja install
           
           cd ..
+          ls -alth
+          ls -alth ~/.iron
+          source /opt/xilinx/xrt/setup.sh
+          source utils/env_setup.sh
+          export PYTHONPATH="$(pwd)/install/python:${PYTHONPATH}"
           xrt-smi examine
 
           python programming_examples/basic/vector_vector_add/vector_vector_add.py

From 40ba1928a976e0e60b3e0e31fac5eae7e88dfd73 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:29:44 -0700
Subject: [PATCH 21/37] Pass device

---
 .github/workflows/buildAndTestRyzenAI.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index ac0cae8c9d8..6bb7e8f0b00 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -135,7 +135,7 @@ jobs:
           export PYTHONPATH="$(pwd)/install/python:${PYTHONPATH}"
           xrt-smi examine
 
-          python programming_examples/basic/vector_vector_add/vector_vector_add.py
+          python programming_examples/basic/vector_vector_add/vector_vector_add.py -d npu2
           pytest -vvv -s test/python/jit_compilation.py
           pytest -vvv -s test/python/jit_extern_functions.py
           pytest -vvv -s test/python/jit_extern_inside_jit.py

From 1611effaa45e0f9a50b13220c6dcf1c164f632a8 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:39:47 -0700
Subject: [PATCH 22/37] Debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index 6bb7e8f0b00..4bf4824e0cb 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -59,7 +59,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        runner_type: [ amdhx370 ]
+        runner_type: [ amdhx370 , amd7940hs ]
     steps:
       - uses: actions/checkout@v4
         with:
@@ -134,7 +134,12 @@ jobs:
           source utils/env_setup.sh
           export PYTHONPATH="$(pwd)/install/python:${PYTHONPATH}"
           xrt-smi examine
-
+          python -c "from aie.iron import get_current_device; from aie.iron.device import NPU1, NPU2; current_device = get_current_device(); target_arch = 'aie2p' if isinstance(current_device, NPU2) else 'aie2' if isinstance(current_device, NPU1) else None; print(f'Device: {current_device}, Target arch: {target_arch}')"
+          
+          # Debug: Check what's available in aie.iron
+          python -c "import aie.iron; print(dir(aie.iron))"
+          cat /home/github/actions-runner/_work/mlir-aie/mlir-aie/aie-venv/lib/python3.12/site-packages/mlir_aie/python/aie/iron/kernel.py
+          
           python programming_examples/basic/vector_vector_add/vector_vector_add.py -d npu2
           pytest -vvv -s test/python/jit_compilation.py
           pytest -vvv -s test/python/jit_extern_functions.py

From ada8709c44a824aa4df6d6881d019d8da2615333 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:49:41 -0700
Subject: [PATCH 23/37] Fix install command

---
 .github/workflows/buildAndTestRyzenAI.yml | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index 4bf4824e0cb..f74e2f3ac9f 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -123,6 +123,13 @@ jobs:
             -DCMAKE_INSTALL_PREFIX=$PWD/../mlir_aie \
             -DCMAKE_MODULE_PATH=$PWD/../cmake/modulesXilinx \
             -DMLIR_DIR=$PWD/../mlir/lib/cmake/mlir \
+            -DAIE_ENABLE_BINDINGS_PYTHON=ON \
+            -DLLVM_ENABLE_RTTI=ON \
+            -DLLVM_ENABLE_ASSERTIONS=ON \
+            -DAIE_VITIS_COMPONENTS=AIE2;AIE2P \
+            -DAIE_RUNTIME_TARGETS=x86_64 \
+            -DAIE_RUNTIME_TEST_TARGET=x86_64 \
+            -DVITIS_VPP=$(which v++) \
             $CMAKE_ARGS
 
           ninja install

From ebd4eed8ae7a92a85d43ecb9eb9afdaa474aedbe Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 22:54:59 -0700
Subject: [PATCH 24/37] Fix cmake command

---
 .github/workflows/buildAndTestRyzenAI.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index f74e2f3ac9f..c41fedc1ce4 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -46,7 +46,7 @@ env:
             -DCMAKE_MODULE_LINKER_FLAGS_INIT="-fuse-ld=lld" \
             -DCMAKE_SHARED_LINKER_FLAGS_INIT="-fuse-ld=lld" \
             -DXRT_ROOT=/opt/xilinx/xrt \
-            -DAIE_VITIS_COMPONENTS=AIE2;AIE2P \
+            -DAIE_VITIS_COMPONENTS="AIE2;AIE2P" \
             -DAIE_ENABLE_PYTHON_PASSES=OFF \
             -DAIE_ENABLE_XRT_PYTHON_BINDINGS=ON \
             -DAIE_INCLUDE_INTEGRATION_TESTS=OFF
@@ -126,7 +126,7 @@ jobs:
             -DAIE_ENABLE_BINDINGS_PYTHON=ON \
             -DLLVM_ENABLE_RTTI=ON \
             -DLLVM_ENABLE_ASSERTIONS=ON \
-            -DAIE_VITIS_COMPONENTS=AIE2;AIE2P \
+            -DAIE_VITIS_COMPONENTS="AIE2;AIE2P" \
             -DAIE_RUNTIME_TARGETS=x86_64 \
             -DAIE_RUNTIME_TEST_TARGET=x86_64 \
             -DVITIS_VPP=$(which v++) \

From 7eb8f3b92e4beb5160c74998303bf6f9ecff2eb0 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 23:09:07 -0700
Subject: [PATCH 25/37] Debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 227 +++++++++++-----------
 1 file changed, 118 insertions(+), 109 deletions(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index c41fedc1ce4..f8eadba528f 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -53,13 +53,115 @@ env:
   LIT_OPTS: -sv --time-tests --timeout 600 --show-unsupported --show-excluded
 
 jobs:
-  build-tests:
-    name: Run Tests on Ryzen AI
+#  build-tests:
+#    name: Run Tests on Ryzen AI
+#    runs-on: ${{ matrix.runner_type }}
+#    strategy:
+#      fail-fast: false
+#      matrix:
+#        runner_type: [ amdhx370 , amd7940hs ]
+#    steps:
+#      - uses: actions/checkout@v4
+#        with:
+#          submodules: "true"
+#
+#      # Launch an ssh session via a proxy server if there is a need
+#      # for debug. This seems to live for 35 min max
+#      # https://github.com/mxschmitt/action-tmate
+#      - name: Setup tmate session
+#        uses: mxschmitt/action-tmate@v3
+#        # To run this, launch it manually on the default branch and
+#        # click on "launch_tmate_terminal_for_debug"
+#        if: inputs.launch_tmate_terminal_for_debug
+#
+#      - name: Run commands
+#        run: |
+#          sudo prlimit -lunlimited --pid $$
+#          pip cache purge
+#          source /opt/xilinx/xrt/setup.sh
+#          python -m venv aie-venv
+#          source aie-venv/bin/activate
+#          pip install -r python/requirements.txt
+#          HOST_MLIR_PYTHON_PACKAGE_PREFIX=aie pip install -r python/requirements_extras.txt
+#          pip install -r python/requirements_notebook.txt
+#          pip install -r python/requirements_ml.txt
+#          pip install -r python/requirements_test.txt
+#
+#          # Install llvm-aie (Peano) which is needed for unittests requiring compiling NPU code
+#          pip install -I llvm-aie -f https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly
+#          export PEANO_INSTALL_DIR="$(pip show llvm-aie | grep ^Location: | awk '{print $2}')/llvm-aie"
+#          
+#          sed -i.bak 's/OUTPUT_TIMEOUT = 10/OUTPUT_TIMEOUT = 100/g' \
+#            $(python -c 'import site; print(site.getsitepackages()[0])')/jupyter_client/runapp.py
+#
+#          VERSION=$(utils/clone-llvm.sh --get-wheel-version)
+#          pip -q download mlir==$VERSION \
+#            -f https://github.com/Xilinx/mlir-aie/releases/expanded_assets/mlir-distro
+#          unzip -q mlir-*.whl
+#          # I have no clue why but the system clock on GHA containers is like 12 hours ahead.
+#          # That means wheels have file with time stamps in the future which makes ninja loop
+#          # forever when configuring. Set the time to some arbitrary stamp in the past just to be safe.
+#          find mlir -exec touch -a -m -t 201108231405.14 {} \;
+#
+#          mkdir build
+#          pushd build
+#
+#          # -j here to reduce the number of parallel chess jobs.
+#          # -j4 for 32GB RAM, -j12 for 64GB RAM
+#          if [ x"${{ matrix.runner_type }}" == x"amdhx370" ]; then
+#            LIT_OPTS="-j4 $LIT_OPTS"
+#          else
+#            LIT_OPTS="-j12 $LIT_OPTS"
+#          fi
+#
+#          export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
+#          
+#          cmake .. -G Ninja \
+#            -DCMAKE_BUILD_TYPE=Release \
+#            -DPython3_EXECUTABLE=$(which python) \
+#            -DLLVM_EXTERNAL_LIT=$(which lit) \
+#            -DCMAKE_INSTALL_PREFIX=$PWD/../mlir_aie \
+#            -DCMAKE_MODULE_PATH=$PWD/../cmake/modulesXilinx \
+#            -DMLIR_DIR=$PWD/../mlir/lib/cmake/mlir \
+#            -DAIE_ENABLE_BINDINGS_PYTHON=ON \
+#            -DLLVM_ENABLE_RTTI=ON \
+#            -DLLVM_ENABLE_ASSERTIONS=ON \
+#            -DAIE_VITIS_COMPONENTS="AIE2;AIE2P" \
+#            -DAIE_RUNTIME_TARGETS=x86_64 \
+#            -DAIE_RUNTIME_TEST_TARGET=x86_64 \
+#            -DVITIS_VPP=$(which v++) \
+#            $CMAKE_ARGS
+#
+#          ninja install
+#          
+#          cd ..
+#          ls -alth
+#          ls -alth ~/.iron
+#          source /opt/xilinx/xrt/setup.sh
+#          source utils/env_setup.sh
+#          export PYTHONPATH="$(pwd)/install/python:${PYTHONPATH}"
+#          xrt-smi examine
+#          python -c "from aie.iron import get_current_device; from aie.iron.device import NPU1, NPU2; current_device = get_current_device(); target_arch = 'aie2p' if isinstance(current_device, NPU2) else 'aie2' if isinstance(current_device, NPU1) else None; print(f'Device: {current_device}, Target arch: {target_arch}')"
+#          
+#          # Debug: Check what's available in aie.iron
+#          python -c "import aie.iron; print(dir(aie.iron))"
+#          cat /home/github/actions-runner/_work/mlir-aie/mlir-aie/aie-venv/lib/python3.12/site-packages/mlir_aie/python/aie/iron/kernel.py
+#          
+#          python programming_examples/basic/vector_vector_add/vector_vector_add.py -d npu2
+#          pytest -vvv -s test/python/jit_compilation.py
+#          pytest -vvv -s test/python/jit_extern_functions.py
+#          pytest -vvv -s test/python/jit_extern_inside_jit.py
+#    
+#          #ninja check-aie
+#          popd
+#
+  build-quick-setup:
+    name: Run Examples on Ryzen AI
     runs-on: ${{ matrix.runner_type }}
     strategy:
       fail-fast: false
       matrix:
-        runner_type: [ amdhx370 , amd7940hs ]
+        runner_type: [ amd7940hs, amdhx370 ]
     steps:
       - uses: actions/checkout@v4
         with:
@@ -79,69 +181,31 @@ jobs:
           sudo prlimit -lunlimited --pid $$
           pip cache purge
           source /opt/xilinx/xrt/setup.sh
-          python -m venv aie-venv
-          source aie-venv/bin/activate
-          pip install -r python/requirements.txt
-          HOST_MLIR_PYTHON_PACKAGE_PREFIX=aie pip install -r python/requirements_extras.txt
-          pip install -r python/requirements_notebook.txt
-          pip install -r python/requirements_ml.txt
+          export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
+          source utils/quick_setup.sh
+          # quick_setup changes directory to programming_examples, so we need to return to mlir-aie
+          cd ..
           pip install -r python/requirements_test.txt
 
-          # Install llvm-aie (Peano) which is needed for unittests requiring compiling NPU code
-          pip install -I llvm-aie -f https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly
-          export PEANO_INSTALL_DIR="$(pip show llvm-aie | grep ^Location: | awk '{print $2}')/llvm-aie"
-          
-          sed -i.bak 's/OUTPUT_TIMEOUT = 10/OUTPUT_TIMEOUT = 100/g' \
-            $(python -c 'import site; print(site.getsitepackages()[0])')/jupyter_client/runapp.py
+          ./utils/build-mlir-aie-from-wheels.sh
 
-          VERSION=$(utils/clone-llvm.sh --get-wheel-version)
-          pip -q download mlir==$VERSION \
-            -f https://github.com/Xilinx/mlir-aie/releases/expanded_assets/mlir-distro
-          unzip -q mlir-*.whl
           # I have no clue why but the system clock on GHA containers is like 12 hours ahead.
           # That means wheels have file with time stamps in the future which makes ninja loop
           # forever when configuring. Set the time to some arbitrary stamp in the past just to be safe.
+          pushd my_install
           find mlir -exec touch -a -m -t 201108231405.14 {} \;
+          popd
 
-          mkdir build
+          # build is created by the build-mlir-aie-from-wheels.sh script
           pushd build
 
           # -j here to reduce the number of parallel chess jobs.
           # -j4 for 32GB RAM, -j12 for 64GB RAM
           if [ x"${{ matrix.runner_type }}" == x"amdhx370" ]; then
             LIT_OPTS="-j4 $LIT_OPTS"
-          else
+            else
             LIT_OPTS="-j12 $LIT_OPTS"
           fi
-
-          export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
-          
-          cmake .. -G Ninja \
-            -DCMAKE_BUILD_TYPE=Release \
-            -DPython3_EXECUTABLE=$(which python) \
-            -DLLVM_EXTERNAL_LIT=$(which lit) \
-            -DCMAKE_INSTALL_PREFIX=$PWD/../mlir_aie \
-            -DCMAKE_MODULE_PATH=$PWD/../cmake/modulesXilinx \
-            -DMLIR_DIR=$PWD/../mlir/lib/cmake/mlir \
-            -DAIE_ENABLE_BINDINGS_PYTHON=ON \
-            -DLLVM_ENABLE_RTTI=ON \
-            -DLLVM_ENABLE_ASSERTIONS=ON \
-            -DAIE_VITIS_COMPONENTS="AIE2;AIE2P" \
-            -DAIE_RUNTIME_TARGETS=x86_64 \
-            -DAIE_RUNTIME_TEST_TARGET=x86_64 \
-            -DVITIS_VPP=$(which v++) \
-            $CMAKE_ARGS
-
-          ninja install
-          
-          cd ..
-          ls -alth
-          ls -alth ~/.iron
-          source /opt/xilinx/xrt/setup.sh
-          source utils/env_setup.sh
-          export PYTHONPATH="$(pwd)/install/python:${PYTHONPATH}"
-          xrt-smi examine
-          python -c "from aie.iron import get_current_device; from aie.iron.device import NPU1, NPU2; current_device = get_current_device(); target_arch = 'aie2p' if isinstance(current_device, NPU2) else 'aie2' if isinstance(current_device, NPU1) else None; print(f'Device: {current_device}, Target arch: {target_arch}')"
           
           # Debug: Check what's available in aie.iron
           python -c "import aie.iron; print(dir(aie.iron))"
@@ -151,64 +215,9 @@ jobs:
           pytest -vvv -s test/python/jit_compilation.py
           pytest -vvv -s test/python/jit_extern_functions.py
           pytest -vvv -s test/python/jit_extern_inside_jit.py
-    
-          #ninja check-aie
-          popd
-
-  # build-quick-setup:
-  #   name: Run Examples on Ryzen AI
-  #   runs-on: ${{ matrix.runner_type }}
-  #   strategy:
-  #     fail-fast: false
-  #     matrix:
-  #       runner_type: [ amd7940hs, amdhx370 ]
-  #   steps:
-  #     - uses: actions/checkout@v4
-  #       with:
-  #         submodules: "true"
-
-  #     # Launch an ssh session via a proxy server if there is a need
-  #     for debug. This seems to live for 35 min max
-  #     # https://github.com/mxschmitt/action-tmate
-  #     - name: Setup tmate session
-  #       uses: mxschmitt/action-tmate@v3
-  #       # To run this, launch it manually on the default branch and
-  #       # click on "launch_tmate_terminal_for_debug"
-  #       if: inputs.launch_tmate_terminal_for_debug
-
-  #     - name: Run commands
-  #       run: |
-  #         sudo prlimit -lunlimited --pid $$
-  #         pip cache purge
-  #         source /opt/xilinx/xrt/setup.sh
-  #         export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
-  #         source utils/quick_setup.sh
-  #         # quick_setup changes directory to programming_examples, so we need to return to mlir-aie
-  #         cd ..
-  #         pip install -r python/requirements_test.txt
 
-  #         ./utils/build-mlir-aie-from-wheels.sh
+          #ninja install
+          #ninja check-reference-designs
+          #ninja check-programming-guide
 
-  #         # I have no clue why but the system clock on GHA containers is like 12 hours ahead.
-  #         # That means wheels have file with time stamps in the future which makes ninja loop
-  #         forever when configuring. Set the time to some arbitrary stamp in the past just to be safe.
-  #         pushd my_install
-  #         find mlir -exec touch -a -m -t 201108231405.14 {} \;
-  #         popd
-
-  #         # build is created by the build-mlir-aie-from-wheels.sh script
-  #         pushd build
-
-  #         # -j here to reduce the number of parallel chess jobs.
-  #         # -j4 for 32GB RAM, -j12 for 64GB RAM
-  #         if [ x"${{ matrix.runner_type }}" == x"amdhx370" ]; then
-  #           LIT_OPTS="-j4 $LIT_OPTS"
-  #           else
-  #           LIT_OPTS="-j12 $LIT_OPTS"
-  #         fi
-
-  #         ninja install
-  #         ninja check-reference-designs
-  #         ninja check-programming-guide
-
-  #         popd
+          popd

From 7e861d51ac79eef60bb77eeab0b52793afb4dbfe Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 23:16:17 -0700
Subject: [PATCH 26/37] Debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index f8eadba528f..b54719dab3b 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -209,7 +209,15 @@ jobs:
           
           # Debug: Check what's available in aie.iron
           python -c "import aie.iron; print(dir(aie.iron))"
-          cat /home/github/actions-runner/_work/mlir-aie/mlir-aie/aie-venv/lib/python3.12/site-packages/mlir_aie/python/aie/iron/kernel.py
+          
+          # Debug: Check Python path and environment
+          echo "Python executable: $(which python)"
+          echo "Python version: $(python --version)"
+          echo "PYTHONPATH: $PYTHONPATH"
+          python -c "import sys; print('Python path:'); [print(p) for p in sys.path]"
+          python -c "import aie.iron; print('aie.iron location:', aie.iron.__file__)"
+
+          cat python/iron/kernel.py 
           
           python programming_examples/basic/vector_vector_add/vector_vector_add.py -d npu2
           pytest -vvv -s test/python/jit_compilation.py

From 44155c6db44f63eb0a3d0b93bad5d16183a771c5 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 23:21:53 -0700
Subject: [PATCH 27/37] Debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index b54719dab3b..c484a0a08ab 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -206,6 +206,11 @@ jobs:
             else
             LIT_OPTS="-j12 $LIT_OPTS"
           fi
+
+          ls -alth
+          cd ..
+          ls -alth
+          git log -1 | cat
           
           # Debug: Check what's available in aie.iron
           python -c "import aie.iron; print(dir(aie.iron))"
@@ -216,7 +221,6 @@ jobs:
           echo "PYTHONPATH: $PYTHONPATH"
           python -c "import sys; print('Python path:'); [print(p) for p in sys.path]"
           python -c "import aie.iron; print('aie.iron location:', aie.iron.__file__)"
-
           cat python/iron/kernel.py 
           
           python programming_examples/basic/vector_vector_add/vector_vector_add.py -d npu2

From ec2a777386a3afeeb0658ac6bf475ec979ca96a4 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 23:35:31 -0700
Subject: [PATCH 28/37] debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 262 +++++++++++-----------
 1 file changed, 133 insertions(+), 129 deletions(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index c484a0a08ab..aebfdcb25c6 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -53,13 +53,118 @@ env:
   LIT_OPTS: -sv --time-tests --timeout 600 --show-unsupported --show-excluded
 
 jobs:
-#  build-tests:
-#    name: Run Tests on Ryzen AI
+  build-tests:
+    name: Run Tests on Ryzen AI
+    runs-on: ${{ matrix.runner_type }}
+    strategy:
+      fail-fast: false
+      matrix:
+        runner_type: [ amdhx370 , amd7940hs ]
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          submodules: "true"
+
+      # Launch an ssh session via a proxy server if there is a need
+      # for debug. This seems to live for 35 min max
+      # https://github.com/mxschmitt/action-tmate
+      - name: Setup tmate session
+        uses: mxschmitt/action-tmate@v3
+        # To run this, launch it manually on the default branch and
+        # click on "launch_tmate_terminal_for_debug"
+        if: inputs.launch_tmate_terminal_for_debug
+
+      - name: Run commands
+        run: |
+          sudo prlimit -lunlimited --pid $$
+          pip cache purge
+          source /opt/xilinx/xrt/setup.sh
+          python -m venv aie-venv
+          source aie-venv/bin/activate
+          pip install -r python/requirements.txt
+          HOST_MLIR_PYTHON_PACKAGE_PREFIX=aie pip install -r python/requirements_extras.txt
+          pip install -r python/requirements_notebook.txt
+          pip install -r python/requirements_ml.txt
+          pip install -r python/requirements_test.txt
+
+          # Install llvm-aie (Peano) which is needed for unittests requiring compiling NPU code
+          pip install -I llvm-aie -f https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly
+          export PEANO_INSTALL_DIR="$(pip show llvm-aie | grep ^Location: | awk '{print $2}')/llvm-aie"
+          
+          sed -i.bak 's/OUTPUT_TIMEOUT = 10/OUTPUT_TIMEOUT = 100/g' \
+            $(python -c 'import site; print(site.getsitepackages()[0])')/jupyter_client/runapp.py
+
+          VERSION=$(utils/clone-llvm.sh --get-wheel-version)
+          pip -q download mlir==$VERSION \
+            -f https://github.com/Xilinx/mlir-aie/releases/expanded_assets/mlir-distro
+          unzip -q mlir-*.whl
+          # I have no clue why but the system clock on GHA containers is like 12 hours ahead.
+          # That means wheels have file with time stamps in the future which makes ninja loop
+          # forever when configuring. Set the time to some arbitrary stamp in the past just to be safe.
+          find mlir -exec touch -a -m -t 201108231405.14 {} \;
+
+          mkdir build
+          pushd build
+
+          # -j here to reduce the number of parallel chess jobs.
+          # -j4 for 32GB RAM, -j12 for 64GB RAM
+          if [ x"${{ matrix.runner_type }}" == x"amdhx370" ]; then
+            LIT_OPTS="-j4 $LIT_OPTS"
+          else
+            LIT_OPTS="-j12 $LIT_OPTS"
+          fi
+
+          export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
+          
+          cmake .. -G Ninja \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DPython3_EXECUTABLE=$(which python) \
+            -DLLVM_EXTERNAL_LIT=$(which lit) \
+            -DCMAKE_INSTALL_PREFIX=$PWD/../mlir_aie \
+            -DCMAKE_MODULE_PATH=$PWD/../cmake/modulesXilinx \
+            -DMLIR_DIR=$PWD/../mlir/lib/cmake/mlir \
+            -DAIE_ENABLE_BINDINGS_PYTHON=ON \
+            -DLLVM_ENABLE_RTTI=ON \
+            -DLLVM_ENABLE_ASSERTIONS=ON \
+            -DAIE_VITIS_COMPONENTS="AIE2;AIE2P" \
+            -DAIE_RUNTIME_TARGETS=x86_64 \
+            -DAIE_RUNTIME_TEST_TARGET=x86_64 \
+            -DVITIS_VPP=$(which v++) \
+            $CMAKE_ARGS
+
+          ninja install
+          
+          cd ..
+          ls -alth
+          ls -alth ~/.iron
+          source /opt/xilinx/xrt/setup.sh
+          source utils/env_setup.sh
+          export PYTHONPATH="$(pwd)/install/python:${PYTHONPATH}"
+          xrt-smi examine
+          python -c "from aie.iron import get_current_device; from aie.iron.device import NPU1, NPU2; current_device = get_current_device(); target_arch = 'aie2p' if isinstance(current_device, NPU2) else 'aie2' if isinstance(current_device, NPU1) else None; print(f'Device: {current_device}, Target arch: {target_arch}')"
+          
+          # Debug: Check what's available in aie.iron
+          python -c "import aie.iron; print(dir(aie.iron))"
+          cat /home/github/actions-runner/_work/mlir-aie/mlir-aie/aie-venv/lib/python3.12/site-packages/mlir_aie/python/aie/iron/kernel.py
+          ls -alt mlir_aie
+          cat mlir_aie/python/aie/iron/__init__.py
+          cat mlir_aie/python/aie/iron/kernel.py
+          export PYTHONPATH="$(pwd)/mlir_aie/python:${PYTHONPATH}"
+          python programming_examples/basic/vector_vector_add/vector_vector_add.py -d npu2
+          pytest -vvv -s test/python/jit_compilation.py
+          pytest -vvv -s test/python/jit_extern_functions.py
+          pytest -vvv -s test/python/jit_extern_inside_jit.py
+    
+          #ninja check-aie
+          popd
+
+#  build-quick-setup:
+#    name: Run Examples on Ryzen AI
 #    runs-on: ${{ matrix.runner_type }}
 #    strategy:
 #      fail-fast: false
 #      matrix:
-#        runner_type: [ amdhx370 , amd7940hs ]
+#        runner_type: [ amd7940hs, amdhx370 ]
 #    steps:
 #      - uses: actions/checkout@v4
 #        with:
@@ -79,157 +184,56 @@ jobs:
 #          sudo prlimit -lunlimited --pid $$
 #          pip cache purge
 #          source /opt/xilinx/xrt/setup.sh
-#          python -m venv aie-venv
-#          source aie-venv/bin/activate
-#          pip install -r python/requirements.txt
-#          HOST_MLIR_PYTHON_PACKAGE_PREFIX=aie pip install -r python/requirements_extras.txt
-#          pip install -r python/requirements_notebook.txt
-#          pip install -r python/requirements_ml.txt
+#          export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
+#          source utils/quick_setup.sh
+#          # quick_setup changes directory to programming_examples, so we need to return to mlir-aie
+#          cd ..
 #          pip install -r python/requirements_test.txt
 #
-#          # Install llvm-aie (Peano) which is needed for unittests requiring compiling NPU code
-#          pip install -I llvm-aie -f https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly
-#          export PEANO_INSTALL_DIR="$(pip show llvm-aie | grep ^Location: | awk '{print $2}')/llvm-aie"
-#          
-#          sed -i.bak 's/OUTPUT_TIMEOUT = 10/OUTPUT_TIMEOUT = 100/g' \
-#            $(python -c 'import site; print(site.getsitepackages()[0])')/jupyter_client/runapp.py
+#          ./utils/build-mlir-aie-from-wheels.sh
 #
-#          VERSION=$(utils/clone-llvm.sh --get-wheel-version)
-#          pip -q download mlir==$VERSION \
-#            -f https://github.com/Xilinx/mlir-aie/releases/expanded_assets/mlir-distro
-#          unzip -q mlir-*.whl
 #          # I have no clue why but the system clock on GHA containers is like 12 hours ahead.
 #          # That means wheels have file with time stamps in the future which makes ninja loop
 #          # forever when configuring. Set the time to some arbitrary stamp in the past just to be safe.
+#          pushd my_install
 #          find mlir -exec touch -a -m -t 201108231405.14 {} \;
+#          popd
 #
-#          mkdir build
+#          # build is created by the build-mlir-aie-from-wheels.sh script
 #          pushd build
 #
 #          # -j here to reduce the number of parallel chess jobs.
 #          # -j4 for 32GB RAM, -j12 for 64GB RAM
 #          if [ x"${{ matrix.runner_type }}" == x"amdhx370" ]; then
 #            LIT_OPTS="-j4 $LIT_OPTS"
-#          else
+#            else
 #            LIT_OPTS="-j12 $LIT_OPTS"
 #          fi
 #
-#          export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
-#          
-#          cmake .. -G Ninja \
-#            -DCMAKE_BUILD_TYPE=Release \
-#            -DPython3_EXECUTABLE=$(which python) \
-#            -DLLVM_EXTERNAL_LIT=$(which lit) \
-#            -DCMAKE_INSTALL_PREFIX=$PWD/../mlir_aie \
-#            -DCMAKE_MODULE_PATH=$PWD/../cmake/modulesXilinx \
-#            -DMLIR_DIR=$PWD/../mlir/lib/cmake/mlir \
-#            -DAIE_ENABLE_BINDINGS_PYTHON=ON \
-#            -DLLVM_ENABLE_RTTI=ON \
-#            -DLLVM_ENABLE_ASSERTIONS=ON \
-#            -DAIE_VITIS_COMPONENTS="AIE2;AIE2P" \
-#            -DAIE_RUNTIME_TARGETS=x86_64 \
-#            -DAIE_RUNTIME_TEST_TARGET=x86_64 \
-#            -DVITIS_VPP=$(which v++) \
-#            $CMAKE_ARGS
-#
-#          ninja install
-#          
+#          ls -alth
 #          cd ..
 #          ls -alth
-#          ls -alth ~/.iron
-#          source /opt/xilinx/xrt/setup.sh
-#          source utils/env_setup.sh
-#          export PYTHONPATH="$(pwd)/install/python:${PYTHONPATH}"
-#          xrt-smi examine
-#          python -c "from aie.iron import get_current_device; from aie.iron.device import NPU1, NPU2; current_device = get_current_device(); target_arch = 'aie2p' if isinstance(current_device, NPU2) else 'aie2' if isinstance(current_device, NPU1) else None; print(f'Device: {current_device}, Target arch: {target_arch}')"
+#          git log -1 | cat
 #          
 #          # Debug: Check what's available in aie.iron
 #          python -c "import aie.iron; print(dir(aie.iron))"
-#          cat /home/github/actions-runner/_work/mlir-aie/mlir-aie/aie-venv/lib/python3.12/site-packages/mlir_aie/python/aie/iron/kernel.py
+#          iron/__init__.py
+#          # Debug: Check Python path and environment
+#          echo "Python executable: $(which python)"
+#          echo "Python version: $(python --version)"
+#          echo "PYTHONPATH: $PYTHONPATH"
+#          python -c "import sys; print('Python path:'); [print(p) for p in sys.path]"
+#          python -c "import aie.iron; print('aie.iron location:', aie.iron.__file__)"
+#          cat python/iron/kernel.py 
 #          
 #          python programming_examples/basic/vector_vector_add/vector_vector_add.py -d npu2
 #          pytest -vvv -s test/python/jit_compilation.py
 #          pytest -vvv -s test/python/jit_extern_functions.py
 #          pytest -vvv -s test/python/jit_extern_inside_jit.py
-#    
-#          #ninja check-aie
-#          popd
 #
-  build-quick-setup:
-    name: Run Examples on Ryzen AI
-    runs-on: ${{ matrix.runner_type }}
-    strategy:
-      fail-fast: false
-      matrix:
-        runner_type: [ amd7940hs, amdhx370 ]
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          submodules: "true"
-
-      # Launch an ssh session via a proxy server if there is a need
-      # for debug. This seems to live for 35 min max
-      # https://github.com/mxschmitt/action-tmate
-      - name: Setup tmate session
-        uses: mxschmitt/action-tmate@v3
-        # To run this, launch it manually on the default branch and
-        # click on "launch_tmate_terminal_for_debug"
-        if: inputs.launch_tmate_terminal_for_debug
-
-      - name: Run commands
-        run: |
-          sudo prlimit -lunlimited --pid $$
-          pip cache purge
-          source /opt/xilinx/xrt/setup.sh
-          export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
-          source utils/quick_setup.sh
-          # quick_setup changes directory to programming_examples, so we need to return to mlir-aie
-          cd ..
-          pip install -r python/requirements_test.txt
-
-          ./utils/build-mlir-aie-from-wheels.sh
-
-          # I have no clue why but the system clock on GHA containers is like 12 hours ahead.
-          # That means wheels have file with time stamps in the future which makes ninja loop
-          # forever when configuring. Set the time to some arbitrary stamp in the past just to be safe.
-          pushd my_install
-          find mlir -exec touch -a -m -t 201108231405.14 {} \;
-          popd
-
-          # build is created by the build-mlir-aie-from-wheels.sh script
-          pushd build
-
-          # -j here to reduce the number of parallel chess jobs.
-          # -j4 for 32GB RAM, -j12 for 64GB RAM
-          if [ x"${{ matrix.runner_type }}" == x"amdhx370" ]; then
-            LIT_OPTS="-j4 $LIT_OPTS"
-            else
-            LIT_OPTS="-j12 $LIT_OPTS"
-          fi
-
-          ls -alth
-          cd ..
-          ls -alth
-          git log -1 | cat
-          
-          # Debug: Check what's available in aie.iron
-          python -c "import aie.iron; print(dir(aie.iron))"
-          
-          # Debug: Check Python path and environment
-          echo "Python executable: $(which python)"
-          echo "Python version: $(python --version)"
-          echo "PYTHONPATH: $PYTHONPATH"
-          python -c "import sys; print('Python path:'); [print(p) for p in sys.path]"
-          python -c "import aie.iron; print('aie.iron location:', aie.iron.__file__)"
-          cat python/iron/kernel.py 
-          
-          python programming_examples/basic/vector_vector_add/vector_vector_add.py -d npu2
-          pytest -vvv -s test/python/jit_compilation.py
-          pytest -vvv -s test/python/jit_extern_functions.py
-          pytest -vvv -s test/python/jit_extern_inside_jit.py
-
-          #ninja install
-          #ninja check-reference-designs
-          #ninja check-programming-guide
-
-          popd
+#          #ninja install
+#          #ninja check-reference-designs
+#          #ninja check-programming-guide
+#
+#          popd
+#
\ No newline at end of file

From c5c9692e408278b359b1fca6f54813cf49afe113 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 23:49:59 -0700
Subject: [PATCH 29/37] Debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 5 +++--
 python/iron/jit.py                        | 2 +-
 test/python/jit_extern_functions.py       | 6 ++++++
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index aebfdcb25c6..c6c6b294e2a 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -139,7 +139,7 @@ jobs:
           ls -alth ~/.iron
           source /opt/xilinx/xrt/setup.sh
           source utils/env_setup.sh
-          export PYTHONPATH="$(pwd)/install/python:${PYTHONPATH}"
+          export PYTHONPATH="$(pwd)/mlir_aie/python:${PYTHONPATH}"
           xrt-smi examine
           python -c "from aie.iron import get_current_device; from aie.iron.device import NPU1, NPU2; current_device = get_current_device(); target_arch = 'aie2p' if isinstance(current_device, NPU2) else 'aie2' if isinstance(current_device, NPU1) else None; print(f'Device: {current_device}, Target arch: {target_arch}')"
           
@@ -150,7 +150,8 @@ jobs:
           cat mlir_aie/python/aie/iron/__init__.py
           cat mlir_aie/python/aie/iron/kernel.py
           export PYTHONPATH="$(pwd)/mlir_aie/python:${PYTHONPATH}"
-          python programming_examples/basic/vector_vector_add/vector_vector_add.py -d npu2
+
+          pyton 
           pytest -vvv -s test/python/jit_compilation.py
           pytest -vvv -s test/python/jit_extern_functions.py
           pytest -vvv -s test/python/jit_extern_inside_jit.py
diff --git a/python/iron/jit.py b/python/iron/jit.py
index 79215c79920..b09b87b1491 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -275,7 +275,7 @@ def compile_external_kernel(func, kernel_dir):
         include_dirs=func._include_dirs,
         compile_args=func._compile_flags,
         cwd=kernel_dir,
-        verbose=False,
+        verbose=True,
     )
 
     # Mark the function as compiled
diff --git a/test/python/jit_extern_functions.py b/test/python/jit_extern_functions.py
index 4a1eeb23a83..0fd766c2f53 100644
--- a/test/python/jit_extern_functions.py
+++ b/test/python/jit_extern_functions.py
@@ -110,11 +110,17 @@ def test_simple_add_one():
     )
 
     # Apply the transform
+    print("Applying transform")
+    print(f"Input tensor: {input_tensor}")
+    print(f"Output tensor: {output_tensor}")
+    print(f"Initial tensor: {initial_tensor}")
     transform(input_tensor, output_tensor, add_one)
 
     # Verify results
     expected = initial_tensor + 1
     actual = output_tensor.numpy()
+    print(f"Actual tensor: {actual}")
+    print(f"Expected tensor: {expected}")
     np.testing.assert_array_equal(actual, expected)
 
 

From 5306493e43c5af4793d6ce2f2c026a8d0f7f01a8 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Tue, 5 Aug 2025 23:53:38 -0700
Subject: [PATCH 30/37] Debug

---
 .github/workflows/buildAndTestRyzenAI.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index c6c6b294e2a..084ddf527df 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -151,7 +151,7 @@ jobs:
           cat mlir_aie/python/aie/iron/kernel.py
           export PYTHONPATH="$(pwd)/mlir_aie/python:${PYTHONPATH}"
 
-          pyton 
+          python test/python/jit_extern_functions.py 
           pytest -vvv -s test/python/jit_compilation.py
           pytest -vvv -s test/python/jit_extern_functions.py
           pytest -vvv -s test/python/jit_extern_inside_jit.py

From ae45e4f9a05d8b7db96b3f448378e0f15f7768a4 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Wed, 6 Aug 2025 00:03:35 -0700
Subject: [PATCH 31/37] Add debugging statments

---
 python/iron/jit.py | 357 +++++++++++++++++++++++++++++++++++++--------
 1 file changed, 296 insertions(+), 61 deletions(-)

diff --git a/python/iron/jit.py b/python/iron/jit.py
index b09b87b1491..9b69d62ba82 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -12,6 +12,8 @@
 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
@@ -42,46 +44,104 @@ def __init__(
             device_index (int, optional): Index of the device. Defaults to 0.
             kernel_name (str, optional): Name of the kernel. Defaults to "PP_FD_PRE".
         """
+        print(f"[DEBUG] NPUKernel.__init__: xclbin_path={xclbin_path}")
+        print(f"[DEBUG] NPUKernel.__init__: insts_path={insts_path}")
+        print(f"[DEBUG] NPUKernel.__init__: device_index={device_index}")
+        print(f"[DEBUG] NPUKernel.__init__: kernel_name={kernel_name}")
 
-        self.__device = xrt.device(device_index)
+        try:
+            self.__device = xrt.device(device_index)
+            print(
+                f"[DEBUG] NPUKernel.__init__: Created device with index {device_index}"
+            )
+        except Exception as e:
+            print(f"[DEBUG] NPUKernel.__init__: Failed to create device: {e}")
+            raise
 
         # Find kernel by name in the xclbin
-        self.__xclbin = xrt.xclbin(xclbin_path)
-        kernels = self.__xclbin.get_kernels()
+        try:
+            self.__xclbin = xrt.xclbin(xclbin_path)
+            print(f"[DEBUG] NPUKernel.__init__: Loaded xclbin from {xclbin_path}")
+            kernels = self.__xclbin.get_kernels()
+            print(f"[DEBUG] NPUKernel.__init__: Found {len(kernels)} kernels in xclbin")
+            for i, k in enumerate(kernels):
+                print(f"[DEBUG] NPUKernel.__init__: Kernel {i}: {k.get_name()}")
+        except Exception as e:
+            print(f"[DEBUG] NPUKernel.__init__: Failed to load xclbin: {e}")
+            raise
 
         try:
             xkernel = [k for k in kernels if kernel_name == k.get_name()][0]
-        except KeyError:
+            print(f"[DEBUG] NPUKernel.__init__: Found kernel '{kernel_name}'")
+        except (KeyError, IndexError) as e:
+            print(
+                f"[DEBUG] NPUKernel.__init__: Failed to find kernel '{kernel_name}': {e}"
+            )
             raise NPUKernel_Error("No such kernel: " + kernel_name)
 
-        self.__device.register_xclbin(self.__xclbin)
-        self.__context = xrt.hw_context(self.__device, self.__xclbin.get_uuid())
-        self.__kernel = xrt.kernel(self.__context, xkernel.get_name())
+        try:
+            self.__device.register_xclbin(self.__xclbin)
+            print(f"[DEBUG] NPUKernel.__init__: Registered xclbin with device")
+            self.__context = xrt.hw_context(self.__device, self.__xclbin.get_uuid())
+            print(f"[DEBUG] NPUKernel.__init__: Created hardware context")
+            self.__kernel = xrt.kernel(self.__context, xkernel.get_name())
+            print(f"[DEBUG] NPUKernel.__init__: Created kernel object")
+        except Exception as e:
+            print(f"[DEBUG] NPUKernel.__init__: Failed to create kernel: {e}")
+            raise
 
         # Set up instruction stream
-        insts = read_insts_binary(insts_path)
-        self.__n_insts = len(insts)
-        insts_buffers_bytes = self.__n_insts * np.dtype(insts.dtype).itemsize
+        try:
+            insts = read_insts_binary(insts_path)
+            print(
+                f"[DEBUG] NPUKernel.__init__: Read {len(insts)} instructions from {insts_path}"
+            )
+            self.__n_insts = len(insts)
+            insts_buffers_bytes = self.__n_insts * np.dtype(insts.dtype).itemsize
+            print(
+                f"[DEBUG] NPUKernel.__init__: Instruction buffer size: {insts_buffers_bytes} bytes"
+            )
+        except Exception as e:
+            print(f"[DEBUG] NPUKernel.__init__: Failed to read instructions: {e}")
+            raise
 
         # Magic number for RyzenAI group id that will be fixed in the future. See same code at XRT:
         # https://github.com/Xilinx/XRT/blob/56222ed5cfd119dff0d5bd920735b87024e8c829/src/runtime_src/core/common/api/xrt_module.cpp#L1621
         group_id = 1
+        print(f"[DEBUG] NPUKernel.__init__: Using group_id={group_id}")
 
-        self.__insts_buffer_bo = xrt.bo(
-            self.__device,
-            insts_buffers_bytes,
-            xrt.bo.cacheable,
-            group_id,
-        )
+        try:
+            self.__insts_buffer_bo = xrt.bo(
+                self.__device,
+                insts_buffers_bytes,
+                xrt.bo.cacheable,
+                group_id,
+            )
+            print(f"[DEBUG] NPUKernel.__init__: Created instruction buffer object")
+        except Exception as e:
+            print(
+                f"[DEBUG] NPUKernel.__init__: Failed to create instruction buffer: {e}"
+            )
+            raise
 
         # Copy into a temporary numpy buffer
-        insts_buffer_bo_np = np.frombuffer(
-            self.__insts_buffer_bo.map(), dtype=insts.dtype
-        ).reshape(insts.shape)
-        insts_buffer_bo_np[:] = insts
+        try:
+            insts_buffer_bo_np = np.frombuffer(
+                self.__insts_buffer_bo.map(), dtype=insts.dtype
+            ).reshape(insts.shape)
+            insts_buffer_bo_np[:] = insts
+            print(f"[DEBUG] NPUKernel.__init__: Copied instructions to buffer")
+        except Exception as e:
+            print(f"[DEBUG] NPUKernel.__init__: Failed to copy instructions: {e}")
+            raise
 
         # Always sync to the device in the constructor
-        self.__insts_buffer_bo.sync(xrt.xclBOSyncDirection.XCL_BO_SYNC_BO_TO_DEVICE)
+        try:
+            self.__insts_buffer_bo.sync(xrt.xclBOSyncDirection.XCL_BO_SYNC_BO_TO_DEVICE)
+            print(f"[DEBUG] NPUKernel.__init__: Synced instructions to device")
+        except Exception as e:
+            print(f"[DEBUG] NPUKernel.__init__: Failed to sync instructions: {e}")
+            raise
 
     # Blocking call.
     def __call__(self, *args):
@@ -91,6 +151,13 @@ def __call__(self, *args):
         Parameters:
             args (IRON Tensors): Arguments to pass to the kernel.
         """
+        print(f"[DEBUG] NPUKernel.__call__: Called with {len(args)} arguments")
+        for i, arg in enumerate(args):
+            print(f"[DEBUG] NPUKernel.__call__: Arg {i}: {type(arg)}")
+            if hasattr(arg, "buffer_object"):
+                print(
+                    f"[DEBUG] NPUKernel.__call__: Arg {i} has buffer_object: {arg.buffer_object()}"
+                )
 
         opcode = 3
         kernel_args = []
@@ -98,23 +165,49 @@ def __call__(self, *args):
         for tensor in args:
             # Skip callable arguments since these are inlined in the kernel
             if callable(tensor):
+                print(f"[DEBUG] NPUKernel.__call__: Skipping callable argument")
                 continue
             if not hasattr(tensor, "buffer_object"):
+                print(
+                    f"[DEBUG] NPUKernel.__call__: Argument {type(tensor)} has no 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:
-            raise NPUKernel_Error(f"Kernel returned {r}")
+            print(f"[DEBUG] NPUKernel.__call__: Added buffer object to kernel args")
+
+        print(
+            f"[DEBUG] NPUKernel.__call__: Calling kernel with opcode={opcode}, n_insts={self.__n_insts}, {len(kernel_args)} buffer args"
+        )
+        try:
+            h = self.__kernel(
+                opcode, self.__insts_buffer_bo, self.__n_insts, *kernel_args
+            )
+            print(f"[DEBUG] NPUKernel.__call__: Kernel execution started")
+            r = h.wait()
+            print(
+                f"[DEBUG] NPUKernel.__call__: Kernel execution completed with result: {r}"
+            )
+            if r != xrt.ert_cmd_state.ERT_CMD_STATE_COMPLETED:
+                print(f"[DEBUG] NPUKernel.__call__: Kernel failed with state: {r}")
+                raise NPUKernel_Error(f"Kernel returned {r}")
+        except Exception as e:
+            print(f"[DEBUG] NPUKernel.__call__: Exception during kernel execution: {e}")
+            print(f"[DEBUG] NPUKernel.__call__: Exception traceback:")
+            traceback.print_exc()
+            raise
 
     def __del__(self):
         """
         Destructor to clean up resources and delete the kernel and device objects.
         """
-        del self.__kernel
-        del self.__device
+        print(f"[DEBUG] NPUKernel.__del__: Cleaning up resources")
+        try:
+            del self.__kernel
+            del self.__device
+        except Exception as e:
+            print(f"[DEBUG] NPUKernel.__del__: Exception during cleanup: {e}")
 
 
 class NPUKernel_Error(Exception):
@@ -139,31 +232,56 @@ def jit(function=None, is_placed=True, use_cache=True):
 
     @functools.wraps(function)
     def decorator(*args, **kwargs):
+        print(
+            f"[DEBUG] jit.decorator: Starting compilation for function {function.__name__}"
+        )
+        print(f"[DEBUG] jit.decorator: is_placed={is_placed}, use_cache={use_cache}")
+        print(
+            f"[DEBUG] jit.decorator: args count={len(args)}, kwargs count={len(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()
+        print(f"[DEBUG] jit.decorator: Cleared ExternalKernel instances")
 
         # Find ExternalKernel instances in arguments and kwargs
         external_kernels = []
         for arg in args:
             if isinstance(arg, ExternalKernel):
                 external_kernels.append(arg)
+                print(
+                    f"[DEBUG] jit.decorator: Found ExternalKernel in args: {arg._name}"
+                )
         for value in kwargs.values():
             if isinstance(value, ExternalKernel):
                 external_kernels.append(value)
+                print(
+                    f"[DEBUG] jit.decorator: Found ExternalKernel in kwargs: {value._name}"
+                )
 
         # Execute the function to generate MLIR
-        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)
+        try:
+            if is_placed:
+                print(f"[DEBUG] jit.decorator: Generating MLIR with placement")
+                with mlir_mod_ctx() as ctx:
+                    function(*args, **kwargs)
+                    print(f"[DEBUG] jit.decorator: Function executed successfully")
+                    assert (
+                        ctx.module.operation.verify()
+                    ), f"Verification failed for '{function.__name__}'"
+                    print(f"[DEBUG] jit.decorator: Module verification passed")
+                    mlir_module = ctx.module
+            else:
+                print(f"[DEBUG] jit.decorator: Generating MLIR without placement")
+                mlir_module = function(*args, **kwargs)
+                print(f"[DEBUG] jit.decorator: Function returned MLIR module")
+        except Exception as e:
+            print(f"[DEBUG] jit.decorator: Exception during MLIR generation: {e}")
+            traceback.print_exc()
+            raise
 
         # Compile all ExternalKernel instances that were created during this JIT compilation
         for func in ExternalKernel._instances:
@@ -171,14 +289,20 @@ def decorator(*args, **kwargs):
                 not hasattr(func, "_compiled") or not func._compiled
             ):  # Don't compile if already compiled
                 external_kernels.append(func)
+                print(
+                    f"[DEBUG] jit.decorator: Added ExternalKernel to compilation list: {func._name}"
+                )
 
         # 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")
+        print(f"[DEBUG] jit.decorator: Module hash: {module_hash}")
+        print(f"[DEBUG] jit.decorator: Kernel directory: {kernel_dir}")
 
         # Ensure cache directory exists
         os.makedirs(kernel_dir, exist_ok=True)
+        print(f"[DEBUG] jit.decorator: Created/verified kernel directory")
 
         # Write MLIR to file if not already cached
         inst_filename = "insts.bin"
@@ -188,16 +312,27 @@ def decorator(*args, **kwargs):
 
         xclbin_exists = os.path.exists(xclbin_path)
         inst_exists = os.path.exists(inst_path)
+        print(f"[DEBUG] jit.decorator: xclbin exists: {xclbin_exists}")
+        print(f"[DEBUG] jit.decorator: inst file exists: {inst_exists}")
 
         if not use_cache or not xclbin_exists or not inst_exists:
+            print(
+                f"[DEBUG] jit.decorator: Need to compile (cache disabled or files missing)"
+            )
             try:
                 with open(mlir_path, "w", encoding="utf-8") as f:
                     print(mlir_module, file=f)
+                print(f"[DEBUG] jit.decorator: Wrote MLIR module to {mlir_path}")
 
                 # Set cache directory for ExternalKernels and compile them
                 for func in external_kernels:
+                    print(
+                        f"[DEBUG] jit.decorator: Compiling ExternalKernel: {func._name}"
+                    )
                     # Compile the ExternalKernel directly in the kernel directory
                     compile_external_kernel(func, kernel_dir)
+
+                print(f"[DEBUG] jit.decorator: Starting MLIR module compilation")
                 # Compile the MLIR module
                 compile_mlir_module_to_binary(
                     mlir_module=mlir_module,
@@ -205,16 +340,36 @@ def decorator(*args, **kwargs):
                     xclbin_path=xclbin_path,
                     work_dir=kernel_dir,
                 )
+                print(f"[DEBUG] jit.decorator: MLIR module compilation completed")
             except Exception as e:
+                print(f"[DEBUG] jit.decorator: Exception during compilation: {e}")
+                traceback.print_exc()
                 # Clean up cache directory on any compilation failure
                 if os.path.exists(kernel_dir):
                     shutil.rmtree(kernel_dir)
+                    print(
+                        f"[DEBUG] jit.decorator: Cleaned up kernel directory after failure"
+                    )
                 raise e
+        else:
+            print(f"[DEBUG] jit.decorator: Using cached compilation results")
 
         kernel_name = "MLIR_AIE"
-        return NPUKernel(xclbin_path, inst_path, kernel_name=kernel_name)(
-            *args, **kwargs
+        print(
+            f"[DEBUG] jit.decorator: Creating NPUKernel with kernel_name={kernel_name}"
         )
+        try:
+            kernel = NPUKernel(xclbin_path, inst_path, kernel_name=kernel_name)
+            print(f"[DEBUG] jit.decorator: NPUKernel created successfully")
+            result = kernel(*args, **kwargs)
+            print(f"[DEBUG] jit.decorator: Kernel execution completed successfully")
+            return result
+        except Exception as e:
+            print(
+                f"[DEBUG] jit.decorator: Exception during kernel creation/execution: {e}"
+            )
+            traceback.print_exc()
+            raise
 
     return decorator
 
@@ -227,77 +382,157 @@ def compile_external_kernel(func, kernel_dir):
         func: ExternalKernel instance to compile
         kernel_dir: Directory to place the compiled object file
     """
+    print(f"[DEBUG] compile_external_kernel: Starting compilation of {func._name}")
+    print(f"[DEBUG] compile_external_kernel: kernel_dir={kernel_dir}")
+
     # Skip if already compiled
     if hasattr(func, "_compiled") and func._compiled:
+        print(f"[DEBUG] compile_external_kernel: Function already compiled, skipping")
         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):
+        print(
+            f"[DEBUG] compile_external_kernel: Object file already exists: {output_file}"
+        )
         return
 
     # Create source file in kernel directory
     source_file = os.path.join(kernel_dir, f"{func._name}.cc")
+    print(f"[DEBUG] compile_external_kernel: Source file path: {source_file}")
 
     # 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)
+        try:
+            with open(source_file, "w") as f:
+                f.write(func._source_string)
+            print(
+                f"[DEBUG] compile_external_kernel: Wrote source from string to {source_file}"
+            )
+            print(
+                f"[DEBUG] compile_external_kernel: Source content:\n{func._source_string}"
+            )
+        except Exception as e:
+            print(f"[DEBUG] compile_external_kernel: Failed to write source file: {e}")
+            raise
     elif func._source_file is not None:
         # Use source_file (copy existing file)
+        print(
+            f"[DEBUG] compile_external_kernel: Copying source file: {func._source_file}"
+        )
 
         # Check if source file exists before copying
         if os.path.exists(func._source_file):
-            shutil.copy2(func._source_file, source_file)
+            try:
+                shutil.copy2(func._source_file, source_file)
+                print(
+                    f"[DEBUG] compile_external_kernel: Copied source file successfully"
+                )
+            except Exception as e:
+                print(
+                    f"[DEBUG] compile_external_kernel: Failed to copy source file: {e}"
+                )
+                raise
         else:
+            print(
+                f"[DEBUG] compile_external_kernel: Source file does not exist: {func._source_file}"
+            )
             return
     else:
+        print(
+            f"[DEBUG] compile_external_kernel: Neither source_string nor source_file provided"
+        )
         raise ValueError("Neither source_string nor source_file is provided")
 
     # Add device-specific flags based on actual device detection
-    current_device = get_current_device()
-
-    # Determine target architecture based on device type
-    if isinstance(current_device, NPU2):
-        target_arch = "aie2p"
-    elif isinstance(current_device, NPU1):
-        target_arch = "aie2"
-    else:
-        raise RuntimeError(f"Unsupported device type: {type(current_device)}")
+    try:
+        current_device = get_current_device()
+        print(f"[DEBUG] compile_external_kernel: Current device: {current_device}")
+        print(f"[DEBUG] compile_external_kernel: Device type: {type(current_device)}")
+
+        # Determine target architecture based on device type
+        if isinstance(current_device, NPU2):
+            target_arch = "aie2p"
+        elif isinstance(current_device, NPU1):
+            target_arch = "aie2"
+        else:
+            print(
+                f"[DEBUG] compile_external_kernel: Unsupported device type: {type(current_device)}"
+            )
+            raise RuntimeError(f"Unsupported device type: {type(current_device)}")
+
+        print(f"[DEBUG] compile_external_kernel: Target architecture: {target_arch}")
+    except Exception as e:
+        print(
+            f"[DEBUG] compile_external_kernel: Failed to determine target architecture: {e}"
+        )
+        raise
 
     from .compile.compile import compile_cxx_core_function
 
-    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=True,
-    )
+    print(f"[DEBUG] compile_external_kernel: Calling compile_cxx_core_function")
+    print(f"[DEBUG] compile_external_kernel: source_path={source_file}")
+    print(f"[DEBUG] compile_external_kernel: target_arch={target_arch}")
+    print(f"[DEBUG] compile_external_kernel: output_path={output_file}")
+    print(f"[DEBUG] compile_external_kernel: include_dirs={func._include_dirs}")
+    print(f"[DEBUG] compile_external_kernel: compile_args={func._compile_flags}")
+    print(f"[DEBUG] compile_external_kernel: cwd={kernel_dir}")
+
+    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=True,
+        )
+        print(f"[DEBUG] compile_external_kernel: Compilation completed successfully")
+    except Exception as e:
+        print(f"[DEBUG] compile_external_kernel: Compilation failed: {e}")
+        traceback.print_exc()
+        raise
 
     # Mark the function as compiled
     func._compiled = True
+    print(f"[DEBUG] compile_external_kernel: Marked function as compiled")
 
 
 def hash_module(module, external_kernels=None):
     """
     Hash the MLIR module and ExternalKernel compiler options to create a unique identifier.
     """
+    print(f"[DEBUG] hash_module: Starting hash computation")
     mlir_str = str(module)
+    print(f"[DEBUG] hash_module: MLIR string length: {len(mlir_str)}")
 
     # Include ExternalKernel compiler options in the hash
     if external_kernels:
+        print(
+            f"[DEBUG] hash_module: Processing {len(external_kernels)} 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)
+            print(
+                f"[DEBUG] hash_module: Added options for {func._name}: include_dirs={func._include_dirs}, compile_flags={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]
+        hash_result = hashlib.sha256(combined_str.encode("utf-8")).hexdigest()[:16]
+        print(
+            f"[DEBUG] hash_module: Computed hash with external kernels: {hash_result}"
+        )
+        return hash_result
     else:
-        return hashlib.sha256(mlir_str.encode("utf-8")).hexdigest()[:16]
+        hash_result = hashlib.sha256(mlir_str.encode("utf-8")).hexdigest()[:16]
+        print(
+            f"[DEBUG] hash_module: Computed hash without external kernels: {hash_result}"
+        )
+        return hash_result

From 668935653a17caebb92433129163cf96f5510d53 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Wed, 6 Aug 2025 00:17:32 -0700
Subject: [PATCH 32/37] Include the device in the hash

---
 python/iron/jit.py | 81 ++++++++++++++++++++++++----------------------
 1 file changed, 42 insertions(+), 39 deletions(-)

diff --git a/python/iron/jit.py b/python/iron/jit.py
index 9b69d62ba82..ef57f7bc7eb 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -293,8 +293,32 @@ def decorator(*args, **kwargs):
                     f"[DEBUG] jit.decorator: Added ExternalKernel to compilation list: {func._name}"
                 )
 
-        # Hash of the IR string and ExternalKernel compiler options
-        module_hash = hash_module(mlir_module, external_kernels)
+        # Determine target architecture based on device type (hoisted from compile_external_kernel)
+        try:
+            current_device = get_current_device()
+            print(f"[DEBUG] jit.decorator: Current device: {current_device}")
+            print(f"[DEBUG] jit.decorator: Device type: {type(current_device)}")
+
+            # Determine target architecture based on device type
+            if isinstance(current_device, NPU2):
+                target_arch = "aie2p"
+            elif isinstance(current_device, NPU1):
+                target_arch = "aie2"
+            else:
+                print(
+                    f"[DEBUG] jit.decorator: Unsupported device type: {type(current_device)}"
+                )
+                raise RuntimeError(f"Unsupported device type: {type(current_device)}")
+
+            print(f"[DEBUG] jit.decorator: Target architecture: {target_arch}")
+        except Exception as e:
+            print(
+                f"[DEBUG] jit.decorator: Failed to determine target architecture: {e}"
+            )
+            raise
+
+        # Hash of the IR string, ExternalKernel 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")
         print(f"[DEBUG] jit.decorator: Module hash: {module_hash}")
@@ -330,7 +354,7 @@ def decorator(*args, **kwargs):
                         f"[DEBUG] jit.decorator: Compiling ExternalKernel: {func._name}"
                     )
                     # Compile the ExternalKernel directly in the kernel directory
-                    compile_external_kernel(func, kernel_dir)
+                    compile_external_kernel(func, kernel_dir, target_arch)
 
                 print(f"[DEBUG] jit.decorator: Starting MLIR module compilation")
                 # Compile the MLIR module
@@ -374,16 +398,18 @@ def decorator(*args, **kwargs):
     return decorator
 
 
-def compile_external_kernel(func, kernel_dir):
+def compile_external_kernel(func, kernel_dir, target_arch):
     """
     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
+        target_arch: Target architecture (e.g., "aie2" or "aie2p")
     """
     print(f"[DEBUG] compile_external_kernel: Starting compilation of {func._name}")
     print(f"[DEBUG] compile_external_kernel: kernel_dir={kernel_dir}")
+    print(f"[DEBUG] compile_external_kernel: target_arch={target_arch}")
 
     # Skip if already compiled
     if hasattr(func, "_compiled") and func._compiled:
@@ -446,30 +472,6 @@ def compile_external_kernel(func, kernel_dir):
         )
         raise ValueError("Neither source_string nor source_file is provided")
 
-    # Add device-specific flags based on actual device detection
-    try:
-        current_device = get_current_device()
-        print(f"[DEBUG] compile_external_kernel: Current device: {current_device}")
-        print(f"[DEBUG] compile_external_kernel: Device type: {type(current_device)}")
-
-        # Determine target architecture based on device type
-        if isinstance(current_device, NPU2):
-            target_arch = "aie2p"
-        elif isinstance(current_device, NPU1):
-            target_arch = "aie2"
-        else:
-            print(
-                f"[DEBUG] compile_external_kernel: Unsupported device type: {type(current_device)}"
-            )
-            raise RuntimeError(f"Unsupported device type: {type(current_device)}")
-
-        print(f"[DEBUG] compile_external_kernel: Target architecture: {target_arch}")
-    except Exception as e:
-        print(
-            f"[DEBUG] compile_external_kernel: Failed to determine target architecture: {e}"
-        )
-        raise
-
     from .compile.compile import compile_cxx_core_function
 
     print(f"[DEBUG] compile_external_kernel: Calling compile_cxx_core_function")
@@ -501,13 +503,14 @@ def compile_external_kernel(func, kernel_dir):
     print(f"[DEBUG] compile_external_kernel: Marked function as compiled")
 
 
-def hash_module(module, external_kernels=None):
+def hash_module(module, external_kernels=None, target_arch=None):
     """
     Hash the MLIR module and ExternalKernel compiler options to create a unique identifier.
     """
     print(f"[DEBUG] hash_module: Starting hash computation")
     mlir_str = str(module)
     print(f"[DEBUG] hash_module: MLIR string length: {len(mlir_str)}")
+    print(f"[DEBUG] hash_module: Target architecture: {target_arch}")
 
     # Include ExternalKernel compiler options in the hash
     if external_kernels:
@@ -525,14 +528,14 @@ def hash_module(module, external_kernels=None):
 
         # Create a combined string for hashing
         combined_str = mlir_str + "|" + "|".join(compiler_options)
-        hash_result = hashlib.sha256(combined_str.encode("utf-8")).hexdigest()[:16]
-        print(
-            f"[DEBUG] hash_module: Computed hash with external kernels: {hash_result}"
-        )
-        return hash_result
     else:
-        hash_result = hashlib.sha256(mlir_str.encode("utf-8")).hexdigest()[:16]
-        print(
-            f"[DEBUG] hash_module: Computed hash without external kernels: {hash_result}"
-        )
-        return hash_result
+        combined_str = mlir_str
+
+    # Include target architecture in the hash
+    if target_arch:
+        combined_str += f"|target_arch={target_arch}"
+        print(f"[DEBUG] hash_module: Added target architecture to hash: {target_arch}")
+
+    hash_result = hashlib.sha256(combined_str.encode("utf-8")).hexdigest()[:16]
+    print(f"[DEBUG] hash_module: Computed hash: {hash_result}")
+    return hash_result

From 561dbe3265ce0a5cc8e874d37f91ba2a37e7ab93 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Wed, 6 Aug 2025 00:30:13 -0700
Subject: [PATCH 33/37] Revert debug code

---
 .github/workflows/buildAndTestRyzenAI.yml | 176 ++++++---------
 python/iron/jit.py                        | 260 +++-------------------
 test/python/jit_extern_functions.py       |   4 -
 3 files changed, 90 insertions(+), 350 deletions(-)

diff --git a/.github/workflows/buildAndTestRyzenAI.yml b/.github/workflows/buildAndTestRyzenAI.yml
index 084ddf527df..a316b11c999 100644
--- a/.github/workflows/buildAndTestRyzenAI.yml
+++ b/.github/workflows/buildAndTestRyzenAI.yml
@@ -22,7 +22,7 @@ on:
         type: boolean
         description: 'Run the build with tmate debugging enabled (https://github.com/marketplace/actions/debugging-with-tmate)'
         required: false
-        default: true
+        default: false
 
 defaults:
   run:
@@ -46,7 +46,7 @@ env:
             -DCMAKE_MODULE_LINKER_FLAGS_INIT="-fuse-ld=lld" \
             -DCMAKE_SHARED_LINKER_FLAGS_INIT="-fuse-ld=lld" \
             -DXRT_ROOT=/opt/xilinx/xrt \
-            -DAIE_VITIS_COMPONENTS="AIE2;AIE2P" \
+            -DAIE_VITIS_COMPONENTS=AIE2;AIE2P \
             -DAIE_ENABLE_PYTHON_PASSES=OFF \
             -DAIE_ENABLE_XRT_PYTHON_BINDINGS=ON \
             -DAIE_INCLUDE_INTEGRATION_TESTS=OFF
@@ -59,7 +59,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        runner_type: [ amdhx370 , amd7940hs ]
+        runner_type: [ amd7940hs, amdhx370 ]
     steps:
       - uses: actions/checkout@v4
         with:
@@ -72,7 +72,8 @@ jobs:
         uses: mxschmitt/action-tmate@v3
         # To run this, launch it manually on the default branch and
         # click on "launch_tmate_terminal_for_debug"
-        if: inputs.launch_tmate_terminal_for_debug
+        if: github.event_name == 'workflow_dispatch'
+                && inputs.launch_tmate_terminal_for_debug
 
       - name: Run commands
         run: |
@@ -123,118 +124,67 @@ jobs:
             -DCMAKE_INSTALL_PREFIX=$PWD/../mlir_aie \
             -DCMAKE_MODULE_PATH=$PWD/../cmake/modulesXilinx \
             -DMLIR_DIR=$PWD/../mlir/lib/cmake/mlir \
-            -DAIE_ENABLE_BINDINGS_PYTHON=ON \
-            -DLLVM_ENABLE_RTTI=ON \
-            -DLLVM_ENABLE_ASSERTIONS=ON \
-            -DAIE_VITIS_COMPONENTS="AIE2;AIE2P" \
-            -DAIE_RUNTIME_TARGETS=x86_64 \
-            -DAIE_RUNTIME_TEST_TARGET=x86_64 \
-            -DVITIS_VPP=$(which v++) \
             $CMAKE_ARGS
 
           ninja install
-          
-          cd ..
-          ls -alth
-          ls -alth ~/.iron
+          ninja check-aie
+          popd
+
+  build-quick-setup:
+    name: Run Examples on Ryzen AI
+    runs-on: ${{ matrix.runner_type }}
+    strategy:
+      fail-fast: false
+      matrix:
+        runner_type: [ amd7940hs, amdhx370 ]
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          submodules: "true"
+
+      # Launch an ssh session via a proxy server if there is a need
+      # for debug. This seems to live for 35 min max
+      # https://github.com/mxschmitt/action-tmate
+      - name: Setup tmate session
+        uses: mxschmitt/action-tmate@v3
+        # To run this, launch it manually on the default branch and
+        # click on "launch_tmate_terminal_for_debug"
+        if: github.event_name == 'workflow_dispatch'
+                && inputs.launch_tmate_terminal_for_debug
+
+      - name: Run commands
+        run: |
+          sudo prlimit -lunlimited --pid $$
+          pip cache purge
           source /opt/xilinx/xrt/setup.sh
-          source utils/env_setup.sh
-          export PYTHONPATH="$(pwd)/mlir_aie/python:${PYTHONPATH}"
-          xrt-smi examine
-          python -c "from aie.iron import get_current_device; from aie.iron.device import NPU1, NPU2; current_device = get_current_device(); target_arch = 'aie2p' if isinstance(current_device, NPU2) else 'aie2' if isinstance(current_device, NPU1) else None; print(f'Device: {current_device}, Target arch: {target_arch}')"
-          
-          # Debug: Check what's available in aie.iron
-          python -c "import aie.iron; print(dir(aie.iron))"
-          cat /home/github/actions-runner/_work/mlir-aie/mlir-aie/aie-venv/lib/python3.12/site-packages/mlir_aie/python/aie/iron/kernel.py
-          ls -alt mlir_aie
-          cat mlir_aie/python/aie/iron/__init__.py
-          cat mlir_aie/python/aie/iron/kernel.py
-          export PYTHONPATH="$(pwd)/mlir_aie/python:${PYTHONPATH}"
-
-          python test/python/jit_extern_functions.py 
-          pytest -vvv -s test/python/jit_compilation.py
-          pytest -vvv -s test/python/jit_extern_functions.py
-          pytest -vvv -s test/python/jit_extern_inside_jit.py
-    
-          #ninja check-aie
+          export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
+          source utils/quick_setup.sh
+          # quick_setup changes directory to programming_examples, so we need to return to mlir-aie
+          cd ..
+          pip install -r python/requirements_test.txt
+
+          ./utils/build-mlir-aie-from-wheels.sh
+
+          # I have no clue why but the system clock on GHA containers is like 12 hours ahead.
+          # That means wheels have file with time stamps in the future which makes ninja loop
+          # forever when configuring. Set the time to some arbitrary stamp in the past just to be safe.
+          pushd my_install
+          find mlir -exec touch -a -m -t 201108231405.14 {} \;
           popd
 
-#  build-quick-setup:
-#    name: Run Examples on Ryzen AI
-#    runs-on: ${{ matrix.runner_type }}
-#    strategy:
-#      fail-fast: false
-#      matrix:
-#        runner_type: [ amd7940hs, amdhx370 ]
-#    steps:
-#      - uses: actions/checkout@v4
-#        with:
-#          submodules: "true"
-#
-#      # Launch an ssh session via a proxy server if there is a need
-#      # for debug. This seems to live for 35 min max
-#      # https://github.com/mxschmitt/action-tmate
-#      - name: Setup tmate session
-#        uses: mxschmitt/action-tmate@v3
-#        # To run this, launch it manually on the default branch and
-#        # click on "launch_tmate_terminal_for_debug"
-#        if: inputs.launch_tmate_terminal_for_debug
-#
-#      - name: Run commands
-#        run: |
-#          sudo prlimit -lunlimited --pid $$
-#          pip cache purge
-#          source /opt/xilinx/xrt/setup.sh
-#          export PATH=$VITIS/bin:$VITIS/aietools/bin:$PATH
-#          source utils/quick_setup.sh
-#          # quick_setup changes directory to programming_examples, so we need to return to mlir-aie
-#          cd ..
-#          pip install -r python/requirements_test.txt
-#
-#          ./utils/build-mlir-aie-from-wheels.sh
-#
-#          # I have no clue why but the system clock on GHA containers is like 12 hours ahead.
-#          # That means wheels have file with time stamps in the future which makes ninja loop
-#          # forever when configuring. Set the time to some arbitrary stamp in the past just to be safe.
-#          pushd my_install
-#          find mlir -exec touch -a -m -t 201108231405.14 {} \;
-#          popd
-#
-#          # build is created by the build-mlir-aie-from-wheels.sh script
-#          pushd build
-#
-#          # -j here to reduce the number of parallel chess jobs.
-#          # -j4 for 32GB RAM, -j12 for 64GB RAM
-#          if [ x"${{ matrix.runner_type }}" == x"amdhx370" ]; then
-#            LIT_OPTS="-j4 $LIT_OPTS"
-#            else
-#            LIT_OPTS="-j12 $LIT_OPTS"
-#          fi
-#
-#          ls -alth
-#          cd ..
-#          ls -alth
-#          git log -1 | cat
-#          
-#          # Debug: Check what's available in aie.iron
-#          python -c "import aie.iron; print(dir(aie.iron))"
-#          iron/__init__.py
-#          # Debug: Check Python path and environment
-#          echo "Python executable: $(which python)"
-#          echo "Python version: $(python --version)"
-#          echo "PYTHONPATH: $PYTHONPATH"
-#          python -c "import sys; print('Python path:'); [print(p) for p in sys.path]"
-#          python -c "import aie.iron; print('aie.iron location:', aie.iron.__file__)"
-#          cat python/iron/kernel.py 
-#          
-#          python programming_examples/basic/vector_vector_add/vector_vector_add.py -d npu2
-#          pytest -vvv -s test/python/jit_compilation.py
-#          pytest -vvv -s test/python/jit_extern_functions.py
-#          pytest -vvv -s test/python/jit_extern_inside_jit.py
-#
-#          #ninja install
-#          #ninja check-reference-designs
-#          #ninja check-programming-guide
-#
-#          popd
-#
\ No newline at end of file
+          # build is created by the build-mlir-aie-from-wheels.sh script
+          pushd build
+
+          # -j here to reduce the number of parallel chess jobs.
+          # -j4 for 32GB RAM, -j12 for 64GB RAM
+          if [ x"${{ matrix.runner_type }}" == x"amdhx370" ]; then
+            LIT_OPTS="-j4 $LIT_OPTS"
+          else
+            LIT_OPTS="-j12 $LIT_OPTS"
+          fi
+
+          ninja install
+          ninja check-reference-designs
+          ninja check-programming-guide
+
+          popd
\ No newline at end of file
diff --git a/python/iron/jit.py b/python/iron/jit.py
index ef57f7bc7eb..e8cccc9833e 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -44,104 +44,46 @@ def __init__(
             device_index (int, optional): Index of the device. Defaults to 0.
             kernel_name (str, optional): Name of the kernel. Defaults to "PP_FD_PRE".
         """
-        print(f"[DEBUG] NPUKernel.__init__: xclbin_path={xclbin_path}")
-        print(f"[DEBUG] NPUKernel.__init__: insts_path={insts_path}")
-        print(f"[DEBUG] NPUKernel.__init__: device_index={device_index}")
-        print(f"[DEBUG] NPUKernel.__init__: kernel_name={kernel_name}")
 
-        try:
-            self.__device = xrt.device(device_index)
-            print(
-                f"[DEBUG] NPUKernel.__init__: Created device with index {device_index}"
-            )
-        except Exception as e:
-            print(f"[DEBUG] NPUKernel.__init__: Failed to create device: {e}")
-            raise
+        self.__device = xrt.device(device_index)
 
         # Find kernel by name in the xclbin
-        try:
-            self.__xclbin = xrt.xclbin(xclbin_path)
-            print(f"[DEBUG] NPUKernel.__init__: Loaded xclbin from {xclbin_path}")
-            kernels = self.__xclbin.get_kernels()
-            print(f"[DEBUG] NPUKernel.__init__: Found {len(kernels)} kernels in xclbin")
-            for i, k in enumerate(kernels):
-                print(f"[DEBUG] NPUKernel.__init__: Kernel {i}: {k.get_name()}")
-        except Exception as e:
-            print(f"[DEBUG] NPUKernel.__init__: Failed to load xclbin: {e}")
-            raise
+        self.__xclbin = xrt.xclbin(xclbin_path)
+        kernels = self.__xclbin.get_kernels()
 
         try:
             xkernel = [k for k in kernels if kernel_name == k.get_name()][0]
-            print(f"[DEBUG] NPUKernel.__init__: Found kernel '{kernel_name}'")
         except (KeyError, IndexError) as e:
-            print(
-                f"[DEBUG] NPUKernel.__init__: Failed to find kernel '{kernel_name}': {e}"
-            )
             raise NPUKernel_Error("No such kernel: " + kernel_name)
 
-        try:
-            self.__device.register_xclbin(self.__xclbin)
-            print(f"[DEBUG] NPUKernel.__init__: Registered xclbin with device")
-            self.__context = xrt.hw_context(self.__device, self.__xclbin.get_uuid())
-            print(f"[DEBUG] NPUKernel.__init__: Created hardware context")
-            self.__kernel = xrt.kernel(self.__context, xkernel.get_name())
-            print(f"[DEBUG] NPUKernel.__init__: Created kernel object")
-        except Exception as e:
-            print(f"[DEBUG] NPUKernel.__init__: Failed to create kernel: {e}")
-            raise
+        self.__device.register_xclbin(self.__xclbin)
+        self.__context = xrt.hw_context(self.__device, self.__xclbin.get_uuid())
+        self.__kernel = xrt.kernel(self.__context, xkernel.get_name())
 
         # Set up instruction stream
-        try:
-            insts = read_insts_binary(insts_path)
-            print(
-                f"[DEBUG] NPUKernel.__init__: Read {len(insts)} instructions from {insts_path}"
-            )
-            self.__n_insts = len(insts)
-            insts_buffers_bytes = self.__n_insts * np.dtype(insts.dtype).itemsize
-            print(
-                f"[DEBUG] NPUKernel.__init__: Instruction buffer size: {insts_buffers_bytes} bytes"
-            )
-        except Exception as e:
-            print(f"[DEBUG] NPUKernel.__init__: Failed to read instructions: {e}")
-            raise
+        insts = read_insts_binary(insts_path)
+        self.__n_insts = len(insts)
+        insts_buffers_bytes = self.__n_insts * np.dtype(insts.dtype).itemsize
 
         # Magic number for RyzenAI group id that will be fixed in the future. See same code at XRT:
         # https://github.com/Xilinx/XRT/blob/56222ed5cfd119dff0d5bd920735b87024e8c829/src/runtime_src/core/common/api/xrt_module.cpp#L1621
         group_id = 1
-        print(f"[DEBUG] NPUKernel.__init__: Using group_id={group_id}")
 
-        try:
-            self.__insts_buffer_bo = xrt.bo(
-                self.__device,
-                insts_buffers_bytes,
-                xrt.bo.cacheable,
-                group_id,
-            )
-            print(f"[DEBUG] NPUKernel.__init__: Created instruction buffer object")
-        except Exception as e:
-            print(
-                f"[DEBUG] NPUKernel.__init__: Failed to create instruction buffer: {e}"
-            )
-            raise
+        self.__insts_buffer_bo = xrt.bo(
+            self.__device,
+            insts_buffers_bytes,
+            xrt.bo.cacheable,
+            group_id,
+        )
 
         # Copy into a temporary numpy buffer
-        try:
-            insts_buffer_bo_np = np.frombuffer(
-                self.__insts_buffer_bo.map(), dtype=insts.dtype
-            ).reshape(insts.shape)
-            insts_buffer_bo_np[:] = insts
-            print(f"[DEBUG] NPUKernel.__init__: Copied instructions to buffer")
-        except Exception as e:
-            print(f"[DEBUG] NPUKernel.__init__: Failed to copy instructions: {e}")
-            raise
+        insts_buffer_bo_np = np.frombuffer(
+            self.__insts_buffer_bo.map(), dtype=insts.dtype
+        ).reshape(insts.shape)
+        insts_buffer_bo_np[:] = insts
 
         # Always sync to the device in the constructor
-        try:
-            self.__insts_buffer_bo.sync(xrt.xclBOSyncDirection.XCL_BO_SYNC_BO_TO_DEVICE)
-            print(f"[DEBUG] NPUKernel.__init__: Synced instructions to device")
-        except Exception as e:
-            print(f"[DEBUG] NPUKernel.__init__: Failed to sync instructions: {e}")
-            raise
+        self.__insts_buffer_bo.sync(xrt.xclBOSyncDirection.XCL_BO_SYNC_BO_TO_DEVICE)
 
     # Blocking call.
     def __call__(self, *args):
@@ -151,13 +93,6 @@ def __call__(self, *args):
         Parameters:
             args (IRON Tensors): Arguments to pass to the kernel.
         """
-        print(f"[DEBUG] NPUKernel.__call__: Called with {len(args)} arguments")
-        for i, arg in enumerate(args):
-            print(f"[DEBUG] NPUKernel.__call__: Arg {i}: {type(arg)}")
-            if hasattr(arg, "buffer_object"):
-                print(
-                    f"[DEBUG] NPUKernel.__call__: Arg {i} has buffer_object: {arg.buffer_object()}"
-                )
 
         opcode = 3
         kernel_args = []
@@ -165,49 +100,27 @@ def __call__(self, *args):
         for tensor in args:
             # Skip callable arguments since these are inlined in the kernel
             if callable(tensor):
-                print(f"[DEBUG] NPUKernel.__call__: Skipping callable argument")
                 continue
             if not hasattr(tensor, "buffer_object"):
-                print(
-                    f"[DEBUG] NPUKernel.__call__: Argument {type(tensor)} has no buffer_object"
-                )
                 raise TypeError(
                     f"Expected Tensor with .buffer_object(), got {type(tensor)}"
                 )
             kernel_args.append(tensor.buffer_object())
-            print(f"[DEBUG] NPUKernel.__call__: Added buffer object to kernel args")
 
-        print(
-            f"[DEBUG] NPUKernel.__call__: Calling kernel with opcode={opcode}, n_insts={self.__n_insts}, {len(kernel_args)} buffer args"
-        )
-        try:
-            h = self.__kernel(
-                opcode, self.__insts_buffer_bo, self.__n_insts, *kernel_args
-            )
-            print(f"[DEBUG] NPUKernel.__call__: Kernel execution started")
-            r = h.wait()
-            print(
-                f"[DEBUG] NPUKernel.__call__: Kernel execution completed with result: {r}"
-            )
-            if r != xrt.ert_cmd_state.ERT_CMD_STATE_COMPLETED:
-                print(f"[DEBUG] NPUKernel.__call__: Kernel failed with state: {r}")
-                raise NPUKernel_Error(f"Kernel returned {r}")
-        except Exception as e:
-            print(f"[DEBUG] NPUKernel.__call__: Exception during kernel execution: {e}")
-            print(f"[DEBUG] NPUKernel.__call__: Exception traceback:")
-            traceback.print_exc()
-            raise
+        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:
+            raise NPUKernel_Error(f"Kernel returned {r}")
 
     def __del__(self):
         """
         Destructor to clean up resources and delete the kernel and device objects.
         """
-        print(f"[DEBUG] NPUKernel.__del__: Cleaning up resources")
         try:
             del self.__kernel
             del self.__device
-        except Exception as e:
-            print(f"[DEBUG] NPUKernel.__del__: Exception during cleanup: {e}")
+        except Exception:
+            pass
 
 
 class NPUKernel_Error(Exception):
@@ -232,55 +145,33 @@ def jit(function=None, is_placed=True, use_cache=True):
 
     @functools.wraps(function)
     def decorator(*args, **kwargs):
-        print(
-            f"[DEBUG] jit.decorator: Starting compilation for function {function.__name__}"
-        )
-        print(f"[DEBUG] jit.decorator: is_placed={is_placed}, use_cache={use_cache}")
-        print(
-            f"[DEBUG] jit.decorator: args count={len(args)}, kwargs count={len(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()
-        print(f"[DEBUG] jit.decorator: Cleared ExternalKernel instances")
 
         # Find ExternalKernel instances in arguments and kwargs
         external_kernels = []
         for arg in args:
             if isinstance(arg, ExternalKernel):
                 external_kernels.append(arg)
-                print(
-                    f"[DEBUG] jit.decorator: Found ExternalKernel in args: {arg._name}"
-                )
         for value in kwargs.values():
             if isinstance(value, ExternalKernel):
                 external_kernels.append(value)
-                print(
-                    f"[DEBUG] jit.decorator: Found ExternalKernel in kwargs: {value._name}"
-                )
 
         # Execute the function to generate MLIR
         try:
             if is_placed:
-                print(f"[DEBUG] jit.decorator: Generating MLIR with placement")
                 with mlir_mod_ctx() as ctx:
                     function(*args, **kwargs)
-                    print(f"[DEBUG] jit.decorator: Function executed successfully")
                     assert (
                         ctx.module.operation.verify()
                     ), f"Verification failed for '{function.__name__}'"
-                    print(f"[DEBUG] jit.decorator: Module verification passed")
                     mlir_module = ctx.module
             else:
-                print(f"[DEBUG] jit.decorator: Generating MLIR without placement")
                 mlir_module = function(*args, **kwargs)
-                print(f"[DEBUG] jit.decorator: Function returned MLIR module")
         except Exception as e:
-            print(f"[DEBUG] jit.decorator: Exception during MLIR generation: {e}")
-            traceback.print_exc()
             raise
 
         # Compile all ExternalKernel instances that were created during this JIT compilation
@@ -289,15 +180,10 @@ def decorator(*args, **kwargs):
                 not hasattr(func, "_compiled") or not func._compiled
             ):  # Don't compile if already compiled
                 external_kernels.append(func)
-                print(
-                    f"[DEBUG] jit.decorator: Added ExternalKernel to compilation list: {func._name}"
-                )
 
         # Determine target architecture based on device type (hoisted from compile_external_kernel)
         try:
             current_device = get_current_device()
-            print(f"[DEBUG] jit.decorator: Current device: {current_device}")
-            print(f"[DEBUG] jit.decorator: Device type: {type(current_device)}")
 
             # Determine target architecture based on device type
             if isinstance(current_device, NPU2):
@@ -305,28 +191,17 @@ def decorator(*args, **kwargs):
             elif isinstance(current_device, NPU1):
                 target_arch = "aie2"
             else:
-                print(
-                    f"[DEBUG] jit.decorator: Unsupported device type: {type(current_device)}"
-                )
                 raise RuntimeError(f"Unsupported device type: {type(current_device)}")
-
-            print(f"[DEBUG] jit.decorator: Target architecture: {target_arch}")
         except Exception as e:
-            print(
-                f"[DEBUG] jit.decorator: Failed to determine target architecture: {e}"
-            )
             raise
 
         # Hash of the IR string, ExternalKernel 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")
-        print(f"[DEBUG] jit.decorator: Module hash: {module_hash}")
-        print(f"[DEBUG] jit.decorator: Kernel directory: {kernel_dir}")
 
         # Ensure cache directory exists
         os.makedirs(kernel_dir, exist_ok=True)
-        print(f"[DEBUG] jit.decorator: Created/verified kernel directory")
 
         # Write MLIR to file if not already cached
         inst_filename = "insts.bin"
@@ -336,27 +211,17 @@ def decorator(*args, **kwargs):
 
         xclbin_exists = os.path.exists(xclbin_path)
         inst_exists = os.path.exists(inst_path)
-        print(f"[DEBUG] jit.decorator: xclbin exists: {xclbin_exists}")
-        print(f"[DEBUG] jit.decorator: inst file exists: {inst_exists}")
 
         if not use_cache or not xclbin_exists or not inst_exists:
-            print(
-                f"[DEBUG] jit.decorator: Need to compile (cache disabled or files missing)"
-            )
             try:
                 with open(mlir_path, "w", encoding="utf-8") as f:
                     print(mlir_module, file=f)
-                print(f"[DEBUG] jit.decorator: Wrote MLIR module to {mlir_path}")
 
                 # Set cache directory for ExternalKernels and compile them
                 for func in external_kernels:
-                    print(
-                        f"[DEBUG] jit.decorator: Compiling ExternalKernel: {func._name}"
-                    )
                     # Compile the ExternalKernel directly in the kernel directory
                     compile_external_kernel(func, kernel_dir, target_arch)
 
-                print(f"[DEBUG] jit.decorator: Starting MLIR module compilation")
                 # Compile the MLIR module
                 compile_mlir_module_to_binary(
                     mlir_module=mlir_module,
@@ -364,35 +229,18 @@ def decorator(*args, **kwargs):
                     xclbin_path=xclbin_path,
                     work_dir=kernel_dir,
                 )
-                print(f"[DEBUG] jit.decorator: MLIR module compilation completed")
             except Exception as e:
-                print(f"[DEBUG] jit.decorator: Exception during compilation: {e}")
-                traceback.print_exc()
                 # Clean up cache directory on any compilation failure
                 if os.path.exists(kernel_dir):
                     shutil.rmtree(kernel_dir)
-                    print(
-                        f"[DEBUG] jit.decorator: Cleaned up kernel directory after failure"
-                    )
                 raise e
-        else:
-            print(f"[DEBUG] jit.decorator: Using cached compilation results")
 
         kernel_name = "MLIR_AIE"
-        print(
-            f"[DEBUG] jit.decorator: Creating NPUKernel with kernel_name={kernel_name}"
-        )
         try:
             kernel = NPUKernel(xclbin_path, inst_path, kernel_name=kernel_name)
-            print(f"[DEBUG] jit.decorator: NPUKernel created successfully")
             result = kernel(*args, **kwargs)
-            print(f"[DEBUG] jit.decorator: Kernel execution completed successfully")
             return result
         except Exception as e:
-            print(
-                f"[DEBUG] jit.decorator: Exception during kernel creation/execution: {e}"
-            )
-            traceback.print_exc()
             raise
 
     return decorator
@@ -407,26 +255,17 @@ def compile_external_kernel(func, kernel_dir, target_arch):
         kernel_dir: Directory to place the compiled object file
         target_arch: Target architecture (e.g., "aie2" or "aie2p")
     """
-    print(f"[DEBUG] compile_external_kernel: Starting compilation of {func._name}")
-    print(f"[DEBUG] compile_external_kernel: kernel_dir={kernel_dir}")
-    print(f"[DEBUG] compile_external_kernel: target_arch={target_arch}")
-
     # Skip if already compiled
     if hasattr(func, "_compiled") and func._compiled:
-        print(f"[DEBUG] compile_external_kernel: Function already compiled, skipping")
         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):
-        print(
-            f"[DEBUG] compile_external_kernel: Object file already exists: {output_file}"
-        )
         return
 
     # Create source file in kernel directory
     source_file = os.path.join(kernel_dir, f"{func._name}.cc")
-    print(f"[DEBUG] compile_external_kernel: Source file path: {source_file}")
 
     # Handle both source_string and source_file cases
     if func._source_string is not None:
@@ -434,54 +273,24 @@ def compile_external_kernel(func, kernel_dir, target_arch):
         try:
             with open(source_file, "w") as f:
                 f.write(func._source_string)
-            print(
-                f"[DEBUG] compile_external_kernel: Wrote source from string to {source_file}"
-            )
-            print(
-                f"[DEBUG] compile_external_kernel: Source content:\n{func._source_string}"
-            )
         except Exception as e:
-            print(f"[DEBUG] compile_external_kernel: Failed to write source file: {e}")
             raise
     elif func._source_file is not None:
         # Use source_file (copy existing file)
-        print(
-            f"[DEBUG] compile_external_kernel: Copying source file: {func._source_file}"
-        )
 
         # Check if source file exists before copying
         if os.path.exists(func._source_file):
             try:
                 shutil.copy2(func._source_file, source_file)
-                print(
-                    f"[DEBUG] compile_external_kernel: Copied source file successfully"
-                )
             except Exception as e:
-                print(
-                    f"[DEBUG] compile_external_kernel: Failed to copy source file: {e}"
-                )
                 raise
         else:
-            print(
-                f"[DEBUG] compile_external_kernel: Source file does not exist: {func._source_file}"
-            )
             return
     else:
-        print(
-            f"[DEBUG] compile_external_kernel: Neither source_string nor source_file provided"
-        )
         raise ValueError("Neither source_string nor source_file is provided")
 
     from .compile.compile import compile_cxx_core_function
 
-    print(f"[DEBUG] compile_external_kernel: Calling compile_cxx_core_function")
-    print(f"[DEBUG] compile_external_kernel: source_path={source_file}")
-    print(f"[DEBUG] compile_external_kernel: target_arch={target_arch}")
-    print(f"[DEBUG] compile_external_kernel: output_path={output_file}")
-    print(f"[DEBUG] compile_external_kernel: include_dirs={func._include_dirs}")
-    print(f"[DEBUG] compile_external_kernel: compile_args={func._compile_flags}")
-    print(f"[DEBUG] compile_external_kernel: cwd={kernel_dir}")
-
     try:
         compile_cxx_core_function(
             source_path=source_file,
@@ -490,41 +299,28 @@ def compile_external_kernel(func, kernel_dir, target_arch):
             include_dirs=func._include_dirs,
             compile_args=func._compile_flags,
             cwd=kernel_dir,
-            verbose=True,
+            verbose=False,
         )
-        print(f"[DEBUG] compile_external_kernel: Compilation completed successfully")
     except Exception as e:
-        print(f"[DEBUG] compile_external_kernel: Compilation failed: {e}")
-        traceback.print_exc()
         raise
 
     # Mark the function as compiled
     func._compiled = True
-    print(f"[DEBUG] compile_external_kernel: Marked function as compiled")
 
 
 def hash_module(module, external_kernels=None, target_arch=None):
     """
     Hash the MLIR module and ExternalKernel compiler options to create a unique identifier.
     """
-    print(f"[DEBUG] hash_module: Starting hash computation")
     mlir_str = str(module)
-    print(f"[DEBUG] hash_module: MLIR string length: {len(mlir_str)}")
-    print(f"[DEBUG] hash_module: Target architecture: {target_arch}")
 
     # Include ExternalKernel compiler options in the hash
     if external_kernels:
-        print(
-            f"[DEBUG] hash_module: Processing {len(external_kernels)} 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)
-            print(
-                f"[DEBUG] hash_module: Added options for {func._name}: include_dirs={func._include_dirs}, compile_flags={func._compile_flags}"
-            )
 
         # Create a combined string for hashing
         combined_str = mlir_str + "|" + "|".join(compiler_options)
@@ -534,8 +330,6 @@ def hash_module(module, external_kernels=None, target_arch=None):
     # Include target architecture in the hash
     if target_arch:
         combined_str += f"|target_arch={target_arch}"
-        print(f"[DEBUG] hash_module: Added target architecture to hash: {target_arch}")
 
     hash_result = hashlib.sha256(combined_str.encode("utf-8")).hexdigest()[:16]
-    print(f"[DEBUG] hash_module: Computed hash: {hash_result}")
     return hash_result
diff --git a/test/python/jit_extern_functions.py b/test/python/jit_extern_functions.py
index 0fd766c2f53..4baf329d5b1 100644
--- a/test/python/jit_extern_functions.py
+++ b/test/python/jit_extern_functions.py
@@ -844,7 +844,3 @@ def test_compiler_flag_combinations(compile_flags, expected_value):
     expected = initial_tensor + expected_value
     actual = output_tensor.numpy()
     np.testing.assert_array_equal(actual, expected)
-
-
-if __name__ == "__main__":
-    test_simple_add_one()

From 85b864e23f8c6a9e2c0df50e335fe174cbe0b7bf Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Wed, 6 Aug 2025 00:45:41 -0700
Subject: [PATCH 34/37] Remove debug code and handke NPU2Col1 and NPU1Col1

---
 python/iron/jit.py | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/python/iron/jit.py b/python/iron/jit.py
index e8cccc9833e..3b44362c62a 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -18,7 +18,7 @@
 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 .device import NPU1, NPU2, NPU1Col1, NPU2Col1
 from .config import get_current_device
 
 
@@ -53,7 +53,7 @@ def __init__(
 
         try:
             xkernel = [k for k in kernels if kernel_name == k.get_name()][0]
-        except (KeyError, IndexError) as e:
+        except KeyError:
             raise NPUKernel_Error("No such kernel: " + kernel_name)
 
         self.__device.register_xclbin(self.__xclbin)
@@ -116,11 +116,8 @@ def __del__(self):
         """
         Destructor to clean up resources and delete the kernel and device objects.
         """
-        try:
-            del self.__kernel
-            del self.__device
-        except Exception:
-            pass
+        del self.__kernel
+        del self.__device
 
 
 class NPUKernel_Error(Exception):
@@ -186,9 +183,9 @@ def decorator(*args, **kwargs):
             current_device = get_current_device()
 
             # Determine target architecture based on device type
-            if isinstance(current_device, NPU2):
+            if isinstance(current_device, (NPU2, NPU2Col1)):
                 target_arch = "aie2p"
-            elif isinstance(current_device, NPU1):
+            elif isinstance(current_device, (NPU1, NPU1Col1)):
                 target_arch = "aie2"
             else:
                 raise RuntimeError(f"Unsupported device type: {type(current_device)}")

From 4a225676a18ba3f782790419d0090341981cb505 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Wed, 6 Aug 2025 01:03:02 -0700
Subject: [PATCH 35/37] handle other device names

---
 python/iron/jit.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/python/iron/jit.py b/python/iron/jit.py
index 3b44362c62a..f2695e3f7d6 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -20,6 +20,7 @@
 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.
@@ -187,6 +188,10 @@ def decorator(*args, **kwargs):
                 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:

From 750d705f8a28ccdea2d6ba880ac344f40b76853a Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Wed, 6 Aug 2025 01:41:46 -0700
Subject: [PATCH 36/37] Replace `ExternalKernel` with `ExternalFunction`

---
 python/iron/__init__.py                       |   2 +-
 python/iron/jit.py                            |  30 ++---
 python/iron/kernel.py                         |  10 +-
 python/iron/worker.py                         |   4 +-
 test/python/jit_extern_functions.py           | 106 +++++++++---------
 ....py => jit_extern_functions_inside_jit.py} |  44 ++++----
 6 files changed, 98 insertions(+), 98 deletions(-)
 rename test/python/{jit_extern_inside_jit.py => jit_extern_functions_inside_jit.py} (87%)

diff --git a/python/iron/__init__.py b/python/iron/__init__.py
index 3c8be7f8a0d..e28c6880819 100644
--- a/python/iron/__init__.py
+++ b/python/iron/__init__.py
@@ -1,5 +1,5 @@
 from .globalbuffer import GlobalBuffer
-from .kernel import ExternalKernel, Kernel
+from .kernel import ExternalFunction, Kernel
 from .localbuffer import LocalBuffer
 from .program import Program
 from .worker import Worker, WorkerRuntimeBarrier
diff --git a/python/iron/jit.py b/python/iron/jit.py
index f2695e3f7d6..2f345affe15 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -143,19 +143,19 @@ 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
+        # Import ExternalFunction at the top
+        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
-        ExternalKernel._instances.clear()
+        ExternalFunction._instances.clear()
 
-        # Find ExternalKernel instances in arguments and kwargs
+        # Find ExternalFunction instances in arguments and kwargs
         external_kernels = []
         for arg in args:
-            if isinstance(arg, ExternalKernel):
+            if isinstance(arg, ExternalFunction):
                 external_kernels.append(arg)
         for value in kwargs.values():
-            if isinstance(value, ExternalKernel):
+            if isinstance(value, ExternalFunction):
                 external_kernels.append(value)
 
         # Execute the function to generate MLIR
@@ -172,8 +172,8 @@ def decorator(*args, **kwargs):
         except Exception as e:
             raise
 
-        # Compile all ExternalKernel instances that were created during this JIT compilation
-        for func in ExternalKernel._instances:
+        # 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
@@ -197,7 +197,7 @@ def decorator(*args, **kwargs):
         except Exception as e:
             raise
 
-        # Hash of the IR string, ExternalKernel compiler options, and target architecture
+        # 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")
@@ -219,9 +219,9 @@ def decorator(*args, **kwargs):
                 with open(mlir_path, "w", encoding="utf-8") as f:
                     print(mlir_module, file=f)
 
-                # Set cache directory for ExternalKernels and compile them
+                # Set cache directory for ExternalFunctions and compile them
                 for func in external_kernels:
-                    # Compile the ExternalKernel directly in the kernel directory
+                    # Compile the ExternalFunction directly in the kernel directory
                     compile_external_kernel(func, kernel_dir, target_arch)
 
                 # Compile the MLIR module
@@ -250,10 +250,10 @@ def decorator(*args, **kwargs):
 
 def compile_external_kernel(func, kernel_dir, target_arch):
     """
-    Compile an ExternalKernel to an object file in the kernel directory.
+    Compile an ExternalFunction to an object file in the kernel directory.
 
     Args:
-        func: ExternalKernel instance to compile
+        func: ExternalFunction instance to compile
         kernel_dir: Directory to place the compiled object file
         target_arch: Target architecture (e.g., "aie2" or "aie2p")
     """
@@ -312,11 +312,11 @@ def compile_external_kernel(func, kernel_dir, target_arch):
 
 def hash_module(module, external_kernels=None, target_arch=None):
     """
-    Hash the MLIR module and ExternalKernel compiler options to create a unique identifier.
+    Hash the MLIR module and ExternalFunction compiler options to create a unique identifier.
     """
     mlir_str = str(module)
 
-    # Include ExternalKernel compiler options in the hash
+    # Include ExternalFunction compiler options in the hash
     if external_kernels:
         compiler_options = []
         for func in external_kernels:
diff --git a/python/iron/kernel.py b/python/iron/kernel.py
index 52c9d37a65c..961242b6108 100644
--- a/python/iron/kernel.py
+++ b/python/iron/kernel.py
@@ -76,7 +76,7 @@ def resolve(
             self._op = external_func(self._name, inputs=self._arg_types)
 
 
-class ExternalKernel(BaseKernel):
+class ExternalFunction(BaseKernel):
     _object_files = set()
     _instances = set()
 
@@ -89,7 +89,7 @@ def __init__(
         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.
+        """An ExternalFunction 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:
@@ -108,7 +108,7 @@ def __init__(
         self._compiled = False
 
         # Track this instance for JIT compilation
-        ExternalKernel._instances.add(self)
+        ExternalFunction._instances.add(self)
 
     def _setup_source(self, source_file: str | None, source_string: str | None) -> None:
         """Set up the source file for compilation."""
@@ -171,7 +171,7 @@ def tile_size(self, arg_index: int = 0) -> int:
         )
 
     def arg_types(self) -> list:
-        """Get the argument types of the ExternalKernel."""
+        """Get the argument types of the ExternalFunction."""
         return self._arg_types.copy()
 
     def resolve(
@@ -185,5 +185,5 @@ def resolve(
 
     def __call__(self, *args, **kwargs):
         if not self._op:
-            raise ValueError("Need to resolve ExternalKernel before it can be called")
+            raise ValueError("Need to resolve ExternalFunction before it can be called")
         call(self._op, args, **kwargs)
diff --git a/python/iron/worker.py b/python/iron/worker.py
index 31835866795..9791f65c72b 100644
--- a/python/iron/worker.py
+++ b/python/iron/worker.py
@@ -16,7 +16,7 @@
 from .device import PlacementTile, AnyComputeTile, Tile
 from .dataflow.objectfifo import ObjectFifoHandle, ObjectFifo
 from .dataflow.endpoint import ObjectFifoEndpoint
-from .kernel import Kernel, ExternalKernel
+from .kernel import Kernel, ExternalFunction
 from .globalbuffer import GlobalBuffer
 from .resolvable import Resolvable
 
@@ -86,7 +86,7 @@ def do_nothing_core_fun(*args) -> None:
 
         # Check arguments to the core. Some information is saved for resolution.
         for arg in self.fn_args:
-            if isinstance(arg, (Kernel, ExternalKernel)):
+            if isinstance(arg, (Kernel, ExternalFunction)):
                 bin_names.add(arg.bin_name)
             elif isinstance(arg, ObjectFifoHandle):
                 arg.endpoint = self
diff --git a/test/python/jit_extern_functions.py b/test/python/jit_extern_functions.py
index 4baf329d5b1..3a7d59919bd 100644
--- a/test/python/jit_extern_functions.py
+++ b/test/python/jit_extern_functions.py
@@ -14,7 +14,7 @@
 import pytest
 
 import aie.iron as iron
-from aie.iron import ExternalKernel, jit
+from aie.iron import ExternalFunction, jit
 from aie.iron import ObjectFifo, Worker, Runtime, Program
 from aie.iron.placers import SequentialPlacer
 from aie.iron.controlflow import range_
@@ -29,7 +29,7 @@ def transform(input, output, func):
         )
     num_elements = np.size(input)
 
-    # Extract tile size from ExternalKernel (using first argument)
+    # Extract tile size from ExternalFunction (using first argument)
     tile_size = func.tile_size(0)
 
     # Assert that input and output arrays have the same tile size
@@ -60,7 +60,7 @@ def transform(input, output, func):
 
     # Define a task that will run on a compute tile
     def core_body(of_in, of_out, func_to_apply):
-        # Extract tile size from ExternalKernel (using first argument)
+        # Extract tile size from ExternalFunction (using first argument)
         tile_size = func_to_apply.tile_size(0)
 
         # Number of sub-vector "tile" iterations
@@ -86,14 +86,14 @@ def core_body(of_in, of_out, func_to_apply):
 
 
 def test_simple_add_one():
-    """Test basic ExternalKernel with simple add_one operation."""
+    """Test basic ExternalFunction with simple add_one operation."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Create ExternalKernel for adding one
-    add_one = ExternalKernel(
+    # Create ExternalFunction for adding one
+    add_one = ExternalFunction(
         "add_one",
         source_string="""extern "C" {
             void add_one(int* input, int* output, int tile_size) {
@@ -126,15 +126,15 @@ def test_simple_add_one():
 
 @pytest.mark.parametrize("tile_size", [8, 16, 32, 64])
 def test_different_tile_sizes(tile_size):
-    """Test ExternalKernel with different tile sizes."""
+    """Test ExternalFunction with different tile sizes."""
     # Create input and output tensors
     num_elements = 1024
     input_tensor = iron.randint(0, 100, (num_elements,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((num_elements,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Create ExternalKernel with specific tile size
-    add_one = ExternalKernel(
+    # Create ExternalFunction with specific tile size
+    add_one = ExternalFunction(
         "add_one",
         source_string="""extern "C" {
             void add_one(int* input, int* output, int tile_size) {
@@ -167,14 +167,14 @@ def test_different_tile_sizes(tile_size):
     ],
 )
 def test_different_data_types(dtype, c_type):
-    """Test ExternalKernel with different data types."""
+    """Test ExternalFunction with different data types."""
     # Create input and output tensors
     input_tensor = iron.rand((1024,), dtype=dtype, device="npu")
     output_tensor = iron.zeros((1024,), dtype=dtype, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Create ExternalKernel with specific data type
-    add_one = ExternalKernel(
+    # Create ExternalFunction with specific data type
+    add_one = ExternalFunction(
         "add_one",
         source_string=f"""extern "C" {{
             void add_one({c_type}* input, {c_type}* output, int tile_size) {{
@@ -201,13 +201,13 @@ def test_different_data_types(dtype, c_type):
 
 @pytest.mark.parametrize("value", [5, 42])
 def test_define_values(value):
-    """Test ExternalKernel with different define values."""
+    """Test ExternalFunction with different define values."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    add_value = ExternalKernel(
+    add_value = ExternalFunction(
         "add_value",
         source_string="""extern "C" {
             void add_value(int* input, int* output, int tile_size) {
@@ -234,14 +234,14 @@ def test_define_values(value):
 
 
 def test_multiple_defines():
-    """Test ExternalKernel with multiple defines."""
+    """Test ExternalFunction with multiple defines."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Create ExternalKernel with multiple defines
-    complex_op = ExternalKernel(
+    # Create ExternalFunction with multiple defines
+    complex_op = ExternalFunction(
         "complex_op",
         source_string="""extern "C" {
             void complex_op(int* input, int* output, int tile_size) {
@@ -272,7 +272,7 @@ def test_multiple_defines():
 
 
 def test_include_directories():
-    """Test ExternalKernel with include directories."""
+    """Test ExternalFunction with include directories."""
     # Create a temporary directory with a header file
     with tempfile.TemporaryDirectory() as temp_dir:
         # Create a header file
@@ -294,8 +294,8 @@ def test_include_directories():
         output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
         initial_tensor = input_tensor.numpy().copy()
 
-        # Create ExternalKernel that includes the header
-        add_value = ExternalKernel(
+        # Create ExternalFunction that includes the header
+        add_value = ExternalFunction(
             "add_value",
             source_string="""extern "C" {
                 #include "math_ops.h"
@@ -323,7 +323,7 @@ def test_include_directories():
 
 
 def test_multiple_include_directories():
-    """Test ExternalKernel with multiple include directories."""
+    """Test ExternalFunction with multiple include directories."""
     # Create temporary directories with header files
     with tempfile.TemporaryDirectory() as temp_dir1, tempfile.TemporaryDirectory() as temp_dir2:
         # Create header files
@@ -340,8 +340,8 @@ def test_multiple_include_directories():
         output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
         initial_tensor = input_tensor.numpy().copy()
 
-        # Create ExternalKernel that includes both headers
-        add_values = ExternalKernel(
+        # Create ExternalFunction that includes both headers
+        add_values = ExternalFunction(
             "add_values",
             source_string="""extern "C" {
                 #include "ops1.h"
@@ -376,8 +376,8 @@ def test_caching_same_source():
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Create two ExternalKernels with identical source
-    add_one_1 = ExternalKernel(
+    # Create two ExternalFunctions with identical source
+    add_one_1 = ExternalFunction(
         "add_one_1",
         source_string="""extern "C" {
             void add_one_1(int* input, int* output, int tile_size) {
@@ -393,7 +393,7 @@ def test_caching_same_source():
         ],
     )
 
-    add_one_2 = ExternalKernel(
+    add_one_2 = ExternalFunction(
         "add_one_2",
         source_string="""extern "C" {
             void add_one_2(int* input, int* output, int tile_size) {
@@ -422,14 +422,14 @@ def test_caching_same_source():
 
 
 def test_context_manager():
-    """Test ExternalKernel with context manager syntax."""
+    """Test ExternalFunction with context manager syntax."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Create ExternalKernel and use it with context manager
-    with ExternalKernel(
+    # Create ExternalFunction and use it with context manager
+    with ExternalFunction(
         "add_one_context",
         source_string="""extern "C" {
             void add_one_context(int* input, int* output, int tile_size) {
@@ -454,14 +454,14 @@ def test_context_manager():
 
 
 def test_context_manager_with_compiler_options():
-    """Test ExternalKernel with context manager and compiler options."""
+    """Test ExternalFunction with context manager and compiler options."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Create ExternalKernel with compiler options using context manager
-    with ExternalKernel(
+    # Create ExternalFunction with compiler options using context manager
+    with ExternalFunction(
         "add_value_context",
         source_string="""extern "C" {
             void add_value_context(int* input, int* output, int tile_size) {
@@ -487,7 +487,7 @@ def test_context_manager_with_compiler_options():
 
 
 def test_source_file():
-    """Test ExternalKernel with source_file instead of source_string."""
+    """Test ExternalFunction with source_file instead of source_string."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
@@ -506,8 +506,8 @@ def test_source_file():
         source_file_path = f.name
 
     try:
-        # Create ExternalKernel using source_file
-        add_one_from_file = ExternalKernel(
+        # Create ExternalFunction using source_file
+        add_one_from_file = ExternalFunction(
             "add_one_from_file",
             source_file=source_file_path,
             arg_types=[
@@ -531,7 +531,7 @@ def test_source_file():
 
 
 def test_source_file_with_compiler_options():
-    """Test ExternalKernel with source_file and compiler options."""
+    """Test ExternalFunction with source_file and compiler options."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
@@ -550,8 +550,8 @@ def test_source_file_with_compiler_options():
         source_file_path = f.name
 
     try:
-        # Create ExternalKernel using source_file with compiler options
-        add_value_from_file = ExternalKernel(
+        # Create ExternalFunction using source_file with compiler options
+        add_value_from_file = ExternalFunction(
             "add_value_from_file",
             source_file=source_file_path,
             arg_types=[
@@ -576,14 +576,14 @@ def test_source_file_with_compiler_options():
 
 
 def test_transform_with_internal_func():
-    """Test transform function that creates ExternalKernel internally."""
+    """Test transform function that creates ExternalFunction internally."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Create ExternalKernel dynamically but pass it as argument
-    internal_func = ExternalKernel(
+    # Create ExternalFunction dynamically but pass it as argument
+    internal_func = ExternalFunction(
         "internal_add_one",
         source_string="""extern "C" {
             void internal_add_one(int* input, int* output, int tile_size) {
@@ -599,7 +599,7 @@ def test_transform_with_internal_func():
         ],
     )
 
-    # Apply the transform (ExternalKernel is passed as argument)
+    # Apply the transform (ExternalFunction is passed as argument)
     transform(input_tensor, output_tensor, internal_func)
 
     # Verify results
@@ -615,8 +615,8 @@ def test_caching_different_flags():
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Create ExternalKernels with same source but different flags
-    add_value_5 = ExternalKernel(
+    # Create ExternalFunctions with same source but different flags
+    add_value_5 = ExternalFunction(
         "add_value",
         source_string="""extern "C" {
             void add_value(int* input, int* output, int tile_size) {
@@ -633,7 +633,7 @@ def test_caching_different_flags():
         compile_flags=["-DADD_VALUE=5"],
     )
 
-    add_value_10 = ExternalKernel(
+    add_value_10 = ExternalFunction(
         "add_value",
         source_string="""extern "C" {
             void add_value(int* input, int* output, int tile_size) {
@@ -703,8 +703,8 @@ def test_invalid_source(invalid_source):
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
 
-    # Create ExternalKernel with invalid C++ source
-    invalid_func = ExternalKernel(
+    # Create ExternalFunction with invalid C++ source
+    invalid_func = ExternalFunction(
         "invalid_func",
         source_string=invalid_source,
         arg_types=[
@@ -733,8 +733,8 @@ def test_mismatched_tile_sizes(input_tile_size, output_tile_size):
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
 
-    # Create ExternalKernel with mismatched tile sizes
-    mismatched_func = ExternalKernel(
+    # Create ExternalFunction with mismatched tile sizes
+    mismatched_func = ExternalFunction(
         "mismatched_func",
         source_string="""extern "C" {
             void mismatched_func(int* input, int* output, int tile_size) {
@@ -769,8 +769,8 @@ def test_invalid_include_directory(invalid_include):
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
 
-    # Create ExternalKernel with invalid include directory
-    invalid_include_func = ExternalKernel(
+    # Create ExternalFunction with invalid include directory
+    invalid_include_func = ExternalFunction(
         "invalid_include_func",
         source_string="""extern "C" {
             #include "nonexistent.h"
@@ -803,7 +803,7 @@ def test_invalid_include_directory(invalid_include):
     ],
 )
 def test_compiler_flag_combinations(compile_flags, expected_value):
-    """Test ExternalKernel with different combinations of compiler flags."""
+    """Test ExternalFunction with different combinations of compiler flags."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
@@ -826,7 +826,7 @@ def test_compiler_flag_combinations(compile_flags, expected_value):
         }
     }"""
 
-    complex_op = ExternalKernel(
+    complex_op = ExternalFunction(
         "complex_op",
         source_string=source_template,
         arg_types=[
diff --git a/test/python/jit_extern_inside_jit.py b/test/python/jit_extern_functions_inside_jit.py
similarity index 87%
rename from test/python/jit_extern_inside_jit.py
rename to test/python/jit_extern_functions_inside_jit.py
index 709569de2dd..0c85d87a293 100644
--- a/test/python/jit_extern_inside_jit.py
+++ b/test/python/jit_extern_functions_inside_jit.py
@@ -14,7 +14,7 @@
 import pytest
 
 import aie.iron as iron
-from aie.iron import ExternalKernel, jit
+from aie.iron import ExternalFunction, jit
 from aie.iron import ObjectFifo, Worker, Runtime, Program
 from aie.iron.placers import SequentialPlacer
 from aie.iron.controlflow import range_
@@ -22,15 +22,15 @@
 
 @jit(is_placed=False)
 def transform_with_internal_func_with_options(input, output):
-    """Transform kernel that creates ExternalKernel internally with compiler options."""
+    """Transform kernel that creates ExternalFunction internally with compiler options."""
     if input.shape != output.shape:
         raise ValueError(
             f"Input shapes are not the equal ({input.shape} != {output.shape})."
         )
     num_elements = np.size(input)
 
-    # Create ExternalKernel inside the transform with compiler options
-    internal_func = ExternalKernel(
+    # Create ExternalFunction inside the transform with compiler options
+    internal_func = ExternalFunction(
         "internal_add_value",
         source_string="""extern "C" {
             void internal_add_value(int* input, int* output, int tile_size) {
@@ -47,7 +47,7 @@ def transform_with_internal_func_with_options(input, output):
         compile_flags=["-DADD_VALUE=1"],
     )
 
-    # Extract tile size from ExternalKernel
+    # Extract tile size from ExternalFunction
     tile_size = internal_func.tile_size(0)
 
     if num_elements % tile_size != 0:
@@ -73,7 +73,7 @@ def transform_with_internal_func_with_options(input, output):
 
     # Define a task that will run on a compute tile
     def core_body(of_in, of_out, func_to_apply):
-        # Extract tile size from ExternalKernel
+        # Extract tile size from ExternalFunction
         tile_size = func_to_apply.tile_size(0)
 
         # Number of sub-vector "tile" iterations
@@ -100,7 +100,7 @@ def core_body(of_in, of_out, func_to_apply):
 
 @jit(is_placed=False)
 def transform_with_internal_func_from_file(input, output):
-    """Transform kernel that creates ExternalKernel internally from a file."""
+    """Transform kernel that creates ExternalFunction internally from a file."""
     if input.shape != output.shape:
         raise ValueError(
             f"Input shapes are not the equal ({input.shape} != {output.shape})."
@@ -120,8 +120,8 @@ def transform_with_internal_func_from_file(input, output):
         )
         temp_file_path = f.name
 
-    # Create ExternalKernel inside the transform from a file
-    internal_func = ExternalKernel(
+    # Create ExternalFunction inside the transform from a file
+    internal_func = ExternalFunction(
         "internal_add_from_file",
         source_file=temp_file_path,
         arg_types=[
@@ -131,7 +131,7 @@ def transform_with_internal_func_from_file(input, output):
         ],
     )
 
-    # Extract tile size from ExternalKernel
+    # Extract tile size from ExternalFunction
     tile_size = internal_func.tile_size(0)
 
     if num_elements % tile_size != 0:
@@ -157,7 +157,7 @@ def transform_with_internal_func_from_file(input, output):
 
     # Define a task that will run on a compute tile
     def core_body(of_in, of_out, func_to_apply):
-        # Extract tile size from ExternalKernel
+        # Extract tile size from ExternalFunction
         tile_size = func_to_apply.tile_size(0)
 
         # Number of sub-vector "tile" iterations
@@ -184,15 +184,15 @@ def core_body(of_in, of_out, func_to_apply):
 
 @jit(is_placed=False)
 def transform_with_internal_func(input, output):
-    """Transform kernel that creates ExternalKernel internally."""
+    """Transform kernel that creates ExternalFunction internally."""
     if input.shape != output.shape:
         raise ValueError(
             f"Input shapes are not the equal ({input.shape} != {output.shape})."
         )
     num_elements = np.size(input)
 
-    # Create ExternalKernel inside the transform
-    internal_func = ExternalKernel(
+    # Create ExternalFunction inside the transform
+    internal_func = ExternalFunction(
         "internal_add_one",
         source_string="""extern "C" {
             void internal_add_one(int* input, int* output, int tile_size) {
@@ -208,7 +208,7 @@ def transform_with_internal_func(input, output):
         ],
     )
 
-    # Extract tile size from ExternalKernel
+    # Extract tile size from ExternalFunction
     tile_size = internal_func.tile_size(0)
 
     if num_elements % tile_size != 0:
@@ -234,7 +234,7 @@ def transform_with_internal_func(input, output):
 
     # Define a task that will run on a compute tile
     def core_body(of_in, of_out, func_to_apply):
-        # Extract tile size from ExternalKernel
+        # Extract tile size from ExternalFunction
         tile_size = func_to_apply.tile_size(0)
 
         # Number of sub-vector "tile" iterations
@@ -260,13 +260,13 @@ def core_body(of_in, of_out, func_to_apply):
 
 
 def test_transform_with_internal_func_with_options_inside():
-    """Test transform function that creates ExternalKernel internally with compiler options."""
+    """Test transform function that creates ExternalFunction internally with compiler options."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Apply the transform (ExternalKernel is created inside with hardcoded compiler options)
+    # Apply the transform (ExternalFunction is created inside with hardcoded compiler options)
     transform_with_internal_func_with_options(input_tensor, output_tensor)
 
     # Verify results
@@ -276,13 +276,13 @@ def test_transform_with_internal_func_with_options_inside():
 
 
 def test_transform_with_internal_func_inside():
-    """Test transform function that creates ExternalKernel internally."""
+    """Test transform function that creates ExternalFunction internally."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Apply the transform (ExternalKernel is created inside)
+    # Apply the transform (ExternalFunction is created inside)
     transform_with_internal_func(input_tensor, output_tensor)
 
     # Verify results
@@ -292,13 +292,13 @@ def test_transform_with_internal_func_inside():
 
 
 def test_transform_with_internal_func_from_file():
-    """Test transform function that creates ExternalKernel from a file."""
+    """Test transform function that creates ExternalFunction from a file."""
     # Create input and output tensors
     input_tensor = iron.randint(0, 100, (1024,), dtype=np.int32, device="npu")
     output_tensor = iron.zeros((1024,), dtype=np.int32, device="npu")
     initial_tensor = input_tensor.numpy().copy()
 
-    # Apply the transform (ExternalKernel is created inside from file)
+    # Apply the transform (ExternalFunction is created inside from file)
     transform_with_internal_func_from_file(input_tensor, output_tensor)
 
     # Verify results

From 2b4db11d0df30a9f23c6a24c9e8edbcb71305484 Mon Sep 17 00:00:00 2001
From: Muhammad Awad <MuhammadAbdelghaffar.Awad@amd.com>
Date: Wed, 6 Aug 2025 01:58:56 -0700
Subject: [PATCH 37/37] Improve comments

---
 python/iron/jit.py    | 10 +++-------
 python/iron/kernel.py |  6 +++---
 2 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/python/iron/jit.py b/python/iron/jit.py
index 2f345affe15..f82b69182d5 100644
--- a/python/iron/jit.py
+++ b/python/iron/jit.py
@@ -143,7 +143,6 @@ def jit(function=None, is_placed=True, use_cache=True):
 
     @functools.wraps(function)
     def decorator(*args, **kwargs):
-        # Import ExternalFunction at the top
         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
@@ -179,7 +178,7 @@ def decorator(*args, **kwargs):
             ):  # Don't compile if already compiled
                 external_kernels.append(func)
 
-        # Determine target architecture based on device type (hoisted from compile_external_kernel)
+        # Determine target architecture based on device type
         try:
             current_device = get_current_device()
 
@@ -219,9 +218,8 @@ def decorator(*args, **kwargs):
                 with open(mlir_path, "w", encoding="utf-8") as f:
                     print(mlir_module, file=f)
 
-                # Set cache directory for ExternalFunctions and compile them
+                # Compile ExternalFunctions from inside the JIT compilation directory
                 for func in external_kernels:
-                    # Compile the ExternalFunction directly in the kernel directory
                     compile_external_kernel(func, kernel_dir, target_arch)
 
                 # Compile the MLIR module
@@ -232,7 +230,7 @@ def decorator(*args, **kwargs):
                     work_dir=kernel_dir,
                 )
             except Exception as e:
-                # Clean up cache directory on any compilation failure
+                # 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
@@ -279,7 +277,6 @@ def compile_external_kernel(func, kernel_dir, target_arch):
             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:
@@ -320,7 +317,6 @@ def hash_module(module, external_kernels=None, target_arch=None):
     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)
 
diff --git a/python/iron/kernel.py b/python/iron/kernel.py
index 961242b6108..85678ffa16f 100644
--- a/python/iron/kernel.py
+++ b/python/iron/kernel.py
@@ -89,13 +89,13 @@ def __init__(
         include_dirs: list[str] = [],
         compile_flags: list[str] = [],
     ) -> None:
-        """An ExternalFunction is a C++ source file that gets compiled to an object file and eventually resolves to a FuncOp.
+        """An ExternalFunction is a C/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
+            source_file (str): Path to the C/C++ source file
+            source_string (str): C/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 [].