From 6c8542786d209466e8716aee57211feca835b9ad Mon Sep 17 00:00:00 2001 From: Youhezhen Date: Mon, 20 Jul 2026 03:59:09 -0700 Subject: [PATCH] Fix: add A5 early-resolve task hint API - Add the generated-code-compatible hint to A5 task arguments without enabling A2/A3 early-dispatch behavior - Expose the hint through dependency-owning arguments and cover its reset semantics --- .../orchestration/pto_arg_with_deps.h | 6 ++- .../runtime/pto_types.h | 7 ++++ tests/ut/cpp/CMakeLists.txt | 1 + tests/ut/cpp/a5/test_task_args.cpp | 40 +++++++++++++++++++ 4 files changed, 52 insertions(+), 2 deletions(-) create mode 100644 tests/ut/cpp/a5/test_task_args.cpp diff --git a/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h b/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h index 20144d3a1..530de5672 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h @@ -54,15 +54,17 @@ class L0TaskArgsWithDeps : private L0TaskArgs { using L0TaskArgs::add_scalar; using L0TaskArgs::add_scalars; using L0TaskArgs::add_scalars_i32; + using L0TaskArgs::allow_early_resolve; // task hint (getter) using L0TaskArgs::copy_scalars_from; + using L0TaskArgs::set_allow_early_resolve; // task hint (setter) + using L0TaskArgs::set_task_timing_slot; // selective task-timing slot (setter) + using L0TaskArgs::task_timing_slot; // selective task-timing slot (getter) // Error / status — forward to Arg using L0TaskArgs::error_msg; using L0TaskArgs::has_error; using L0TaskArgs::launch_spec; using L0TaskArgs::set_error; - using L0TaskArgs::set_task_timing_slot; // selective task-timing slot (setter) - using L0TaskArgs::task_timing_slot; // selective task-timing slot (getter) // NOT exposed: set_dependencies, explicit_dep_count, explicit_dep, // explicit_deps_data — these are the primitive-layer dep API. Users of diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h index bd0cd66be..14957c2b4 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h @@ -242,6 +242,12 @@ struct Arg : TaskArgsTpl { const char *error_msg{nullptr}; PTO2LaunchSpec launch_spec; // SPMD launch parameters (block_num, etc.) + // Cross-architecture codegen may set this task hint. A5 does not currently + // change dispatch behavior based on it. + bool allow_early_resolve_{false}; + void set_allow_early_resolve(bool value = true) { allow_early_resolve_ = value; } + bool allow_early_resolve() const { return allow_early_resolve_; } + // Dispatch predicate (codegen-author set; default op == NONE = always // dispatch). A FALSE result at the dispatch point retires the task inline // through the dep-only path — never dispatched to an AICore — while still @@ -273,6 +279,7 @@ struct Arg : TaskArgsTpl { #endif explicit_deps_ = nullptr; explicit_dep_count_ = 0; + allow_early_resolve_ = false; predicate_ = L0TaskPredicate{}; task_timing_slot_ = TASK_TIMING_SLOT_NONE; } diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index a49e418c2..2afd90461 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -606,6 +606,7 @@ add_a2a3_runtime_test(test_a2a3_scope_stats_collector a2a3/test_scope_stats_coll # A5 tests (src/a5/runtime/tensormap_and_ringbuffer/) # --------------------------------------------------------------------------- add_a5_test(test_a5_fatal a5/test_a5_fatal.cpp) +add_a5_test(test_a5_task_args a5/test_task_args.cpp) # A5 trb runtime UTs — mirror of a2a3 trb runtime UTs, link against a5_rt_objs. # Target names carry the a5_ prefix because hierarchical/test_tensormap (and diff --git a/tests/ut/cpp/a5/test_task_args.cpp b/tests/ut/cpp/a5/test_task_args.cpp new file mode 100644 index 000000000..68a94f5e8 --- /dev/null +++ b/tests/ut/cpp/a5/test_task_args.cpp @@ -0,0 +1,40 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include "pto_orchestration_api.h" + +TEST(A5TaskArgs, EarlyResolveHintRoundTripsAndResets) { + L0TaskArgs args; + + EXPECT_FALSE(args.allow_early_resolve()); + args.set_allow_early_resolve(); + EXPECT_TRUE(args.allow_early_resolve()); + + args.set_allow_early_resolve(false); + EXPECT_FALSE(args.allow_early_resolve()); + + args.set_allow_early_resolve(true); + args.reset(); + EXPECT_FALSE(args.allow_early_resolve()); +} + +TEST(A5TaskArgs, EarlyResolveHintIsExposedWithDeps) { + L0TaskArgsWithDeps<> args; + + EXPECT_FALSE(args.allow_early_resolve()); + args.set_allow_early_resolve(true); + EXPECT_TRUE(args.allow_early_resolve()); + + args.reset(); + EXPECT_FALSE(args.allow_early_resolve()); +}