Update mlir-aie to v1.4.0 and migrate the operator library to its APIs

aie_kernels/aie2/rms_norm.cc

@@ -38,20 +38,20 @@ void rms_norm_general(const T *restrict input,
 
     float rms = sum_sq / cols + epsilon;
     float inv_rms = invsqrt(rms);
-    // Normalize in f32 and round once/ Mirrors layer_norm.cc's accfloat path.
-    ::aie::accum<accfloat, N> inv_rms_v;
-    inv_rms_v.from_vector(::aie::broadcast<float, N>(inv_rms), 0);
+    // Peano has no f32 vector multiply for AIE2, so the f32 scale rides in a
+    // bf16 pair and is applied as two exact products accumulated in f32. A
+    // single bf16 scale would shift every element of a norm the same way.
+    T inv_rms_hi = static_cast<T>(inv_rms);
+    T inv_rms_lo = static_cast<T>(inv_rms - static_cast<float>(inv_rms_hi));
 
     for (int i = 0; i < vector_chunks; i++) {
-        ::aie::accum<accfloat, N> reg_a;
-        reg_a.from_vector(::aie::load_v<N>(input + i * N), 0);
-        reg_a = ::aie::mul(reg_a.template to_vector<float>(), inv_rms_v.template to_vector<float>());
+        ::aie::vector<T, N> reg_a = ::aie::load_v<N>(input + i * N);
+        ::aie::accum<accfloat, N> acc = ::aie::mul(reg_a, inv_rms_hi);
+        acc = ::aie::mac(acc, reg_a, inv_rms_lo);
         if (input2) {
-            ::aie::accum<accfloat, N> reg_b;
-            reg_b.from_vector(::aie::load_v<N>(input2 + i * N), 0);
-            reg_a = ::aie::mul(reg_a.template to_vector<float>(), reg_b.template to_vector<float>());
+            acc = ::aie::mul(acc.template to_vector<T>(), ::aie::load_v<N>(input2 + i * N));
         }
-        ::aie::store_v(output + i * N, reg_a.template to_vector<T>());
+        ::aie::store_v(output + i * N, acc.template to_vector<T>());
     }
 
     if (remaining > 0) {

iron/common/compilation/base.py

@@ -518,7 +518,6 @@ def compile(self, graph):
                 str(self.aiecc_path),
                 "-v",
                 f"-j{os.environ.get('AIECC_JOBS', '1')}",
-                "--no-compile-host",
             ]
             if self.use_chess:
                 compile_cmd += [
@@ -534,9 +533,16 @@ def compile(self, graph):
                 ]
             compile_cmd += [
                 "--expand-load-pdis",
-                "--generate-full-elf",
+                "--get-full-elf",
                 "--full-elf-name",
                 os.path.abspath(artifact.filename),
+                # Unrequested, params.txt is never written and
+                # SequenceFullELFCallable.params silently returns None. It has
+                # no name flag of its own, so --output-dir is what puts it in
+                # the project directory that property reads.
+                "--get-scratchpad-parameters",
+                "--output-dir",
+                os.path.abspath(artifact.mlir_input.filename) + ".prj",
                 *artifact.extra_flags,
                 os.path.abspath(artifact.mlir_input.filename),
             ]
@@ -573,7 +579,6 @@ def compile(self, graph):
                 str(self.aiecc_path),
                 "-v",
                 f"-j{os.environ.get('AIECC_JOBS', '1')}",
-                "--no-compile-host",
             ]
             if self.use_chess:
                 compile_cmd += [
@@ -597,7 +602,7 @@ def compile(self, graph):
                     0
                 ]  # TODO: this does not handle the case of multiple xclbins with different kernel names or flags from the same MLIR
                 compile_cmd += first_xclbin.extra_flags + [
-                    "--aie-generate-xclbin",
+                    "--get-xclbin",
                     "--xclbin-name=" + os.path.abspath(first_xclbin.filename),
                     "--xclbin-kernel-name=" + first_xclbin.kernel_name,
                 ]
@@ -610,10 +615,10 @@ def compile(self, graph):
                 first_insts_bin = mlir_sources_to_insts[mlir_source][
                     0
                 ]  # TODO: this does not handle the case of multiple insts.bins with different flags from the same MLIR
-                if not do_compile_xclbin:
-                    compile_cmd += ["--no-compile"]
+                # Outputs are selected by --get-<name>; asking only for the insts is what
+                # "--no-compile" used to mean, so there is nothing to opt out of here.
                 compile_cmd += first_insts_bin.extra_flags + [
-                    "--aie-generate-npu-insts",
+                    "--get-npu-insts",
                     "--npu-insts-name=" + os.path.abspath(first_insts_bin.filename),
                 ]
             compile_cmd += [os.path.abspath(mlir_source.filename)]

iron/common/sequence.py

@@ -574,14 +574,17 @@ def params(self):
         The ``params.txt`` describing the runtime parameters is written by
         ``aie-lower-parameters`` into the ``<mlir>.prj`` project directory next
         to the fused MLIR source. Returns ``None`` if the sequence declared no
-        runtime parameters (in which case the file is not written).
+        runtime parameters: the file is still written, but holds a count of
+        zero and there is no ctrl scratchpad buffer object to bind to.
         """
         if self._params is not None:
             return self._params
         mlir_filename = self.op.artifacts[0].mlir_input.filename
         params_path = Path(mlir_filename + ".prj") / "params.txt"
         if not params_path.exists():
             return None
+        if params_path.read_text().split("\n", 1)[0].strip() == "0":
+            return None
         from aie.utils.hostruntime.xrtruntime.parameter_scratchpad import (
             ParameterScratchpad,
         )

iron/operators/_trace.py

@@ -39,13 +39,11 @@ def _default_coretile_events():
     ]
 
 
-def maybe_enable_trace(rt, trace_size, workers, coretile_events=None):
-    """Configure per-op hardware trace on ``rt`` if tracing is requested.
-
-    Call inside the ``rt.sequence(...)`` block, before ``rt.start(...)``.
+def maybe_enable_trace(prog, trace_size, workers, coretile_events=None):
+    """Configure per-op hardware trace if tracing is requested.
 
     Args:
-        rt: the ``Runtime`` being built.
+        prog: the ``Program`` being built.
         trace_size: the design's ``trace_size`` argument (may be None/0).
         workers: the design's workers; the first ``IRON_TRACE_NTILES`` are traced.
         coretile_events: override the default core-tile event set.
@@ -61,7 +59,7 @@ def maybe_enable_trace(rt, trace_size, workers, coretile_events=None):
     # meaningless (a negative slice index would silently drop the LAST worker).
     ntiles = max(0, int(os.environ.get("IRON_TRACE_NTILES", "1")))
 
-    rt.enable_trace(
+    prog.enable_trace(
         ts,
         workers=list(workers)[:ntiles],
         coretile_events=(

iron/operators/axpy/design.py

@@ -4,7 +4,7 @@
 from ml_dtypes import bfloat16
 import numpy as np
 
-from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker
+from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker
 from aie.helpers.taplib.tap import TensorAccessPattern
 from aie.iron.controlflow import range_
 from iron.operators._trace import maybe_enable_trace
@@ -84,38 +84,45 @@ def core_body(of_in1, of_in2, of_out, axpy):
     ]
 
     # Runtime operations to move data to/from the AIE-array
-    rt = Runtime()
-    with rt.sequence(tensor_ty, tensor_ty, tensor_ty) as (A, B, C):
-        maybe_enable_trace(rt, trace_size, my_workers)
-        rt.start(*my_workers)
-
+    def sequence(A, B, C, in1_prods, in2_prods, out_conses):
         # Initialize a group for parallel drain tasks, with fill resources free'd when drains complete.
-        tg = rt.task_group()
+        tg = TaskGroup()
 
         # Fill the input objectFIFOs with data
         for i in range(num_columns):
-            rt.fill(
-                of_in1s[i].prod(),
+            in1_prods[i].fill(
                 A,
                 taps[i],
-                task_group=tg,
+                group=tg,
             )
-            rt.fill(
-                of_in2s[i].prod(),
+            in2_prods[i].fill(
                 B,
                 taps[i],
-                task_group=tg,
+                group=tg,
             )
         # Drain the output objectFIFOs with data
         for i in range(num_columns):
-            rt.drain(
-                of_outs[i].cons(),
+            out_conses[i].drain(
                 C,
                 taps[i],
                 wait=True,  # wait for the transfer to complete and data to be available
-                task_group=tg,
+                group=tg,
             )
-        rt.finish_task_group(tg)
+        tg.finish()
+
+    rt = Runtime(
+        sequence,
+        [
+            tensor_ty,
+            tensor_ty,
+            tensor_ty,
+            [of_in1s[i].prod() for i in range(num_columns)],
+            [of_in2s[i].prod() for i in range(num_columns)],
+            [of_outs[i].cons() for i in range(num_columns)],
+        ],
+    )
 
     # Place program components (assign them resources on the device) and generate an MLIR module
-    return Program(dev, rt).resolve_program()
+    prog = Program(dev, rt, workers=my_workers)
+    maybe_enable_trace(prog, trace_size, my_workers)
+    return prog.resolve_program()

iron/operators/binary_elementwise_design.py

@@ -4,7 +4,7 @@
 from ml_dtypes import bfloat16
 import numpy as np
 
-from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker
+from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker
 from aie.helpers.taplib.tap import TensorAccessPattern
 from aie.iron.controlflow import range_
 from iron.operators._trace import maybe_enable_trace
@@ -83,37 +83,44 @@ def core_body(of_in1, of_in2, of_out, eltwise_fn):
     ]
 
     # Runtime operations to move data to/from the AIE-array
-    rt = Runtime()
-    with rt.sequence(tensor_ty, tensor_ty, tensor_ty) as (A, B, C):
-        maybe_enable_trace(rt, trace_size, my_workers)
-        rt.start(*my_workers)
-
-        tg = rt.task_group()
+    def sequence(A, B, C, in1_prods, in2_prods, out_conses):
+        tg = TaskGroup()
 
         # Fill the input objectFIFOs with data
         for i in range(num_columns):
-            rt.fill(
-                of_in1s[i].prod(),
+            in1_prods[i].fill(
                 A,
                 taps[i],
-                task_group=tg,
+                group=tg,
             )
-            rt.fill(
-                of_in2s[i].prod(),
+            in2_prods[i].fill(
                 B,
                 taps[i],
-                task_group=tg,
+                group=tg,
             )
         # Drain the output objectFIFOs with data
         for i in range(num_columns):
-            rt.drain(
-                of_outs[i].cons(),
+            out_conses[i].drain(
                 C,
                 taps[i],
                 wait=True,
-                task_group=tg,
+                group=tg,
             )
-        rt.finish_task_group(tg)
+        tg.finish()
+
+    rt = Runtime(
+        sequence,
+        [
+            tensor_ty,
+            tensor_ty,
+            tensor_ty,
+            [of_in1s[i].prod() for i in range(num_columns)],
+            [of_in2s[i].prod() for i in range(num_columns)],
+            [of_outs[i].cons() for i in range(num_columns)],
+        ],
+    )
 
     # Place program components and generate an MLIR module
-    return Program(dev, rt).resolve_program()
+    prog = Program(dev, rt, workers=my_workers)
+    maybe_enable_trace(prog, trace_size, my_workers)
+    return prog.resolve_program()

iron/operators/channeled_unary_design.py

@@ -4,7 +4,7 @@
 from ml_dtypes import bfloat16
 import numpy as np
 
-from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker
+from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker
 from aie.helpers.taplib.tap import TensorAccessPattern
 from aie.iron.controlflow import range_
 from iron.operators._trace import maybe_enable_trace
@@ -97,33 +97,39 @@ def core_fn(of_in, of_out, kernel_line):
     ]
 
     # Runtime operations to move data to/from the AIE-array
-    rt = Runtime()
-    with rt.sequence(transfer_type, transfer_type) as (a_in, b_out):
-        maybe_enable_trace(rt, trace_size, my_workers)
-        rt.start(*my_workers)
-
-        tg = rt.task_group()
+    def sequence(a_in, b_out, in_prods, out_conses):
+        tg = TaskGroup()
 
         # Fill the input objectFIFOs with data
         for i in range(num_columns):
             for j in range(num_channels):
-                rt.fill(
-                    of_ins[i * num_channels + j].prod(),
+                in_prods[i * num_channels + j].fill(
                     a_in,
                     taps[i * num_channels + j],
-                    task_group=tg,
+                    group=tg,
                 )
         # Drain the output objectFIFOs with data
         for i in range(num_columns):
             for j in range(num_channels):
-                rt.drain(
-                    of_outs[i * num_channels + j].cons(),
+                out_conses[i * num_channels + j].drain(
                     b_out,
                     taps[i * num_channels + j],
                     wait=True,
-                    task_group=tg,
+                    group=tg,
                 )
-        rt.finish_task_group(tg)
+        tg.finish()
+
+    rt = Runtime(
+        sequence,
+        [
+            transfer_type,
+            transfer_type,
+            [of.prod() for of in of_ins],
+            [of.cons() for of in of_outs],
+        ],
+    )
 
     # Place components and generate an MLIR module
-    return Program(dev, rt).resolve_program()
+    prog = Program(dev, rt, workers=my_workers)
+    maybe_enable_trace(prog, trace_size, my_workers)
+    return prog.resolve_program()