#include <NEEmbeddingLookupKernel.h>

Collaboration diagram for arm_compute::NEEmbeddingLookupKernel:

Public Member Functions
const char *	name () const override

	NEEmbeddingLookupKernel ()

	NEEmbeddingLookupKernel (const NEEmbeddingLookupKernel &)=delete

NEEmbeddingLookupKernel &	operator= (const NEEmbeddingLookupKernel &)=delete

	NEEmbeddingLookupKernel (NEEmbeddingLookupKernel &&)=default

NEEmbeddingLookupKernel &	operator= (NEEmbeddingLookupKernel &&)=default

void	configure (const ITensor input, ITensor output, const ITensor *lookups)

void	run (const Window &window, const ThreadInfo &info) override

Static Public Member Functions
static Status	validate (const ITensorInfo input, const ITensorInfo output, const ITensorInfo *lookups)

Detailed Description

NEON kernel to perform EmbeddingLookup operation

Definition at line 52 of file NEEmbeddingLookupKernel.h.

Constructor & Destructor Documentation

◆ NEEmbeddingLookupKernel() [1/3]

NEEmbeddingLookupKernel::NEEmbeddingLookupKernel ( )

Default constructor

Definition at line 55 of file NEEmbeddingLookupKernel.cpp.

  : _input(nullptr), _lookups(nullptr), _output(nullptr)
{
}

◆ NEEmbeddingLookupKernel() [2/3]

arm_compute::NEEmbeddingLookupKernel::NEEmbeddingLookupKernel ( const NEEmbeddingLookupKernel & )

delete

Prevent instances of this class from being copied (As this class contains pointers).

◆ NEEmbeddingLookupKernel() [3/3]

arm_compute::NEEmbeddingLookupKernel::NEEmbeddingLookupKernel ( NEEmbeddingLookupKernel && )

default

Allow instances of this class to be moved

Member Function Documentation

◆ configure()

void NEEmbeddingLookupKernel::configure	(	const ITensor *	input,
		ITensor *	output,
		const ITensor *	lookups
	)

Initialize the kernel's input, output.

Parameters

[in]	input	Source tensor. Data types supported: U8/S8/QASYMM8/U16/S16/F16/U32/S32/F32.
[out]	output	Destination tensor. Data types supported: same as `input`.
[in]	lookups	Lookups are 1D tensor that values are indices into the first dimension of input.

Definition at line 60 of file NEEmbeddingLookupKernel.cpp.

{
  ARM_COMPUTE_ERROR_ON_NULLPTR(input, output);
  ARM_COMPUTE_ERROR_THROW_ON(validate(input->info(), output->info(), lookups->info()));
 
  _input = input;
  _output = output;
  _lookups = lookups;
 
  // Auto initialize output if not initialized
  auto out_shape = input->info()->tensor_shape();
  out_shape.set(out_shape.num_dimensions() - 1, lookups->info()->num_dimensions());
  auto_init_if_empty(*output->info(), out_shape, 1, input->info()->data_type(),
                     input->info()->quantization_info());
 
  INEKernel::configure(calculate_max_window(*output->info()));
}

References validate().

◆ name()

const char * arm_compute::NEEmbeddingLookupKernel::name ( ) const

inlineoverride

Definition at line 55 of file NEEmbeddingLookupKernel.h.

55{ return "NEEmbeddingLookupKernel"; }

◆ operator=() [1/2]

NEEmbeddingLookupKernel & arm_compute::NEEmbeddingLookupKernel::operator= ( const NEEmbeddingLookupKernel & )

delete

Prevent instances of this class from being copied (As this class contains pointers).

◆ operator=() [2/2]

NEEmbeddingLookupKernel & arm_compute::NEEmbeddingLookupKernel::operator= ( NEEmbeddingLookupKernel && )

default

Allow instances of this class to be moved

References validate().

◆ run()

void NEEmbeddingLookupKernel::run	(	const Window &	window,
		const ThreadInfo &	info
	)

override

Definition at line 107 of file NEEmbeddingLookupKernel.cpp.

{
  ARM_COMPUTE_UNUSED(info);
  ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
  ARM_COMPUTE_ERROR_ON_INVALID_SUBWINDOW(INEKernel::window(), window);
 
  const size_t lookup_dim = _output->info()->num_dimensions() - 1;
 
  Window output_window{window};
  output_window.set(Window::DimX,
                    Window::Dimension(output_window.x().start(), output_window.x().end(),
                                      _input->info()->dimension(0)));
 
  Window out_slice = output_window.first_slice_window_4D();
  do
  {
    Iterator output_it(_output, out_slice);
 
    execute_window_loop(
      out_slice,
      [&](const Coordinates &id) {
        const int32_t lookup =
          *reinterpret_cast<int32_t *>(_lookups->ptr_to_element(Coordinates{id[lookup_dim]}));
        Coordinates input_id{id};
        input_id.set(lookup_dim, lookup);
        memcpy(output_it.ptr(), _input->ptr_to_element(input_id),
               _output->info()->dimension(0) * _output->info()->element_size());
      },
      output_it);
 
  } while (window.slide_window_slice_4D(out_slice));
}

References info.

Referenced by package.infer.session::inference().

◆ validate()

Status NEEmbeddingLookupKernel::validate	(	const ITensorInfo *	input,
		const ITensorInfo *	output,
		const ITensorInfo *	lookups
	)

static

Static function to check if given info will lead to a valid configuration of NEEmbeddingLookupKernel