#include <BatchMatMul.h>

Collaboration diagram for luci_interpreter::kernels::BatchMatMul:

[legend]

Public Member Functions
	BatchMatMul (const Tensor x, const Tensor y, Tensor output, Tensor x_tmp, Tensor *y_tmp, const BatchMatMulParams &params)

const Tensor *	x () const

const Tensor *	y () const

Tensor *	output () const

void	configure () override

void	execute () const override

	BatchMatMul (const Tensor x, const Tensor y, Tensor output, Tensor x_tmp, Tensor *y_tmp, const BatchMatMulParams &params)

const Tensor *	x () const

const Tensor *	y () const

Tensor *	output () const

void	configure () override

void	execute () const override

Public Member Functions inherited from luci_interpreter::KernelWithParams< BatchMatMulParams >
const BatchMatMulParams &	params () const

Public Member Functions inherited from luci_interpreter::Kernel
virtual	~Kernel ()=default

const std::vector< const Tensor * > &	getInputTensors () const

const std::vector< Tensor * > &	getOutputTensors () const

Additional Inherited Members
Protected Member Functions inherited from luci_interpreter::KernelWithParams< BatchMatMulParams >
	KernelWithParams (std::vector< const Tensor * > inputs, std::vector< Tensor * > outputs, const BatchMatMulParams &params)

Protected Member Functions inherited from luci_interpreter::Kernel
	Kernel (std::vector< const Tensor * > inputs, std::vector< Tensor * > outputs)

Protected Attributes inherited from luci_interpreter::KernelWithParams< BatchMatMulParams >
const BatchMatMulParams	_params

Protected Attributes inherited from luci_interpreter::Kernel
const std::vector< const Tensor * >	_inputs

const std::vector< Tensor * >	_outputs

Detailed Description

Definition at line 28 of file BatchMatMul.h.

Constructor & Destructor Documentation

◆ BatchMatMul() [1/2]

luci_interpreter::kernels::BatchMatMul::BatchMatMul	(	const Tensor *	x,
		const Tensor *	y,
		Tensor *	output,
		Tensor *	x_tmp,
		Tensor *	y_tmp,
		const BatchMatMulParams &	params
	)

Definition at line 46 of file BatchMatMul.cpp.

  : KernelWithParams({x, y}, {output, x_tmp, y_tmp}, params)
{
}

References x(), and y().

◆ BatchMatMul() [2/2]

luci_interpreter::kernels::BatchMatMul::BatchMatMul	(	const Tensor *	x,
		const Tensor *	y,
		Tensor *	output,
		Tensor *	x_tmp,
		Tensor *	y_tmp,
		const BatchMatMulParams &	params
	)

Member Function Documentation

◆ configure() [1/2]

void luci_interpreter::kernels::BatchMatMul::configure ( )

overridevirtual

Implements luci_interpreter::Kernel.

Definition at line 52 of file BatchMatMul.cpp.

{
  auto lhs = x();
  auto rhs = y();
  auto adj_x = params().adj_x;
  auto adj_y = params().adj_y;
 
  // TODO Support non-float types
  if (lhs->element_type() != DataType::FLOAT32 || rhs->element_type() != DataType::FLOAT32)
    throw std::runtime_error("luci-intp BatchMatMul(1) Unsupported type.");
 
  LUCI_INTERPRETER_CHECK(lhs->element_type() == rhs->element_type());
 
  auto lhs_rank = lhs->shape().num_dims();
  auto rhs_rank = rhs->shape().num_dims();
  LUCI_INTERPRETER_CHECK(lhs_rank >= 2 && lhs_rank <= 4);
  LUCI_INTERPRETER_CHECK(rhs_rank >= 2 && rhs_rank <= 4);
 
  auto lhs_scratchpad = temp_lhs();
  auto rhs_scratchpad = temp_rhs();
  luci_interpreter_pal::SetupScratchpadTensor(lhs_scratchpad, rhs_scratchpad, getTensorShape(lhs),
                                              getTensorShape(rhs));
 
  auto output_rank = std::max(lhs_rank, rhs_rank);
 
  auto extended_lhs_shape = tflite::RuntimeShape::ExtendedShape(output_rank, getTensorShape(lhs));
  auto extended_rhs_shape = tflite::RuntimeShape::ExtendedShape(output_rank, getTensorShape(rhs));
 
  // Ensure any batch dimensions obey broacasting rules.
  for (int i = 0; i < output_rank - 2; ++i)
  {
    const int lhs_dim = extended_lhs_shape.Dims(i);
    const int rhs_dim = extended_rhs_shape.Dims(i);
    if (lhs_dim != rhs_dim)
    {
      if (lhs_dim != 1)
      {
        LUCI_INTERPRETER_CHECK(rhs_dim == 1);
      }
    }
  }
 
  // Ensure other dimensions work for matrix multiplication.
  int accum_dim_lhs =
    adj_x ? extended_lhs_shape.Dims(output_rank - 2) : extended_lhs_shape.Dims(output_rank - 1);
  int accum_dim_rhs =
    adj_y ? extended_rhs_shape.Dims(output_rank - 1) : extended_rhs_shape.Dims(output_rank - 2);
  LUCI_INTERPRETER_CHECK(accum_dim_lhs == accum_dim_rhs);
 
  Shape output_shape(output_rank);
  // Fill in any broadcast dimensions.
  for (int i = 0; i < output_rank - 2; ++i)
  {
    const int lhs_dim = extended_lhs_shape.Dims(i);
    const int rhs_dim = extended_rhs_shape.Dims(i);
    int broadcast_dim = lhs_dim;
    if ((lhs_dim != rhs_dim) && (lhs_dim == 1))
    {
      broadcast_dim = rhs_dim;
    }
    output_shape.dim(i) = broadcast_dim;
  }
  // Fill in the matmul dimensions.
  int lhs_rows_index = adj_x ? output_rank - 1 : output_rank - 2;
  int rhs_cols_index = adj_y ? output_rank - 2 : output_rank - 1;
 
  output_shape.dim(output_rank - 2) = extended_lhs_shape.Dims(lhs_rows_index);
  output_shape.dim(output_rank - 1) = extended_rhs_shape.Dims(rhs_cols_index);
 
  output()->resize(output_shape);
}

References luci_interpreter::BatchMatMulParams::adj_x, luci_interpreter::BatchMatMulParams::adj_y, luci_interpreter::kernels::getTensorShape(), LUCI_INTERPRETER_CHECK, luci::must_cast(), output(), output_shape, luci_interpreter::KernelWithParams< BatchMatMulParams >::params(), luci_interpreter::Tensor::resize(), x(), and y().

◆ configure() [2/2]

void luci_interpreter::kernels::BatchMatMul::configure ( )

overridevirtual

Implements luci_interpreter::Kernel.

◆ execute() [1/2]

void luci_interpreter::kernels::BatchMatMul::execute ( ) const

overridevirtual

Implements luci_interpreter::Kernel.

Definition at line 151 of file BatchMatMul.cpp.

{
  auto lhs = x();
  auto rhs = y();
 
  bool adj_x = params().adj_x;
  bool adj_y = params().adj_y;
 
  auto orig_lhs_shape = getTensorShape(lhs);
  auto orig_rhs_shape = getTensorShape(rhs);
 
  auto rhs_tensor = adj_y ? rhs : temp_rhs();
  auto lhs_tensor = adj_x ? temp_lhs() : lhs;
  if (not adj_y)
  {
    TransposeRowsColumns(rhs, temp_rhs());
  }
  if (adj_x)
  {
    TransposeRowsColumns(lhs, temp_lhs());
  }
  tflite::RuntimeShape rhs_shape = adj_y ? orig_rhs_shape : SwapRowColumnDims(orig_rhs_shape);
  tflite::RuntimeShape lhs_shape = adj_x ? orig_lhs_shape : SwapRowColumnDims(orig_lhs_shape);
 
  switch (x()->element_type())
  {
    case DataType::FLOAT32:
      luci_interpreter_pal::BatchMatMul(rhs_shape, getTensorData<float>(rhs_tensor), lhs_shape,
                                        getTensorData<float>(lhs_tensor), getTensorShape(output()),
                                        getTensorData<float>(output()));
      break;
    default:
      throw std::runtime_error("luci-intp BatchMatMul(2) Unsupported type.");
  }
}