"""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)
|
