Revert "Revert #2957 -- Remove is_placed flag from JIT decorator" (#3004)

hunhoffe · web-flow · commit a504a6293388 · 2026-03-30T18:58:52.000Z
diff --git a/programming_examples/algorithms/README.md b/programming_examples/algorithms/README.md
@@ -59,7 +59,7 @@ my_kernel = ExternalFunction(
 
 # Pass in full size tensors, algorithm will tile it to tile_size
 # tile_size must be passed as a keyword argument
-iron.jit(is_placed=False)(transform)(
+iron.jit(transform)(
     my_kernel, input, output, factor, tile_size=TILE_SIZE
 )
 ```
diff --git a/programming_examples/algorithms/for_each.py b/programming_examples/algorithms/for_each.py
@@ -35,7 +35,7 @@ def main():
     initial_tensor = tensor.numpy().copy()
 
     # JIT compile the algorithm
-    iron.jit(is_placed=False)(for_each)(lambda a: a + 1, tensor, tile_size=16)
+    iron.jit(for_each)(lambda a: a + 1, tensor, tile_size=16)
 
     # Check the correctness of the result
     e = np.equal(initial_tensor + 1, tensor.numpy())
diff --git a/programming_examples/algorithms/transform.py b/programming_examples/algorithms/transform.py
@@ -36,7 +36,7 @@ def main():
     output = iron.zeros_like(input)
 
     # JIT compile the algorithm
-    iron.jit(is_placed=False)(transform)(lambda a: a + 1, input, output, tile_size=16)
+    iron.jit(transform)(lambda a: a + 1, input, output, tile_size=16)
 
     # Check the correctness of the result
     e = np.equal(input.numpy() + 1, output.numpy())
diff --git a/programming_examples/algorithms/transform_binary.py b/programming_examples/algorithms/transform_binary.py
@@ -36,9 +36,7 @@ def main():
     output = iron.zeros_like(input0)
 
     # JIT compile the algorithm
-    iron.jit(is_placed=False)(transform_binary)(
-        lambda a, b: a + b, input0, input1, output, tile_size=16
-    )
+    iron.jit(transform_binary)(lambda a, b: a + b, input0, input1, output, tile_size=16)
 
     # Check the correctness of the result
     e = np.equal(input0.numpy() + input1.numpy(), output.numpy())
diff --git a/programming_examples/algorithms/transform_parallel.py b/programming_examples/algorithms/transform_parallel.py
@@ -36,9 +36,7 @@ def main():
     output = iron.zeros_like(input)
 
     # JIT compile the algorithm
-    iron.jit(is_placed=False)(transform_parallel)(
-        lambda a: a + 1, input, output, tile_size=16
-    )
+    iron.jit(transform_parallel)(lambda a: a + 1, input, output, tile_size=16)
 
     # Check the correctness of the result
     e = np.equal(input.numpy() + 1, output.numpy())
diff --git a/programming_examples/algorithms/transform_parallel_binary.py b/programming_examples/algorithms/transform_parallel_binary.py
@@ -36,7 +36,7 @@ def main():
     output = iron.zeros_like(input0)
 
     # JIT compile the algorithm
-    iron.jit(is_placed=False)(transform_parallel_binary)(
+    iron.jit(transform_parallel_binary)(
         lambda a, b: a + b, input0, input1, output, tile_size=16
     )
 
diff --git a/programming_examples/basic/vector_reduce_min/vector_reduce_min_jit.py b/programming_examples/basic/vector_reduce_min/vector_reduce_min_jit.py
@@ -17,7 +17,7 @@
 from aie.iron import ExternalFunction
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def my_reduce_min(input_tensor, output_tensor):
 
     num_elements = input_tensor.numel()
diff --git a/programming_examples/basic/vector_scalar_mul/vector_scalar_mul_jit.py b/programming_examples/basic/vector_scalar_mul/vector_scalar_mul_jit.py
@@ -61,9 +61,7 @@ def vector_scalar_mul(input, factor, output):
     # Pass scale kernel to the transform algorithm
     # tile_size is passed as keyword argument
     # iron.jit compiles and runs the program
-    iron.jit(is_placed=False)(transform)(
-        scale_kernel, input, output, factor, tile_size=tile_size
-    )
+    iron.jit(transform)(scale_kernel, input, output, factor, tile_size=tile_size)
 
 
 def main():
diff --git a/programming_examples/basic/vector_vector_add/vector_vector_add.py b/programming_examples/basic/vector_vector_add/vector_vector_add.py
@@ -17,7 +17,7 @@
 from aie.iron.controlflow import range_
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def vector_vector_add(input0, input1, output):
     if input0.shape != input1.shape:
         raise ValueError(
diff --git a/programming_examples/getting_started/00_memcpy/memcpy.py b/programming_examples/getting_started/00_memcpy/memcpy.py
@@ -27,9 +27,8 @@
 # JIT decorator for IRON
 # Decorator to compile an IRON kernel into a binary to run on the NPU.
 # Parameters:
-#     - is_placed (bool): Whether the kernel is using explicit or deferred placement API. Defaults to True.
 #     - use_cache (bool): Use cached MLIR module if available. Defaults to True.
-@iron.jit(is_placed=False)
+@iron.jit
 def my_memcpy(input0, output):
     # --------------------------------------------------------------------------
     # Configuration
diff --git a/programming_examples/getting_started/01_SAXPY/saxpy.py b/programming_examples/getting_started/01_SAXPY/saxpy.py
@@ -19,9 +19,8 @@
 # JIT decorator for IRON
 # Decorator to compile an IRON kernel into a binary to run on the NPU.
 # Parameters:
-#     - is_placed (bool): Whether the kernel is using explicit or deferred placement API. Defaults to True.
 #     - use_cache (bool): Use cached MLIR module if available. Defaults to True.
-@iron.jit(is_placed=False)
+@iron.jit
 def saxpy(input0, input1, output):
     N = input0.shape[0]  # Tensor size
     element_type = output.dtype
diff --git a/programming_examples/getting_started/02_vector_reduce_max/vector_reduce_max_1col.py b/programming_examples/getting_started/02_vector_reduce_max/vector_reduce_max_1col.py
@@ -22,9 +22,8 @@
 # JIT decorator for IRON
 # Decorator to compile an IRON kernel into a binary to run on the NPU.
 # Parameters:
-#     - is_placed (bool): Whether the kernel is using explicit or deferred placement API. Defaults to True.
 #     - use_cache (bool): Use cached MLIR module if available. Defaults to True.
-@iron.jit(is_placed=False)
+@iron.jit
 def vector_reduce_max(input0, output):
     element_type = output.dtype
 
diff --git a/programming_examples/getting_started/03_matrix_multiplication_single_core/matrix_multiplication_single_core.py b/programming_examples/getting_started/03_matrix_multiplication_single_core/matrix_multiplication_single_core.py
@@ -20,9 +20,8 @@
 # JIT decorator for IRON
 # Decorator to compile an IRON kernel into a binary to run on the NPU.
 # Parameters:
-#     - is_placed (bool): Whether the kernel is using explicit or deferred placement API. Defaults to True.
 #     - use_cache (bool): Use cached MLIR module if available. Defaults to True.
-@iron.jit(is_placed=False)
+@iron.jit
 def matrix_multiplication_single_core(input0, input1, output):
     # Problem size
     # - matrix0 shapes: (M, K)
diff --git a/programming_guide/mini_tutorial/aie2.py b/programming_guide/mini_tutorial/aie2.py
@@ -26,9 +26,8 @@
 # JIT decorator for IRON
 # Decorator to compile an IRON kernel into a binary to run on the NPU.
 # Parameters:
-#     - is_placed (bool): Whether the kernel is using explicit or deferred placement API. Defaults to True.
 #     - use_cache (bool): Use cached MLIR module if available. Defaults to True.
-@iron.jit(is_placed=False)
+@iron.jit
 def aie2p(input0, output):
     # Dataflow with ObjectFifos
     # ObjectFifos represent a dataflow connection between endpoints in the AIE array.
diff --git a/programming_guide/mini_tutorial/exercise_1/answer_2.py b/programming_guide/mini_tutorial/exercise_1/answer_2.py
@@ -16,7 +16,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_1(input0, output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_1/answer_3.py b/programming_guide/mini_tutorial/exercise_1/answer_3.py
@@ -15,7 +15,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_1(input0, output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_1/answer_4.py b/programming_guide/mini_tutorial/exercise_1/answer_4.py
@@ -16,7 +16,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_1(input0, output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_1/answer_5.py b/programming_guide/mini_tutorial/exercise_1/answer_5.py
@@ -16,7 +16,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_1(input0, output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_1/exercise_1.py b/programming_guide/mini_tutorial/exercise_1/exercise_1.py
@@ -16,7 +16,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_1(output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_2/answer_1.py b/programming_guide/mini_tutorial/exercise_2/answer_1.py
@@ -16,7 +16,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_2(input0, output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_2/answer_2.py b/programming_guide/mini_tutorial/exercise_2/answer_2.py
@@ -16,7 +16,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_2(input0, output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_2/exercise_2.py b/programming_guide/mini_tutorial/exercise_2/exercise_2.py
@@ -16,7 +16,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_2(input0, output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_3/answer_2.py b/programming_guide/mini_tutorial/exercise_3/answer_2.py
@@ -16,7 +16,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_3(input0, input1, output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_3/exercise_3.py b/programming_guide/mini_tutorial/exercise_3/exercise_3.py
@@ -16,7 +16,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_3(input0, output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_4/answer_2.py b/programming_guide/mini_tutorial/exercise_4/answer_2.py
@@ -23,7 +23,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_4(output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_4/exercise_4.py b/programming_guide/mini_tutorial/exercise_4/exercise_4.py
@@ -16,7 +16,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_4(output):
     data_size = output.numel()
     element_type = output.dtype
diff --git a/programming_guide/mini_tutorial/exercise_5/exercise_5a/answer_1.py b/programming_guide/mini_tutorial/exercise_5/exercise_5a/answer_1.py
@@ -17,7 +17,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_5a(input0, output):
     # Define tile size
     tile_height = 3
diff --git a/programming_guide/mini_tutorial/exercise_5/exercise_5a/answer_2.py b/programming_guide/mini_tutorial/exercise_5/exercise_5a/answer_2.py
@@ -17,7 +17,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_5a(input0, output):
     # Define tile size
     tile_height = 3
diff --git a/programming_guide/mini_tutorial/exercise_5/exercise_5a/answer_3.py b/programming_guide/mini_tutorial/exercise_5/exercise_5a/answer_3.py
@@ -16,7 +16,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_5a(input0, output):
     data_size = input0.numel()
     element_type = input0.dtype
diff --git a/programming_guide/mini_tutorial/exercise_5/exercise_5a/exercise_5a.py b/programming_guide/mini_tutorial/exercise_5/exercise_5a/exercise_5a.py
@@ -17,7 +17,7 @@
 import aie.iron as iron
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_5a(input0, output):
     # Define tile size
     tile_height = 3
diff --git a/programming_guide/mini_tutorial/exercise_5/exercise_5b/exercise_5b.py b/programming_guide/mini_tutorial/exercise_5/exercise_5b/exercise_5b.py
@@ -23,7 +23,7 @@
 data_width = 16
 
 
-@iron.jit(is_placed=False)
+@iron.jit
 def exercise_5b(input0, output):
     # Define tile size
     tile_height = 3
diff --git a/python/utils/jit.py b/python/utils/jit.py
@@ -27,20 +27,24 @@
 _compiled_kernels = CircularCache(max_size=1)
 
 
-def jit(function=None, is_placed=True, use_cache=True):
+def jit(function=None, use_cache=True):
     """
     Decorator to compile an NPU kernel into a binary to run on the NPU.
 
+    The decorated function may either return an MLIR module directly (unplaced
+    style, using the IRON API) or return None and populate the module implicitly
+    through the active ``mlir_mod_ctx`` context (placed style, using low-level
+    dialects).  The mode is detected automatically from the return value.
+
     Args:
         function (callable, optional): The function to compile.
-        is_placed (bool, optional): Whether the kernel is using explicit or implicit placement. Defaults to True.
         use_cache (bool, optional): Use cached MLIR module if available. Defaults to True.
 
     Returns:
         callable: The decorated function.
     """
     if function is None:
-        return functools.partial(jit, is_placed=is_placed, use_cache=use_cache)
+        return functools.partial(jit, use_cache=use_cache)
 
     @functools.wraps(function)
     def decorator(*args, **kwargs):
@@ -89,26 +93,21 @@ def decorator(*args, **kwargs):
         ExternalFunction._instances.clear()
 
         # Execute the function to generate MLIR.
-        # Placed designs use the raw @device(...) DSL and populate the context
-        # as a side effect (returning nothing), so we must provide an outer
-        # mlir_mod_ctx() and capture ctx.module.
-        # Non-placed designs use Program.resolve_program(), which opens its own
-        # mlir_mod_ctx() internally and returns the module directly; wrapping
-        # them in an outer context would give an empty module.
-        if is_placed:
-            with mlir_mod_ctx() as ctx:
-                function(*args, **kwargs)
-                if not ctx.module.operation.verify():
-                    raise RuntimeError(
-                        f"MLIR verification failed for '{function.__name__}'"
-                    )
-                mlir_module = ctx.module
+        # Always wrap in mlir_mod_ctx so that placed-style functions (which
+        # populate the module implicitly) work correctly.  If the function
+        # returns a module directly (unplaced style) we use that instead.
+        with mlir_mod_ctx() as ctx:
+            result = function(*args, **kwargs)
+
+        if result is None:
+            # Placed style: module was built implicitly via the context.
+            assert (
+                ctx.module.operation.verify()
+            ), f"Verification failed for '{function.__name__}'"
+            mlir_module = ctx.module
         else:
-            mlir_module = function(*args, **kwargs)
-            if not mlir_module.operation.verify():
-                raise RuntimeError(
-                    f"MLIR verification failed for '{function.__name__}'"
-                )
+            # Unplaced style: function returned the module directly.
+            mlir_module = result
 
         # Also collect ExternalFunction instances created during function()
         # execution (e.g. inside algorithm helpers that construct them internally).
diff --git a/test/python-concurrency/jit_parallel_compilation.py b/test/python-concurrency/jit_parallel_compilation.py
@@ -32,7 +32,7 @@ def test_parallel_compilation_subprocess():
 from aie.iron.placers import SequentialPlacer
 from aie.iron.controlflow import range_
 
-@iron.jit(is_placed=False)
+@iron.jit
 def simple_add(input0, input1, output):
     if input0.shape != input1.shape:
         raise ValueError(f"Input shapes are not equal ({input0.shape} != {input1.shape}).")
diff --git a/test/python/npu-xrt/test_algorithms.py b/test/python/npu-xrt/test_algorithms.py
diff --git a/test/python/npu-xrt/test_cached_xrt_runtime.py b/test/python/npu-xrt/test_cached_xrt_runtime.py
diff --git a/test/python/npu-xrt/test_cached_xrt_runtime_insts.py b/test/python/npu-xrt/test_cached_xrt_runtime_insts.py
diff --git a/test/python/npu-xrt/test_compile_cache_functionality.py b/test/python/npu-xrt/test_compile_cache_functionality.py
diff --git a/test/python/npu-xrt/test_jit_compilation.py b/test/python/npu-xrt/test_jit_compilation.py
diff --git a/test/python/npu-xrt/test_jit_extern_functions.py b/test/python/npu-xrt/test_jit_extern_functions.py
diff --git a/test/python/npu-xrt/test_jit_extern_functions_inside_jit.py b/test/python/npu-xrt/test_jit_extern_functions_inside_jit.py
diff --git a/test/python/npu-xrt/test_jit_placed_style.py b/test/python/npu-xrt/test_jit_placed_style.py
diff --git a/test/python/npu-xrt/test_jit_trace.py b/test/python/npu-xrt/test_jit_trace.py
diff --git a/test/python/npu-xrt/test_jit_two_extern_functions.py b/test/python/npu-xrt/test_jit_two_extern_functions.py
diff --git a/test/python/npu-xrt/test_jit_utils.py b/test/python/npu-xrt/test_jit_utils.py

Original file line number	Diff line number	Diff line change
`@@ -59,7 +59,7 @@ my_kernel = ExternalFunction(`
`59`	`59`
`60`	`60`	`# Pass in full size tensors, algorithm will tile it to tile_size`
`61`	`61`	`# tile_size must be passed as a keyword argument`
`62`		`-iron.jit(is_placed=False)(transform)(`
	`62`	`+iron.jit(transform)(`
`63`	`63`	`my_kernel, input, output, factor, tile_size=TILE_SIZE`
`64`	`64`	`)`
`65`	`65`	```
Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,7 @@ def main():`
`36`	`36`	`output = iron.zeros_like(input0)`
`37`	`37`
`38`	`38`	`# JIT compile the algorithm`
`39`		`- iron.jit(is_placed=False)(transform_parallel_binary)(`
	`39`	`+ iron.jit(transform_parallel_binary)(`
`40`	`40`	`lambda a, b: a + b, input0, input1, output, tile_size=16`
`41`	`41`	`)`
`42`	`42`