EZ API: lazy kernel creation, eliminate done()

Refactor ProgramBuilder so kernels are created lazily at build() time
instead of eagerly in reader()/writer()/compute(). This enables
defines() and named_args() to be called on KernelRef after the kernel
method, and eliminates the need for done() to chain back to the builder.

KernelRef now stores all kernel configuration as a deferred descriptor
and provides forwarding methods (cb, reader, writer, compute, kernel,
on, semaphore, build) that implicitly transition back to the builder.

Author: olofj

tests/tt_metal/tt_metal/ez/kernels/tiles_add_named.cpp

@@ -0,0 +1,39 @@
+// SPDX-FileCopyrightText: (c) 2026 Olof Johansson <olof@lixom.net>
+//
+// SPDX-License-Identifier: Apache-2.0
+
+// Variant of tiles_add that reads CB indices from named compile-time args
+// instead of hardcoding them. Used to test per-kernel named_args on KernelRef.
+
+#include <cstdint>
+#include "api/compute/common.h"
+#include "api/compute/tile_move_copy.h"
+#include "api/compute/eltwise_binary.h"
+#include "api/compute/compute_kernel_api.h"
+
+void kernel_main() {
+    uint32_t n_tiles = get_arg_val<uint32_t>(0);
+
+    constexpr auto cb_in0 = get_named_compile_time_arg_val("cb_in0");
+    constexpr auto cb_in1 = get_named_compile_time_arg_val("cb_in1");
+    constexpr auto cb_out0 = get_named_compile_time_arg_val("cb_out0");
+    constexpr uint32_t dst_reg = 0;
+
+    binary_op_init_common(cb_in0, cb_in1, cb_out0);
+    add_tiles_init(cb_in0, cb_in1);
+
+    for (uint32_t i = 0; i < n_tiles; i++) {
+        cb_wait_front(cb_in0, 1);
+        cb_wait_front(cb_in1, 1);
+        tile_regs_acquire();
+        add_tiles(cb_in0, cb_in1, 0, 0, dst_reg);
+        tile_regs_commit();
+        tile_regs_wait();
+        cb_reserve_back(cb_out0, 1);
+        pack_tile(dst_reg, cb_out0);
+        cb_push_back(cb_out0, 1);
+        cb_pop_front(cb_in0, 1);
+        cb_pop_front(cb_in1, 1);
+        tile_regs_release();
+    }
+}

tests/tt_metal/tt_metal/ez/test_ez_api.cpp

@@ -20,7 +20,6 @@ TEST_F(EzApiTest, HelloWorld) {
     DeviceContext ctx(0);
     auto program = ProgramBuilder(CoreCoord{0, 0})
                        .compute("tests/tt_metal/tt_metal/ez/kernels/void_compute.cpp")
-                       .done()
                        .build();
     ctx.run(std::move(program));
 }
@@ -79,15 +78,12 @@ TEST_F(EzApiTest, EltwiseBinary) {
                            "tests/tt_metal/tt_metal/ez/kernels/read_tiles.cpp",
                            {src0, src1})
                        .runtime_args({src0->address(), src1->address(), n_tiles})
-                       .done()
                        .compute("tests/tt_metal/tt_metal/ez/kernels/tiles_add.cpp")
                        .runtime_args({n_tiles})
-                       .done()
                        .writer(
                            "tests/tt_metal/tt_metal/ez/kernels/write_tile.cpp",
                            {dst})
                        .runtime_args({dst->address(), n_tiles})
-                       .done()
                        .build();
 
     ctx.run(std::move(program));
@@ -111,7 +107,6 @@ TEST_F(EzApiTest, MultiCore) {
     auto& k = builder.compute("tests/tt_metal/tt_metal/ez/kernels/void_compute.cpp");
     k.runtime_args_at(CoreCoord{0, 0}, {});
     k.runtime_args_at(CoreCoord{1, 0}, {});
-    k.done();
 
     auto program = builder.build();
     ctx.run(std::move(program));
@@ -127,8 +122,7 @@ TEST_F(EzApiTest, PerCoreLambdaArgs) {
         .runtime_args([](const CoreCoord& core) -> std::vector<uint32_t> {
             // Each core gets a unique arg based on its x coordinate.
             return {core.x};
-        })
-        .done();
+        });
 
     auto program = builder.build();
     ctx.run(std::move(program));
@@ -142,10 +136,8 @@ TEST_F(EzApiTest, CoreOverride) {
 
     auto program = ProgramBuilder(default_core)
                        .compute("tests/tt_metal/tt_metal/ez/kernels/void_compute.cpp")
-                       .done()
                        .on(other_core)
                        .compute("tests/tt_metal/tt_metal/ez/kernels/void_compute.cpp")
-                       .done()
                        .build();
 
     ctx.run(std::move(program));
@@ -179,7 +171,6 @@ TEST_F(EzApiTest, L1BackedCircularBuffer) {
     auto program = ProgramBuilder(CoreCoord{0, 0})
                        .cb(tt::CBIndex::c_0, l1_buf)
                        .compute("tests/tt_metal/tt_metal/ez/kernels/void_compute.cpp")
-                       .done()
                        .build();
     ctx.run(std::move(program));
 }
@@ -208,8 +199,7 @@ TEST_F(EzApiTest, Semaphore) {
     uint32_t sem_addr = builder.semaphore(0);
 
     builder.compute("tests/tt_metal/tt_metal/ez/kernels/void_compute.cpp")
-        .runtime_args([&](const CoreCoord&) -> std::vector<uint32_t> { return {sem_addr}; })
-        .done();
+        .runtime_args([&](const CoreCoord&) -> std::vector<uint32_t> { return {sem_addr}; });
     auto program = builder.build();
     ctx.run(std::move(program));
 }
@@ -229,8 +219,63 @@ TEST_F(EzApiTest, NonBlockingLaunchAndFinish) {
     DeviceContext ctx(0);
     auto program = ProgramBuilder(CoreCoord{0, 0})
                        .compute("tests/tt_metal/tt_metal/ez/kernels/void_compute.cpp")
-                       .done()
                        .build();
     ctx.launch(std::move(program));
     ctx.finish();
 }
+
+TEST_F(EzApiTest, PerKernelNamedArgs) {
+    // Verify that named_args() called on a KernelRef applies to that kernel only.
+    // The compute kernel reads CB indices from named compile-time args; if they
+    // were routed to the wrong kernel (or lost), compilation would fail.
+    DeviceContext ctx(0);
+    constexpr uint32_t n_tiles = 4;
+    constexpr uint32_t elements = n_tiles * tt::constants::TILE_WIDTH * tt::constants::TILE_HEIGHT;
+
+    auto src0 = ctx.dram_tile_buffer(n_tiles);
+    auto src1 = ctx.dram_tile_buffer(n_tiles);
+    auto dst = ctx.dram_tile_buffer(n_tiles);
+
+    std::mt19937 rng(999);
+    std::uniform_real_distribution<float> dist(0.0f, 1.0f);
+    std::vector<bfloat16> a_data(elements), b_data(elements);
+    for (size_t i = 0; i < elements; ++i) {
+        a_data[i] = bfloat16(dist(rng));
+        b_data[i] = bfloat16(dist(rng));
+    }
+
+    ctx.write(src0, a_data);
+    ctx.write(src1, b_data);
+
+    // named_args() is called on the KernelRef returned by .compute(), not on the builder.
+    constexpr CoreCoord core = {0, 0};
+    auto program = ProgramBuilder(core)
+                       .cb(tt::CBIndex::c_0)
+                       .cb(tt::CBIndex::c_1)
+                       .cb(tt::CBIndex::c_16)
+                       .reader(
+                           "tests/tt_metal/tt_metal/ez/kernels/read_tiles.cpp",
+                           {src0, src1})
+                       .runtime_args({src0->address(), src1->address(), n_tiles})
+                       .compute("tests/tt_metal/tt_metal/ez/kernels/tiles_add_named.cpp")
+                       .named_args({{"cb_in0", (uint32_t)tt::CBIndex::c_0},
+                                    {"cb_in1", (uint32_t)tt::CBIndex::c_1},
+                                    {"cb_out0", (uint32_t)tt::CBIndex::c_16}})
+                       .runtime_args({n_tiles})
+                       .writer(
+                           "tests/tt_metal/tt_metal/ez/kernels/write_tile.cpp",
+                           {dst})
+                       .runtime_args({dst->address(), n_tiles})
+                       .build();
+
+    ctx.run(std::move(program));
+    auto result = ctx.read<bfloat16>(dst);
+
+    ASSERT_EQ(result.size(), elements);
+    constexpr float eps = 1e-2f;
+    for (size_t i = 0; i < elements; ++i) {
+        float expected = static_cast<float>(a_data[i]) + static_cast<float>(b_data[i]);
+        float actual = static_cast<float>(result[i]);
+        EXPECT_NEAR(expected, actual, eps) << "Mismatch at index " << i;
+    }
+}

tt_metal/api/tt-metalium/experimental/ez/program_builder.hpp

@@ -27,9 +27,23 @@ using CoreSpec = std::variant<CoreCoord, CoreRange, CoreRangeSet>;
 
 class ProgramBuilder;
 
-// Handle to a kernel added via ProgramBuilder, supporting fluent runtime_args configuration.
+// Deferred kernel descriptor added via ProgramBuilder. Stores all configuration
+// and defers CreateKernel + SetRuntimeArgs to build() time.
+// Supports fluent chaining: methods for kernel configuration (defines, named_args,
+// runtime_args) return KernelRef&, while builder methods (cb, reader, writer,
+// compute, kernel, on, semaphore, build) forward to the parent ProgramBuilder.
 class KernelRef {
 public:
+    enum class Type { Reader, Writer, Compute, Custom };
+
+    // Kernel configuration methods (return KernelRef& for chaining).
+
+    // Set preprocessor defines for this kernel.
+    KernelRef& defines(const std::map<std::string, std::string>& defs);
+
+    // Set named compile-time args for this kernel.
+    KernelRef& named_args(const std::unordered_map<std::string, uint32_t>& args);
+
     // Set runtime args applied uniformly across all cores in the kernel's core spec.
     KernelRef& runtime_args(std::initializer_list<uint32_t> args);
     KernelRef& runtime_args(const std::vector<uint32_t>& args);
@@ -40,19 +54,70 @@ class KernelRef {
     // Set runtime args for a specific core.
     KernelRef& runtime_args_at(const CoreCoord& core, const std::vector<uint32_t>& args);
 
-    // Return to the ProgramBuilder for chaining.
-    ProgramBuilder& done();
+    // Forwarding methods to ProgramBuilder (return appropriate type for chaining).
+
+    ProgramBuilder& cb(tt::CBIndex index, uint32_t num_tiles = 2,
+                       tt::DataFormat fmt = tt::DataFormat::Float16_b);
+    ProgramBuilder& cb(tt::CBIndex index, tt::DataFormat fmt, uint32_t num_tiles, uint32_t page_size);
+    ProgramBuilder& cb(tt::CBIndex index, const std::shared_ptr<distributed::MeshBuffer>& l1_buffer,
+                       uint32_t num_tiles = 0, tt::DataFormat fmt = tt::DataFormat::Float16_b);
+    ProgramBuilder& cb(const CircularBufferConfig& config);
+
+    KernelRef& reader(
+        const std::string& path,
+        const std::vector<std::shared_ptr<distributed::MeshBuffer>>& buffers = {},
+        const std::vector<uint32_t>& compile_args = {});
+    KernelRef& writer(
+        const std::string& path,
+        const std::vector<std::shared_ptr<distributed::MeshBuffer>>& buffers = {},
+        const std::vector<uint32_t>& compile_args = {});
+    KernelRef& compute(
+        const std::string& path,
+        MathFidelity fidelity = MathFidelity::HiFi4,
+        const std::vector<uint32_t>& compile_args = {});
+    KernelRef& compute(const std::string& path, const ComputeConfig& config);
+    KernelRef& kernel(
+        const std::string& path,
+        const std::variant<DataMovementConfig, ComputeConfig, EthernetConfig>& config);
 
-    // Access the underlying kernel handle.
-    KernelHandle handle() const;
+    ProgramBuilder& on(const CoreSpec& core_spec);
+    ProgramBuilder& defines_next(const std::map<std::string, std::string>& defs);
+    ProgramBuilder& named_args_next(const std::unordered_map<std::string, uint32_t>& args);
+    uint32_t semaphore(uint32_t initial_value = 0);
+    uint32_t semaphore(const CoreSpec& cores, uint32_t initial_value = 0);
+    Program build();
 
 private:
     friend class ProgramBuilder;
-    KernelRef(ProgramBuilder& builder, KernelHandle handle, CoreSpec core_spec);
+
+    using DeferredRuntimeArgs = std::function<void(Program&, KernelHandle, const CoreSpec&)>;
+
+    KernelRef(ProgramBuilder& builder, Type type, std::string path, CoreSpec core_spec);
 
     ProgramBuilder& builder_;
-    KernelHandle handle_;
+    Type type_;
+    std::string path_;
     CoreSpec core_spec_;
+
+    // Reader/Writer specific.
+    std::vector<std::shared_ptr<distributed::MeshBuffer>> buffers_;
+    std::vector<uint32_t> compile_args_;
+
+    // Compute specific.
+    MathFidelity fidelity_ = MathFidelity::HiFi4;
+
+    // Custom kernel config (for kernel() with full config).
+    std::optional<std::variant<DataMovementConfig, ComputeConfig, EthernetConfig>> custom_config_;
+
+    // Shared.
+    std::map<std::string, std::string> defines_;
+    std::unordered_map<std::string, uint32_t> named_compile_args_;
+
+    // Deferred runtime args — replayed at build() time.
+    std::vector<DeferredRuntimeArgs> deferred_runtime_args_;
+
+    // Materialize this kernel into the program. Called by ProgramBuilder::build().
+    void materialize(Program& program);
 };
 
 // Fluent builder for constructing a Program with circular buffers and kernels.