tenstorrent / tt-metal

:metal: TT-NN operator library, and TT-Metalium low level kernel programming model.
Apache License 2.0
471 stars 74 forks source link

Add TensorAsnc to handle promise feature and remove this feature from Tensor #8813

Open eyonland opened 5 months ago

eyonland commented 5 months ago

Problem: We do not currently have a way to cleanly separate function calls that are intended to launch an async operation from those that are intended to run within the device thread. This means that when we compose operations, the unintended consequences is that the operation will sometimes invoke launch_op when in fact we intended for the operation to simply run as-is on the device.

Solution: Introduce the TensorAsync and remove the existing promise-like features from Tensor. By having a distinct type, the same function can be overloaded such that a function receiving arguments with TorchAsync would naturally be expected to invoke launch_op whereas the function receiving Tensor would simply be expected to run on device. By having this distinction, we would then be able to easily compose operations that are expected to run in within one invocation of launch_op.

This issue will be the global issue should this work be accepted and will drive sub-issue(s) per operation that will update the underlying ops to differentiate with these overloaded functions.

tt-asaigal commented 5 months ago

Proposal: https://docs.google.com/presentation/d/1Cs-8Rjq8duPq_jIkwi8W3erRXF5KlWzq/edit#slide=id.p1