xref: /external/gemmlowp/meta/generators/mul_Nx8_Mx8_neon.py

"""Mul primitive used by the GEMM function.

The Mul primitive takes 1-3 zipped rows and 1-3 zipped columns and performs
matrix multiplication on those resulting in a small 1x1 to 3x3 block of results.
"""

import neon_emitter


class Error(Exception):
  """Module level error."""


class ConfigurationError(Error):
  """Unsupported configuration."""


class MulLanes(object):

  def __init__(self, input_address):
    self.input_address = input_address
    self.lanes = []

  def AddLane(self, lane):
    self.lanes.append(lane)

  def FreeRegisters(self, registers):
    for i in range(0, len(self.lanes)):
      registers.FreeRegister(self.lanes[i])
      self.lanes[i] = None


def GenerateMulLanes(registers, lane_count, address):
  lanes = MulLanes(address)
  for unused_i in range(0, lane_count):
    lanes.AddLane(registers.DoubleRegister())
  return lanes


def Generate3MulLanes(quad_register, registers, address):
  lanes = MulLanes(address)
  lanes.AddLane(registers.Low(quad_register))
  lanes.AddLane(registers.High(quad_register))
  lanes.AddLane(registers.DoubleRegister())
  return lanes


def GenerateAndClearAggregators(emitter, registers, aggregator_count):
  """Prepare aggregators and emit aggregator clear code."""
  emitter.EmitComment('Clear aggregators.')
  aggregators = []
  for i in range(0, aggregator_count):
    aggregator = registers.QuadRegister()
    aggregators.append(aggregator)
    if i < 3:
      emitter.EmitVMov('i32', aggregator, emitter.ImmediateConstant(0))
    else:
      emitter.EmitVMov('i32', aggregator, aggregators[i - 3])
  emitter.EmitNewline()
  return aggregators


def GenerateNxMLoadMultiplyAggregate(emitter, registers, left_lanes,
                                     right_lanes, aggregators, count):
  """Emit inner loop for N rows x M cols multiplication."""
  emitter.EmitComment('General NxM lanes loop.')
  emitter.EmitNumericalLabel(1)
  emitter.EmitNewline()
  emitter.EmitComment('Subtract counter.')
  emitter.EmitSubs(count, count, emitter.ImmediateConstant(8))
  emitter.EmitNewline()

  emitter.EmitVLoadA('1.8', left_lanes.lanes,
                     emitter.DereferenceIncrement(left_lanes.input_address, 64))
  emitter.EmitVLoadA(
      '1.8', right_lanes.lanes,
      emitter.DereferenceIncrement(right_lanes.input_address, 64))

  emitter.EmitPldOffset(left_lanes.input_address, emitter.ImmediateConstant(64))
  emitter.EmitPldOffset(right_lanes.input_address,
                        emitter.ImmediateConstant(64))

  rows = len(left_lanes.lanes)
  cols = len(right_lanes.lanes)

  multiply_results = []
  for i in range(0, rows * cols):
    multiply_results.append(registers.QuadRegister())

  for row in range(0, rows):
    for col in range(0, cols):
      index = row * cols + col
      emitter.EmitVMull('u8', multiply_results[index], right_lanes.lanes[col],
                        left_lanes.lanes[row])

  for i in range(0, rows * cols):
    emitter.EmitVPadal('u16', aggregators[i], multiply_results[i])

  emitter.EmitNewline()
  emitter.EmitComment('Loop break.')
  emitter.EmitBneBack(1)
  emitter.EmitNewline()

  for register in multiply_results:
    registers.FreeRegister(register)


def Generate3x3LoadMultiplyAggregate(emitter, registers, left_lanes,
                                     right_lanes, aggregators, count,
                                     backup_register):
  """Emit inner loop for 3 rows x 3 cols multiplication (register trick)."""
  emitter.EmitComment('3x3 lanes loop.')
  emitter.EmitNumericalLabel(1)
  emitter.EmitNewline()
  emitter.EmitComment('Subtract counter.')
  emitter.EmitSubs(count, count, emitter.ImmediateConstant(8))
  emitter.EmitNewline()

  emitter.EmitVLoadA('1.8', left_lanes.lanes,
                     emitter.DereferenceIncrement(left_lanes.input_address, 64))
  emitter.EmitVLoadA(
      '1.8', right_lanes.lanes,
      emitter.DereferenceIncrement(right_lanes.input_address, 64))

  emitter.EmitPldOffset(left_lanes.input_address, emitter.ImmediateConstant(64))
  emitter.EmitPldOffset(right_lanes.input_address,
                        emitter.ImmediateConstant(64))

  temp = []
  for unused_i in range(0, 4):
    temp.append(registers.QuadRegister())

  emitter.EmitVMull('u8', temp[0], left_lanes.lanes[0], right_lanes.lanes[0])
  emitter.EmitVMull('u8', temp[1], left_lanes.lanes[0], right_lanes.lanes[1])
  emitter.EmitVMull('u8', temp[2], left_lanes.lanes[0], right_lanes.lanes[2])
  emitter.EmitVMull('u8', temp[3], left_lanes.lanes[1], right_lanes.lanes[0])

  emitter.EmitVPadal('u16', aggregators[0], temp[0])
  emitter.EmitVPadal('u16', aggregators[1], temp[1])
  emitter.EmitVPadal('u16', aggregators[2], temp[2])
  emitter.EmitVPadal('u16', aggregators[3], temp[3])

  emitter.EmitVMull('u8', temp[0], left_lanes.lanes[1], right_lanes.lanes[1])
  emitter.EmitVMull('u8', temp[1], left_lanes.lanes[1], right_lanes.lanes[2])
  emitter.EmitVMull('u8', temp[2], left_lanes.lanes[2], right_lanes.lanes[0])
  emitter.EmitVMull('u8', temp[3], left_lanes.lanes[2], right_lanes.lanes[1])
  emitter.EmitVMull('u8', backup_register, left_lanes.lanes[2],
                    right_lanes.lanes[2])

  emitter.EmitVPadal('u16', aggregators[4], temp[0])
  emitter.EmitVPadal('u16', aggregators[5], temp[1])
  emitter.EmitVPadal('u16', aggregators[6], temp[2])
  emitter.EmitVPadal('u16', aggregators[7], temp[3])
  emitter.EmitVPadal('u16', aggregators[8], backup_register)

  emitter.EmitNewline()
  emitter.EmitComment('Loop break.')
  emitter.EmitBneBack(1)
  emitter.EmitNewline()

  for register in temp:
    registers.FreeRegister(register)


def ReadParams(emitter, registers, input_address, elements, min_reg):
  if elements == 1 or elements == 2:
    register = registers.DoubleRegister(min_reg * 2)
    emitter.EmitVLoad('1.32', register, emitter.Dereference(input_address, 64))
    return register
  elif elements == 3 or elements == 4:
    register = registers.QuadRegister(min_reg)
    emitter.EmitVLoad('1.32', register, emitter.Dereference(input_address, 64))
    return register
  else:
    raise ConfigurationError('Unsupported elements no: %d' % elements)


def Duplicate(emitter, registers, rows, cols, min_register, values):
  """Populate a grid of registers duplicating provided values."""
  duplicated = []
  if cols == 1 or cols == 2:
    for unused_i in range(0, rows):
      duplicated.append(registers.DoubleRegister(min_register))
  elif cols == 3 or cols == 4:
    for unused_i in range(0, rows):
      duplicated.append(registers.QuadRegister(min_register))
  else:
    raise ConfigurationError('Unsupported duplicate amount: %d' % cols)

  if rows == 1:
    emitter.EmitVDup('32', duplicated[0], emitter.Lane(values, 0))
  elif rows == 2:
    emitter.EmitVDup('32', duplicated[0], emitter.Lane(values, 0))
    emitter.EmitVDup('32', duplicated[1], emitter.Lane(values, 1))
  elif rows == 3:
    emitter.EmitVDup('32', duplicated[0], emitter.Lane(
        registers.Low(values), 0))
    emitter.EmitVDup('32', duplicated[1], emitter.Lane(
        registers.Low(values), 1))
    emitter.EmitVDup('32', duplicated[2], emitter.Lane(
        registers.High(values), 0))
  elif rows == 4:
    emitter.EmitVDup('32', duplicated[0], emitter.Lane(
        registers.Low(values), 0))
    emitter.EmitVDup('32', duplicated[1], emitter.Lane(
        registers.Low(values), 1))
    emitter.EmitVDup('32', duplicated[2], emitter.Lane(
        registers.High(values), 0))
    emitter.EmitVDup('32', duplicated[3], emitter.Lane(
        registers.High(values), 1))

  return duplicated


def DuplicateGeneralRegister(emitter, registers, cols, general_register,
                             min_register):
  if cols == 1 or cols == 2:
    duplicated = registers.DoubleRegister(min_register)
  elif cols == 3 or cols == 4:
    duplicated = registers.QuadRegister(min_register)
  else:
    raise ConfigurationError('Unsupported duplicate amount: %d' % cols)

  emitter.EmitVDup('32', duplicated, general_register)
  return duplicated


def ReduceAggregator(emitter, registers, aggregators, row, cols):
  if cols == 1:
    register = registers.Low(aggregators[row])
    emitter.EmitVPadd('u32', register, register, register)
    return register
  elif cols == 2:
    register = registers.Low(aggregators[row * 2])
    emitter.EmitVPadd('u32', register, register,
                      registers.Low(aggregators[row * 2 + 1]))
    return register
  elif cols == 3:
    register = aggregators[row * 3]
    emitter.EmitVPadd('u32', registers.Low(register), registers.Low(register),
                      registers.Low(aggregators[row * 3 + 1]))
    emitter.EmitVPadd('u32', registers.High(register),
                      registers.Low(aggregators[row * 3 + 2]),
                      registers.Low(aggregators[row * 3 + 2]))
    return register
  elif cols == 4:
    register = aggregators[row * 3]
    emitter.EmitVPadd('u32', registers.Low(register), registers.Low(register),
                      registers.Low(aggregators[row * 3 + 1]))
    emitter.EmitVPadd('u32', registers.High(register),
                      registers.Low(aggregators[row * 3 + 2]),
                      registers.Low(aggregators[row * 3 + 3]))
    return register
  else:
    raise ConfigurationError('Unsupported columns no: %d' % cols)


def StoreAggregator(emitter, registers, aggregator, cols, result_address,
                    result_stride):
  if cols == 1:
    emitter.EmitVStoreOffset('1.32', emitter.Lane(aggregator, 0),
                             emitter.Dereference(result_address, None),
                             result_stride)
  elif cols == 2:
    emitter.EmitVStoreOffset('1.32', aggregator,
                             emitter.Dereference(result_address, None),
                             result_stride)
  elif cols == 3:
    emitter.EmitVStore('1.32', registers.Low(aggregator),
                       emitter.DereferenceIncrement(result_address, None))
    emitter.EmitVStoreOffset('1.32', emitter.Lane(
        registers.High(aggregator),
        0), emitter.Dereference(result_address, None), result_stride)
    emitter.EmitNewline()
  elif cols == 4:
    emitter.EmitVStoreOffsetA(
        '1.32', [registers.Low(aggregator), registers.High(aggregator)],
        emitter.Dereference(result_address, None), result_stride)
  else:
    raise ConfigurationError('Unsupported columns no: %d' % cols)


def GenerateAggregatorReduceStore(emitter, registers, aggregators, result_type,
                                  lhs_add, rhs_add, left_lanes, right_lanes,
                                  results, results_stride):
  """Emit code that reduces 4 lane aggregators to 1 value, and stores them."""
  rows = len(left_lanes.lanes)
  cols = len(right_lanes.lanes)

  if lhs_add:
    left_offset = ReadParams(emitter, registers, left_lanes.input_address, rows,
                             4)
    left_offsets = Duplicate(emitter, registers, rows, cols, 4, left_offset)
  else:
    left_offsets = None

  if rhs_add:
    right_offset = ReadParams(emitter, registers, right_lanes.input_address,
                              cols, 4)
  else:
    right_offset = None

  if result_type is 'float':
    result_scale = DuplicateGeneralRegister(
        emitter, registers, cols, registers.MapParameter('result_scale'), 4)
  else:
    result_scale = None

  if cols == 3:
    emitter.EmitNewline()
    emitter.EmitComment('Change stride because storing in two ops.')
    emitter.EmitSub(results_stride, results_stride,
                    emitter.ImmediateConstant(8))

  emitter.EmitNewline()
  emitter.EmitComment('Horizontal reduce aggregators.')
  for aggregator in aggregators:
    emitter.EmitVPadd('u32', registers.Low(aggregator),
                      registers.Low(aggregator), registers.High(aggregator))

  emitter.EmitNewline()
  emitter.EmitComment('Reduce rows.')
  row_temps = []
  for i in range(0, rows):
    row_temps.append(ReduceAggregator(emitter, registers, aggregators, i, cols))

  if lhs_add:
    emitter.EmitNewline()
    emitter.EmitComment('Add lhs offsets to aggregated rows.')
    for (row_temp, left_offset) in zip(row_temps, left_offsets):
      emitter.EmitVAdd('s32', row_temp, row_temp, left_offset)

  if rhs_add:
    emitter.EmitNewline()
    emitter.EmitComment('Add rhs offset to aggregated rows.')
    for row_temp in row_temps:
      emitter.EmitVAdd('s32', row_temp, row_temp, right_offset)

  if result_type is 'float':
    emitter.EmitNewline()
    emitter.EmitComment('Convert to float. Multiply by result scale.')
    for row_temp in row_temps:
      emitter.EmitVCvt('f32', 's32', row_temp, row_temp)
    for row_temp in row_temps:
      emitter.EmitVMul('f32', row_temp, row_temp, result_scale)

  emitter.EmitNewline()
  emitter.EmitComment('Store reduced rows.')
  for row_temp in row_temps:
    StoreAggregator(emitter, registers, row_temp, cols, results, results_stride)


def BuildName(result_type, lhs_add, rhs_add, left, right):
  name = 'mul_%dx8_%dx8_%s' % (left, right, result_type)
  if lhs_add:
    name += '_lhsadd'
  if rhs_add:
    name += '_rhsadd'
  return name


def CppResultType(result_type):
  if result_type is 'int32':
    return 'std::int32_t*'
  elif result_type is 'float':
    return 'float*'
  else:
    raise ConfigurationError('Unsupported result type: %s' % result_type)


def GetParameters(result_type):
  params = [['const std::uint8_t*', 'lhs'], ['const std::uint8_t*', 'rhs'],
            ['std::int32_t', 'count'], [CppResultType(result_type), 'result'],
            ['std::int32_t', 'result_stride']]
  if result_type is 'float':
    params.append(['float', 'result_scale'])
  return params


def GenerateMulNx8Mx8(emitter, result_type, lhs_add, rhs_add, left_lanes_count,
                      right_lanes_count):
  """Emit the multiply code for given rows and cols counts."""
  if left_lanes_count < 1 or left_lanes_count > 4:
    raise ConfigurationError('Left_lanes should be: 1, 2, 3 or 4.')
  if right_lanes_count < 1 or right_lanes_count > 4:
    raise ConfigurationError('Right_lanes should be: 1, 2, 3 or 4.')

  emitter.EmitFunctionBeginA(
      BuildName(result_type, lhs_add, rhs_add, left_lanes_count,
                right_lanes_count), GetParameters(result_type), 'inline void')

  emitter.EmitAssert('count % 8 == 0')
  emitter.EmitAssert('count >= 8')
  emitter.EmitAsmBegin()

  registers = neon_emitter.NeonRegisters()

  count = registers.MapParameter('count')

  size = left_lanes_count * right_lanes_count

  lhs = registers.MapParameter('lhs')
  rhs = registers.MapParameter('rhs')

  emitter.EmitPld(lhs)
  emitter.EmitPld(rhs)

  aggregators = GenerateAndClearAggregators(emitter, registers, size)

  if size < 9:
    left_lanes = GenerateMulLanes(registers, left_lanes_count, lhs)
    right_lanes = GenerateMulLanes(registers, right_lanes_count, rhs)

    GenerateNxMLoadMultiplyAggregate(emitter, registers, left_lanes,
                                     right_lanes, aggregators, count)

  else:  # left == 3 and right == 3
    backup_register = registers.QuadRegister()
    left_lanes = Generate3MulLanes(backup_register, registers, lhs)
    right_lanes = GenerateMulLanes(registers, right_lanes_count, rhs)

    Generate3x3LoadMultiplyAggregate(emitter, registers, left_lanes,
                                     right_lanes, aggregators, count,
                                     backup_register)
  left_lanes.FreeRegisters(registers)
  right_lanes.FreeRegisters(registers)

  GenerateAggregatorReduceStore(emitter, registers, aggregators, result_type,
                                lhs_add, rhs_add, left_lanes, right_lanes,
                                registers.MapParameter('result'),
                                registers.MapParameter('result_stride'))

  emitter.EmitAsmEnd(registers.MappedParameters(), [],
                     registers.Clobbers() + ['cc', 'memory'])
  emitter.EmitFunctionEnd()


def GenerateFunctions(emitter, result_type, lhs_add, rhs_add):
  for left_lanes in range(1, 4):
    for right_lanes in range(1, 4):
      GenerateMulNx8Mx8(emitter, result_type, lhs_add, rhs_add, left_lanes,
                        right_lanes)
      emitter.EmitNewline()

  GenerateMulNx8Mx8(emitter, result_type, lhs_add, rhs_add, 1, 4)
  emitter.EmitNewline()


if __name__ == '__main__':
  GenerateFunctions(neon_emitter.NeonEmitter(), 'int32', True, True)