Skip to content

Commit

Permalink
update backend
Browse files Browse the repository at this point in the history
  • Loading branch information
zetwhite committed Jul 17, 2024
1 parent 022c79c commit a6ed87c
Show file tree
Hide file tree
Showing 5 changed files with 14 additions and 130 deletions.
120 changes: 7 additions & 113 deletions runtime/onert/backend/train/ExtraTensorGenerator.cc
Original file line number Diff line number Diff line change
Expand Up @@ -67,16 +67,6 @@ ExtraTensorGenerator::ExtraTensorGenerator(const ir::train::TrainableGraph &tgra
std::shared_ptr<TensorRegistry> &tensor_registry)
: _tgraph(tgraph), _tensor_builder(tensor_builder), _tensor_reg(tensor_registry){};

/* void ExtraTensorGenerator::generate()
{
_tgraph.operations().iterate([&](const ir::OperationIndex &, const ir::IOperation &op) {
const auto trainable_op = dynamic_cast<const ir::train::TrainableOperation *>(&op);
trainable_op->accept(*this);
});
handle_requests();
}*/

void ExtraTensorGenerator::notify_first_use(ir::OperationIndex op_idx,
const ExtraTensorRequests &reqs,
bool (*cond)(const ExtraTensorLifeTime))
Expand All @@ -101,120 +91,24 @@ void ExtraTensorGenerator::notify_last_use(ir::OperationIndex op_idx,
return;
}

std::vector<ExtraTensor *> ExtraTensorGenerator::register_requests(ir::OperationIndex op_idx,
const ExtraTensorRequests &reqs)
void ExtraTensorGenerator::generate_extra_tensor(ir::OperationIndex op_idx,
const ExtraTensorRequests &reqs)
{
std::vector<ExtraTensor *> registered_tensors;

// save requests
// save request, _idx_to_reuqests used for memory planning
_idx_to_requests[op_idx] = reqs;

// register tensor
for (size_t i = 0; i < reqs.size(); i++)
{
// register tensor
ExtraTensorIndex tensor_idx(op_idx, i);
_tensor_builder->registerExtraTensorInfo(tensor_idx, reqs[i].info, reqs[i].layout);
registered_tensors.push_back(_tensor_reg->getExtraTensor(tensor_idx));
}
return registered_tensors;
}
/*
void ExtraTensorGenerator::handle_requests()
{
// register tensor
for (const auto &pair : _idx_to_requests)
{
const auto &reqs = pair.second;
for (size_t i = 0; i < reqs.size(); ++i)
{
ExtraTensorIndex idx(pair.first, i);
_tensor_builder->registerExtraTensorInfo(idx, reqs[i].info, reqs[i].layout);
}
}

// forward
for (const auto &op_index : _tgraph.topolSortOperations())
{
if (_idx_to_requests.find(op_index) == _idx_to_requests.end())
continue;
const auto &reqs = _idx_to_requests[op_index];
notify_first_use(op_index, reqs, is_fwd_or_fwd2bwd);
notify_last_use(op_index, reqs, is_fwd);
}
// backward
for (const auto &op_index : _tgraph.btopolSortOperations())
{
if (_idx_to_requests.find(op_index) == _idx_to_requests.end())
continue;
const auto &reqs = _idx_to_requests[op_index];
notify_first_use(op_index, reqs, is_bwd);
notify_last_use(op_index, reqs, is_bwd_or_fwd2bwd);
// return registered tensor
auto generated_tensor = _tensor_reg->getExtraTensor(tensor_idx);
*reqs[i].address = generated_tensor;
}
}
void ExtraTensorGenerator::visit(const ir::train::operation::FullyConnected &node)
{
using ir::train::operation::FullyConnected;
const auto out_index{node.getOutputs().at(0)};
const auto in_index{node.getInputs().at(FullyConnected::Input::INPUT)};
const auto weights_index{node.getInputs().at(FullyConnected::Input::WEIGHT)};
auto in_tensor = _tensor_reg->getPortableTensor(in_index);
auto weights_tensor = _tensor_reg->getTrainableTensor(weights_index);
auto out_back_prop_tensor = _tensor_reg->getBackPropTensor(out_index);
const auto activation = node.param().activation;
auto requests = ops::FullyConnectedLayer::requestExtraTensors(weights_tensor, in_tensor,
out_back_prop_tensor, activation);
auto op_idx = _node_to_idx[&node];
_idx_to_requests.emplace(op_idx, requests);
return;
}
void ExtraTensorGenerator::visit(const ir::train::operation::Conv2D &node)
{
using ir::train::operation::Conv2D;
const auto weights_index{node.getInputs().at(Conv2D::Input::KERNEL)};
auto weights_tensor = _tensor_reg->getTrainableTensor(weights_index);
const auto out_index{node.getOutputs().at(0)};
auto out_back_prop_tensor = _tensor_reg->getBackPropTensor(out_index);
const auto activation = node.param().activation;
auto requests =
ops::ConvolutionLayer::requestExtraTensors(weights_tensor, out_back_prop_tensor, activation);
auto op_idx = _node_to_idx[&node];
_idx_to_requests.emplace(op_idx, requests);
return;
}*/

/*
void ExtraTensorGenerator::visit(const ir::train::operation::Pool2D &node)
{
using ir::train::operation::Pool2D;
const auto out_index{node.getOutputs().at(0)};
auto out_back_prop_tensor = _tensor_reg->getBackPropTensor(out_index);
const auto activation = node.param().activation;
auto requests = ops::ConvolutionLayer::requestExtraTensors(out_back_prop_tensor, activation);
auto op_idx = _node_to_idx[&node];
_idx_to_requests.emplace(op_idx, requests);
return;
}
*/

} // namespace train
} // namespace backend
Expand Down
4 changes: 1 addition & 3 deletions runtime/onert/backend/train/ExtraTensorGenerator.h
Original file line number Diff line number Diff line change
Expand Up @@ -50,14 +50,12 @@ class ExtraTensorGenerator
bool (*cond)(const ExtraTensorLifeTime));

public:
std::vector<ExtraTensor *> register_requests(ir::OperationIndex idx,
const ExtraTensorRequests &requests);
void generate_extra_tensor(ir::OperationIndex idx, const ExtraTensorRequests &requests);

private:
const ir::train::TrainableGraph &_tgraph;
std::shared_ptr<TensorBuilder> _tensor_builder;
std::shared_ptr<TensorRegistry> _tensor_reg;
// std::unordered_map<const ir::train::ITrainableOperation *, ir::OperationIndex> _node_to_idx;
std::unordered_map<ir::OperationIndex, ExtraTensorRequests> _idx_to_requests;
};

Expand Down
7 changes: 5 additions & 2 deletions runtime/onert/backend/train/ExtraTensorRequest.h
Original file line number Diff line number Diff line change
Expand Up @@ -36,18 +36,21 @@ enum class ExtraTensorLifeTime

struct ExtraTensorRequest
{
static ExtraTensorRequest createRequestLike(const IPortableTensor *origin)
static ExtraTensorRequest createRequestLike(const IPortableTensor *origin,
ExtraTensor **const addr)
{
ExtraTensorRequest r = {.info = origin->get_info(),
.layout = origin->layout(),
.lifetime =
ExtraTensorLifeTime::BACKWARD}; // default lifetime is BACKWARD
ExtraTensorLifeTime::BACKWARD, // < default lifetime is BACWKARD
.address = addr};
return r;
}

ir::OperandInfo info;
ir::Layout layout;
ExtraTensorLifeTime lifetime;
ExtraTensor **address;
};

using ExtraTensorRequests = std::vector<ExtraTensorRequest>;
Expand Down
3 changes: 1 addition & 2 deletions runtime/onert/backend/train/KernelGenerator.cc
Original file line number Diff line number Diff line change
Expand Up @@ -113,8 +113,7 @@ void KernelGenerator::genExtraTensors(ir::OperationIndex idx,
onert::exec::train::ITrainableFunction *layer)
{
auto reqs = layer->requestExtraTensors();
auto extra_tensors = _extra_tensor_generator->register_requests(idx, reqs);
layer->ConfigureExtraTensors(extra_tensors);
_extra_tensor_generator->generate_extra_tensor(idx, reqs);
}

std::unique_ptr<exec::train::TrainableFnSequence> KernelGenerator::generate(ir::OperationIndex idx)
Expand Down
10 changes: 0 additions & 10 deletions runtime/onert/core/src/compiler/ExecutorFactory.cc
Original file line number Diff line number Diff line change
Expand Up @@ -877,16 +877,6 @@ exec::IExecutor *ExecutorFactory::createTrainableExecutor(
// allocate extra tensors
}

// [idea2] : Generate Extra Tensors by traversing code_map

// Problem 1 :
// Since code_map lost its' backend context,
// it is hard to specify where estra tensor should be registered.

// Problme 2:
// Since ITrainableFunction is in exec area, it isn't a good idea to use member function of *Layer
// here, Istead of exec area interface, we have to implement another backend interface

if (order.size() != code_map.size())
{
throw std::runtime_error("ExecutorFactory: Some kernels are not generated");
Expand Down

0 comments on commit a6ed87c

Please sign in to comment.