diff options
Diffstat (limited to 'runtimes/neurun/backend/acl_cl')
-rw-r--r-- | runtimes/neurun/backend/acl_cl/StageGenerator.cc | 157 | ||||
-rw-r--r-- | runtimes/neurun/backend/acl_cl/StageGenerator.h | 3 |
2 files changed, 160 insertions, 0 deletions
diff --git a/runtimes/neurun/backend/acl_cl/StageGenerator.cc b/runtimes/neurun/backend/acl_cl/StageGenerator.cc index cc09d564f..73530aa51 100644 --- a/runtimes/neurun/backend/acl_cl/StageGenerator.cc +++ b/runtimes/neurun/backend/acl_cl/StageGenerator.cc @@ -2458,6 +2458,163 @@ void StageGenerator::visit(const model::operation::LogicalNotNode &node) }); } +void StageGenerator::visit(const model::operation::EqualNode &node) +{ + const auto output_index{node.getOutputs().at(0)}; + const auto input0_index{node.getInputs().at(model::operation::EqualNode::Input::INPUT0)}; + const auto input1_index{node.getInputs().at(model::operation::EqualNode::Input::INPUT1)}; + + if (!(_ctx.at(input0_index).shape() == _ctx.at(input1_index).shape())) + { + const auto broadcast_rank = + std::max(_ctx.at(input0_index).shape().rank(), _ctx.at(input1_index).shape().rank()); + const_cast<neurun::model::operand::Shape &>(_ctx.at(input0_index).shape()) + .extendRank(broadcast_rank); + const_cast<neurun::model::operand::Shape &>(_ctx.at(input1_index).shape()) + .extendRank(broadcast_rank); + } + + // Construct operation parameters + struct Param + { + model::operand::Index output_index; + model::operand::Index input0_index; + model::operand::Index input1_index; + }; + + Param param; + + param.output_index = output_index; + param.input0_index = input0_index; + param.input1_index = input1_index; + + auto tensors = _tensor_builder; + + returnStage([tensors, param](IExecutionBuilder &builder) { + auto output_alloc = tensors->at(param.output_index).get(); + auto input0_alloc = tensors->at(param.input0_index).get(); + auto input1_alloc = tensors->at(param.input1_index).get(); + + std::unique_ptr<::arm_compute::IFunction> fn; + + auto l = make_layer<::arm_compute::CLComparison>(); + + l->configure(input0_alloc->handle(), input1_alloc->handle(), output_alloc->handle(), + ::arm_compute::ComparisonOperation::Equal); + + fn = std::move(l); + + auto acl_fn = make_cl_function(std::move(fn)); + + builder.append(std::move(acl_fn)); + }); +} + +void StageGenerator::visit(const model::operation::SquaredDifferenceNode &node) +{ + const auto ofm_index{node.getOutputs().at(0)}; + const auto lhs_index{node.getInputs().at(model::operation::SquaredDifferenceNode::Input::LHS)}; + const auto rhs_index{node.getInputs().at(model::operation::SquaredDifferenceNode::Input::RHS)}; + + if (!(_ctx.at(lhs_index).shape() == _ctx.at(rhs_index).shape())) + { + const auto broadcast_rank = + std::max(_ctx.at(lhs_index).shape().rank(), _ctx.at(rhs_index).shape().rank()); + const_cast<neurun::model::operand::Shape &>(_ctx.at(lhs_index).shape()) + .extendRank(broadcast_rank); + const_cast<neurun::model::operand::Shape &>(_ctx.at(rhs_index).shape()) + .extendRank(broadcast_rank); + } + + // Construct operation parameters + struct Param + { + model::operand::Index ofm_index; + model::operand::Index lhs_index; + model::operand::Index rhs_index; + }; + + Param param; + + param.ofm_index = ofm_index; + param.lhs_index = lhs_index; + param.rhs_index = rhs_index; + + auto tensors = _tensor_builder; + + returnStage([tensors, param](IExecutionBuilder &builder) { + auto ofm_alloc = tensors->at(param.ofm_index).get(); + auto lhs_alloc = tensors->at(param.lhs_index).get(); + auto rhs_alloc = tensors->at(param.rhs_index).get(); + + std::unique_ptr<::arm_compute::IFunction> fn; + + auto l = make_layer<::arm_compute::CLSquaredDifference>(); + + l->configure(lhs_alloc->handle(), rhs_alloc->handle(), ofm_alloc->handle()); + + fn = std::move(l); + + auto acl_fn = make_cl_function(std::move(fn)); + + builder.append(std::move(acl_fn)); + }); +} + +void StageGenerator::visit(const model::operation::TopKV2Node &node) +{ + const auto outputValues_index{ + node.getOutputs().at(model::operation::TopKV2Node::Output::OUTPUT_VALUES)}; + const auto outputIndices_index{ + node.getOutputs().at(model::operation::TopKV2Node::Output::OUTPUT_INDICES)}; + + const auto inputData_index{node.getInputs().at(model::operation::TopKV2Node::Input::INPUT)}; + const auto k_index{node.param().k_index}; + + // Currently, we only support the vector input. + assert(_ctx.at(inputData_index).shape().rank() == 1 || + _ctx.at(inputData_index).shape().rank() == 2); + + const int32_t k = _ctx.at(k_index).asScalar<int32_t>(); + + // Construct operation parameters + struct Param + { + model::operand::Index outputValues_index; + model::operand::Index outputIndices_index; + + model::operand::Index inputData_index; + int32_t k; + }; + + Param param; + + param.outputValues_index = outputValues_index; + param.outputIndices_index = outputIndices_index; + param.inputData_index = inputData_index; + param.k = k; + + auto tensors = _tensor_builder; + + returnStage([tensors, param](IExecutionBuilder &builder) { + auto values_alloc = tensors->at(param.outputValues_index).get(); + auto indices_alloc = tensors->at(param.outputIndices_index).get(); + auto input_alloc = tensors->at(param.inputData_index).get(); + + std::unique_ptr<::arm_compute::IFunction> fn; + + auto l = make_layer<::arm_compute::CLTopKV2>(); + + l->configure(input_alloc->handle(), param.k, values_alloc->handle(), indices_alloc->handle()); + + fn = std::move(l); + + auto acl_fn = make_cl_function(std::move(fn)); + + builder.append(std::move(acl_fn)); + }); +} + } // namespace acl_cl } // namespace backend } // namespace neurun diff --git a/runtimes/neurun/backend/acl_cl/StageGenerator.h b/runtimes/neurun/backend/acl_cl/StageGenerator.h index 3a05104c9..5dad5c707 100644 --- a/runtimes/neurun/backend/acl_cl/StageGenerator.h +++ b/runtimes/neurun/backend/acl_cl/StageGenerator.h @@ -76,6 +76,9 @@ public: virtual void visit(const model::operation::SQRTNode &) override; virtual void visit(const model::operation::LogicalOrNode &) override; virtual void visit(const model::operation::LogicalNotNode &) override; + virtual void visit(const model::operation::EqualNode &) override; + virtual void visit(const model::operation::SquaredDifferenceNode &) override; + virtual void visit(const model::operation::TopKV2Node &) override; private: const neurun::model::operand::Set &_ctx; |