def scalar_affine(
op: ScalarAffine,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
x = variables_layout[op.inputs["x"]]
y = variables_layout[op.outputs["y"]]
assert x.variable.shape == y.variable.shape
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
metabuffer_injector.register({
"affine_transform_X_offset": x.offset,
"affine_transform_Y_offset": y.offset,
"affine_transform_N": y.variable.size,
"affine_transform_scale": op.scale,
"affine_transform_bias": op.bias
})
source = metabuffer_injector.inject(template)
func_name = util.add_canonical_suffix("scalar_affine", source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel({func_name: source}, func_name,
metabuffer_injector.generate_buffer())
return [kernel]
def tanh(op: Tanh,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
x = variables_layout[op.inputs["x"]]
y = variables_layout[op.outputs["y"]]
assert x.variable.order == y.variable.order
assert x.variable.shape == y.variable.shape
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
metabuffer_injector.register({
"relu_X_offset": x.offset,
"relu_Y_offset": y.offset,
"relu_N": y.variable.size
})
source = metabuffer_injector.inject(template)
func_name = util.add_canonical_suffix("tanh", source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel({func_name: source}, func_name,
metabuffer_injector.generate_buffer())
return [kernel]
def concat(op: Concat,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
xs = [variables_layout[op.inputs[f"x{str(i)}"]] for i in range(len(op.inputs))]
y = variables_layout[op.outputs["y"]]
target_axis = op.axis
x_offsets = [x.offset for x in xs]
x_shapes = [x.variable.shape for x in xs]
y_strides = []
stride = 1
for s in reversed(y.variable.shape):
y_strides.insert(0, stride)
stride *= s
# x_strides[i][j] is stride size of xs[i].order.axes[j] in y
x_strides_in_y = [[] for _ in xs]
for x, strides in zip(xs, x_strides_in_y):
for axis in x.variable.order.axes:
strides.append(y_strides[y.variable.order.axes_dict[axis]])
# x_offsets[i] is memory offset of xs[i]'s data in y.
y_offsets = []
target_axis_offset = 0
for x in xs:
y_offsets.append(target_axis_offset * y_strides[y.variable.order.axes_dict[target_axis]])
target_axis_offset += x.variable.shape_dict[target_axis]
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
metabuffer_injector.register({
"concat_y_offset": y.offset,
"concat_D": len(y.variable.shape),
"concat_N": len(xs),
"concat_x_offsets": np.array(x_offsets, dtype=np.int32).tobytes(),
"concat_x_strides_in_y": np.array(x_strides_in_y, dtype=np.int32).tobytes(),
"concat_x_shapes": np.array(x_shapes, dtype=np.int32).tobytes(),
"concat_y_offsets": np.array(y_offsets, dtype=np.int32).tobytes(),
})
source = metabuffer_injector.inject(template)
func_name = util.add_canonical_suffix("concat", source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel(
{func_name: source},
func_name,
metabuffer_injector.generate_buffer()
)
return [kernel]
def sgemm(op: Sgemm,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
A = variables_layout[op.inputs["A"]] if op.inputs[
"A"] in variables_layout else constants_layout[op.inputs["A"]]
B = variables_layout[op.inputs["B"]] if op.inputs[
"B"] in variables_layout else constants_layout[op.inputs["B"]]
C = variables_layout[op.outputs["C"]]
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
metabuffer_injector.register({
"sgemm_A_offset": A.offset,
"sgemm_B_offset": B.offset,
"sgemm_C_offset": C.offset,
"sgemm_M": op.M,
"sgemm_N": op.N,
"sgemm_K": op.K
})
if op.parameters["eigen"]:
source = generate_template_eigen(op.transpose_A, op.transpose_B, op.M,
op.N, op.K)
metabuffer_injector.register({
"sgemm_A_offset": A.offset,
"sgemm_B_offset": B.offset,
"sgemm_C_offset": C.offset
})
else:
source = generate_template(op.transpose_A, op.transpose_B)
metabuffer_injector.register({
"sgemm_A_offset": A.offset,
"sgemm_B_offset": B.offset,
"sgemm_C_offset": C.offset,
"sgemm_M": op.M,
"sgemm_N": op.N,
"sgemm_K": op.K
})
source = metabuffer_injector.inject(source)
func_name = util.add_canonical_suffix("sgemm", source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel({func_name: source}, func_name,
metabuffer_injector.generate_buffer())
return [kernel]
def average_pooling_2d(
op: AveragePooling2D,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
x = variables_layout[op.inputs["x"]]
y = variables_layout[op.outputs["y"]]
assert x.variable.order == OrderNHWC
assert y.variable.order == OrderNHWC
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
metabuffer_injector.register({
"average_pooling_2d_X_offset":
x.offset,
"average_pooling_2d_Y_offset":
y.offset,
"average_pooling_2d_N":
x.variable.shape_dict[Axis.N],
"average_pooling_2d_H1":
x.variable.shape_dict[Axis.H],
"average_pooling_2d_W1":
x.variable.shape_dict[Axis.W],
"average_pooling_2d_C":
x.variable.shape_dict[Axis.C],
"average_pooling_2d_H2":
y.variable.shape_dict[Axis.H],
"average_pooling_2d_W2":
y.variable.shape_dict[Axis.W],
"average_pooling_2d_K":
op.parameters["ksize"][0],
"average_pooling_2d_S":
op.parameters["stride"][0],
"average_pooling_2d_P":
op.parameters["padding"][0],
})
source = metabuffer_injector.inject(template)
func_name = util.add_canonical_suffix("average_pooling_2d", source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel({func_name: source}, func_name,
metabuffer_injector.generate_buffer())
return [kernel]
def local_response_normalization(
op: LocalResponseNormalization,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
x = variables_layout[op.inputs["x"]]
y = variables_layout[op.outputs["y"]]
assert x.variable.order == OrderNHWC
assert y.variable.order == OrderNHWC
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
metabuffer_injector.register({
"local_response_normalization_X_offset":
x.offset,
"local_response_normalization_Y_offset":
y.offset,
"local_response_normalization_N":
x.variable.shape_dict[Axis.N],
"local_response_normalization_H":
x.variable.shape_dict[Axis.H],
"local_response_normalization_W":
x.variable.shape_dict[Axis.W],
"local_response_normalization_C":
x.variable.shape_dict[Axis.C],
"local_response_normalization_param_half_n":
int(op.parameters["n"] // 2),
"local_response_normalization_param_k":
float(op.parameters["k"]),
"local_response_normalization_param_alpha":
float(op.parameters["alpha"]),
"local_response_normalization_param_minus_beta":
float(-op.parameters["beta"])
})
source = metabuffer_injector.inject(template)
func_name = util.add_canonical_suffix("local_response_normalization",
source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel({func_name: source}, func_name,
metabuffer_injector.generate_buffer())
return [kernel]
def axiswise_bias(
op: AxiswiseBias,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
x = variables_layout[op.inputs["x"]]
b = constants_layout[op.inputs["b"]]
y = variables_layout[op.outputs["y"]]
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
assert x.variable.order == OrderNC or x.variable.order == OrderNHWC or x.variable.order == OrderHWNC
assert y.variable.shape == x.variable.shape
assert op.parameters[
"axis"] == Axis.C, "[Webassembly] AxiswiseBias supports only channelwise bias."
metabuffer_injector.register({
"axiswise_bias_X_offset":
x.offset,
"axiswise_bias_Y_offset":
y.offset,
"axiswise_bias_B_offset":
b.offset,
"axiswise_bias_N":
y.variable.size // y.variable.shape_dict[Axis.C],
"axiswise_bias_C":
y.variable.shape_dict[Axis.C],
})
inline_injector = InlineInjector()
if "inline_elementwise" in op.parameters:
inline_injector.delegate = op.parameters["inline_elementwise"]
source = template
source = metabuffer_injector.inject(source)
source = inline_injector.inject(source)
func_name = util.add_canonical_suffix("axiswise_bias", source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel({func_name: source}, func_name,
metabuffer_injector.generate_buffer())
return [kernel]
def im2col(op: Im2Col,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
im = variables_layout[op.inputs["im"]]
col = variables_layout[op.outputs["col"]]
assert im.variable.order == OrderNHWC
assert col.variable.order == OrderNHWC or col.variable.order == OrderCNHW
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
metabuffer_injector.register({
"im2col_im_offset": im.offset,
"im2col_col_offset": col.offset,
"im2col_N": col.variable.shape_dict[Axis.N],
"im2col_C1": im.variable.shape_dict[Axis.C],
"im2col_H1": im.variable.shape_dict[Axis.H],
"im2col_W1": im.variable.shape_dict[Axis.W],
"im2col_H2": col.variable.shape_dict[Axis.H],
"im2col_W2": col.variable.shape_dict[Axis.W],
"im2col_KH": op.KH,
"im2col_KW": op.KW,
"im2col_SH": op.SH,
"im2col_SW": op.SW,
"im2col_PH": op.PH,
"im2col_PW": op.PW,
})
source = template_CNHW if col.variable.order == OrderCNHW else template_NHWC
source = metabuffer_injector.inject(source)
func_name = util.add_canonical_suffix("im2col", source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel(
{func_name: source},
func_name,
metabuffer_injector.generate_buffer()
)
return [kernel]
def linear(op: Linear,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
x = variables_layout[op.inputs["x"]]
w = constants_layout[op.inputs["w"]]
y = variables_layout[op.outputs["y"]]
assert x.variable.order == OrderNC or x.variable.order == OrderNHWC
assert w.variable.order == OrderCN or w.variable.order == OrderHWCN
assert y.variable.order == OrderNC or y.variable.order == OrderNHWC
assert w.variable.ndim == x.variable.ndim
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
metabuffer_injector.register({
"linear_X_offset":
x.offset,
"linear_Y_offset":
y.offset,
"linear_W_offset":
w.offset,
"linear_M":
y.variable.shape_dict[Axis.N],
"linear_N":
y.variable.size // y.variable.shape_dict[Axis.N],
"linear_K":
x.variable.size // x.variable.shape_dict[Axis.N],
})
source = metabuffer_injector.inject(template)
func_name = util.add_canonical_suffix("linear", source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel({func_name: source}, func_name,
metabuffer_injector.generate_buffer())
return [kernel]
def elementwise_sum(op: AxiswiseScale,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
x0 = variables_layout[op.inputs["x0"]]
x1 = variables_layout[op.inputs["x1"]]
y = variables_layout[op.outputs["y"]]
assert len(op.inputs) == 2, "[Webassembly] ElementwiseSum operator currently supported only 2 inputs."
assert x0.variable.shape == x1.variable.shape == y.variable.shape
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
metabuffer_injector.register({
"elementwise_sum_X0_offset": x0.offset,
"elementwise_sum_X1_offset": x1.offset,
"elementwise_sum_Y_offset": y.offset,
"elementwise_sum_N": y.variable.size
})
inline_injector = InlineInjector()
if "inline_elementwise" in op.parameters:
inline_injector.delegate = op.parameters["inline_elementwise"]
source = template
source = metabuffer_injector.inject(source)
source = inline_injector.inject(source)
func_name = util.add_canonical_suffix("elementwise_sum", source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel(
{func_name: source},
func_name,
metabuffer_injector.generate_buffer()
)
return [kernel]
def axiswise_scale(
op: AxiswiseScale,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
x = variables_layout[op.inputs["x"]]
s = constants_layout[op.inputs["s"]]
y = variables_layout[op.outputs["y"]]
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
assert x.variable.order == OrderNC or x.variable.order == OrderNHWC or x.variable.order == OrderHWNC
assert y.variable.order == OrderNC or y.variable.order == OrderNHWC or y.variable.order == OrderHWNC
assert op.parameters[
"axis"] == Axis.C, "[Webassembly] AxiswiseScale supports only channelwise bias."
metabuffer_injector.register({
"axiswise_scale_X_offset":
x.offset,
"axiswise_scale_Y_offset":
y.offset,
"axiswise_scale_S_offset":
s.offset,
"axiswise_scale_N":
y.variable.size,
"axiswise_scale_C":
y.variable.shape_dict[Axis.C],
})
source = metabuffer_injector.inject(template)
func_name = util.add_canonical_suffix("axiswise_scale", source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel({func_name: source}, func_name,
metabuffer_injector.generate_buffer())
return [kernel]
def flatten(op: Flatten,
constants_layout: MemoryLayout,
variables_layout: MemoryLayout,
metabuffer_injector: MetaBufferInjector = None) -> List[Kernel]:
x = variables_layout[op.inputs["x"]]
y = variables_layout[op.outputs["y"]]
if metabuffer_injector is None:
metabuffer_injector = MetaBufferInjector()
metabuffer_injector.register({
"flatten_x_offset": x.offset,
"flatten_y_offset": y.offset,
"flatten_N": y.variable.size,
})
source = metabuffer_injector.inject(template)
func_name = util.add_canonical_suffix("flatten", source)
source = source.replace("%%FUNC_NAME%%", func_name)
kernel = Kernel({func_name: source}, func_name,
metabuffer_injector.generate_buffer())
return [kernel]
