Skip to content

[Feature] tmr: let get/set_tensor_data declare the tasks an access is ordered against #2084

Description

@poursoul

Summary

Extend get_tensor_data / set_tensor_data in tensormap_and_ringbuffer with an optional list of task ids the access must be ordered against, mirroring what Arg::set_dependencies does for tasks. The accessor would wait on the declared tasks — a read until they complete, a write until they retire — in addition to whatever it discovers today through the tensor's owner_task_id and the TensorMap.

Scope is tensormap_and_ringbuffer only; see "Why host_build_graph is out of scope" below.

Motivation / Use Case

Scalar access discovers a producer through exactly two channels: the tensor's owner_task_id, and an overlapping entry in the TensorMap. Both are products of automatic dependency generation, and both are switched off in two cases (src/a5/runtime/tensormap_and_ringbuffer/runtime/dep_compute.h):

  • inside a manual scope, register_task_outputs() returns immediately, so nothing a task writes there is registered;
  • a tensor carrying manual_dep is skipped by the same function even in an automatic scope.

So an external, host-staged tensor (invalid owner_task_id) submitted as INOUT / OUTPUT_EXISTING through either path is invisible to both accessors. get_tensor_data returns the pre-kernel bytes and set_tensor_data writes a value the kernel later overwrites — with no error, and no way for the caller to ask for the ordering it needs.

add_dep cannot close this. It is a task-to-task edge and carries no buffer identity, so from a tensor there is no path back to the task that writes it. docs/war-anti-dependency.md already records the same limitation from the other side: a reader wired only through add_dep is not in the set set_tensor_data waits on, and the doc's only advice is to promote that reader to add_inout. What is missing is a channel where the caller declares the edge, which is why it belongs on the accessor rather than in more runtime inference.

The WAR case that doc steers to add_inout falls out of the same extension: a host write that must wait for readers the TensorMap does not hold becomes expressible without serializing independent readers.

Why host_build_graph is out of scope

host_build_graph does not support a tensor having a producer at all. Its orchestration runs to completion before the device starts, so a produced buffer simply has no value during the call; #2080 makes both accessors reject any tensor with a producer with INVALID_ARGS. A declared dependency could express nothing there — it could only ever turn a rejection into the same rejection.

Following the precedent of the deliberately absent TensorCreateInfo::set_initial_value in that runtime, the parameter should not exist in hbg's orchestration/orchestration_api.h, so an orchestration that passes one fails to compile against that runtime rather than behaving differently at run time. The two runtimes keep separate orchestration_api.h files, so this needs no compatibility shim.

Proposed API / Behavior

// src/{arch}/runtime/tensormap_and_ringbuffer/orchestration/orchestration_api.h
template <typename T = uint64_t>
T get_tensor_data(const ChipTensor &tensor, uint32_t ndims, const uint32_t indices[],
                  const TaskId *deps = nullptr, uint32_t dep_count = 0);

template <typename T>
void set_tensor_data(const ChipTensor &tensor, uint32_t ndims, const uint32_t indices[], T value,
                     const TaskId *deps = nullptr, uint32_t dep_count = 0);
  • get: wait for every declared task to reach COMPLETED, on top of the producers found through owner_task_id / TensorMap.
  • set: the same, plus the consumer wait the runtime already performs for discovered producers — a declared task is treated as a live accessor of the buffer, which is what makes the WAR case expressible.
  • Default arguments keep every existing call site source-compatible. RuntimeOps's function-pointer signature changes, but orchestration shared objects are compiled by KernelCompiler against the same headers, so both sides move together.
  • Validation should mirror append_fanin_or_fail: ids must be valid, in the GLOBAL id space, and below the submitted task count; anything else is INVALID_ARGS.
  • SCALAR_DATA_ACCESS.md for the runtime gains the contract: the accessors screen what automatic dependency generation registered, and a caller in a manual scope or using manual_dep owns the declaration.

Alternatives Considered

  • Promote the reader/writer to add_inout — the workaround docs/war-anti-dependency.md recommends today. It serializes otherwise independent readers, and it does not help inside a manual scope, where nothing is registered at all.
  • Infer more in the runtime (e.g. register manual-scope outputs anyway) — contradicts the manual-scope contract, whose whole point is that the caller owns dependency construction, and re-introduces the tensormap cost manual scopes exist to avoid.
  • Track reader accesses automatically in the TensorMap ([Feature] Track reader accesses in TensorMap and derive WAR dependencies automatically #1388) — complementary rather than a substitute: that proposal is explicitly scoped to automatic-scope, non-manual_dep arguments, which is exactly the set this issue does not cover.

Additional Context

Raised by @zhusy54 in review of #2080 (inline comment on src/common/host_build_graph/shared/runtime_core.cpp). That PR removes hbg's producer/consumer wait, which could never observe device progress; the discovery gap named here predates it and is unchanged by it — the removed wait_for_tensor_ready() read the same two channels, so a manual-scope or manual_dep tensor passed through silently before and after.

Related: #1388, docs/war-anti-dependency.md.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions