Skip to content

Commit

Permalink
[onert] Add lifetime enum to LayerScopeTensor (#13861)
Browse files Browse the repository at this point in the history
This PR adds lifetime info to the LayerScopeTensor.
Lifetime info will be used to plan and allocate buffer for LayerScopeTensor.

ONE-DCO-1.0-Signed-off-by: seunghui youn <[email protected]>
  • Loading branch information
zetwhite authored Sep 2, 2024
1 parent dc0a540 commit ae411be
Showing 1 changed file with 19 additions and 1 deletion.
20 changes: 19 additions & 1 deletion runtime/onert/core/include/backend/train/LayerScopeTensor.h
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,12 @@ namespace backend
namespace train
{

enum class LayerScopeTensorLifeTime : unsigned char
{
BACKWARD, // alive during backward()
FORWARD_TO_BACKWARD, // alive from forward() to backward()
};

// LayerScopeTensor is a tensor that is not shown in graph but required by each layer.
// It is accessed within one operation layer.
class LayerScopeTensor final : public basic::Tensor
Expand All @@ -34,10 +40,22 @@ class LayerScopeTensor final : public basic::Tensor
LayerScopeTensor() = delete;

public:
LayerScopeTensor(const ir::OperandInfo &info) : basic::Tensor(info, nullptr)
LayerScopeTensor(const ir::OperandInfo &info, LayerScopeTensorLifeTime lt)
: basic::Tensor(info, nullptr), _lifetime(lt)
{
// DO NOTHING
}

LayerScopeTensor(const ir::OperandInfo &info)
: basic::Tensor(info, nullptr), _lifetime(LayerScopeTensorLifeTime::BACKWARD)
{
// DO NOTHING
}

LayerScopeTensorLifeTime lifetime() const { return _lifetime; }

private:
LayerScopeTensorLifeTime _lifetime;
};

} // namespace train
Expand Down

0 comments on commit ae411be

Please sign in to comment.