17#include "../KernelGenerator.h"
18#include "../Validator.h"
25void Validator::visit(
const ir::operation::ArgMinMax &) {
_supported =
true; }
27void KernelGenerator::visit(
const ir::operation::ArgMinMax &node)
29 const auto ofm_index{node.getOutputs().at(0)};
33 auto ifm_shape = _ctx.
at(ifm_index).shape();
34 auto ofm_shape = _ctx.
at(ofm_index).shape();
36 assert((ifm_shape.rank() - 1) == ofm_shape.rank());
38 auto ofm_tensor = _tensor_reg->getAclTensor(ofm_index);
39 auto ifm_tensor = _tensor_reg->getAclTensor(ifm_index);
40 const auto ifm_rank = _ctx.
at(ifm_index).shape().rank();
42 int axis_value = _ctx.
at(axis_index).asScalar<int32_t>();
45 axis_value += ifm_rank;
49 auto reduce_type = node.param().is_arg_max ? ::arm_compute::ReductionOperation::ARG_IDX_MAX
50 : ::arm_compute::ReductionOperation::ARG_IDX_MIN;
51 auto fn = acl_common::generateLayer<arm_compute::CLArgMinMaxLayer>(
52 ifm_tensor->handle(), acl_axis, ofm_tensor->handle(), reduce_type);
uint32_t value(void) const
std::unique_ptr< exec::IFunction > _return_fn
const Object & at(const Index &index) const
Get the object that is associated with the given index.
ARMComputeAxis ToARMComputeAxis(uint32_t rank, uint32_t axis)
std::unique_ptr< AclFunction > asAclFunction(std::unique_ptr<::arm_compute::IFunction > &&layer)