| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| """Core kernel of dot product of 4 Int8 operations""" |
| # pylint: disable=invalid-name |
| import tvm |
| from tvm import te |
| import tvm.target.codegen |
| |
| |
| def dot_16x1x16_uint8_int8_int32(): |
| """Dispatch the most optimized intrin depending on the target""" |
| mcpu = tvm.target.Target.current().mcpu |
| |
| assert mcpu in ( |
| "skylake-avx512", |
| "cascadelake", |
| ), "An old Intel machine that does not have fast Int8 support." |
| if mcpu == "skylake-avx512": |
| return dot_16x1x16_uint8_int8_int32_skylake() |
| # cascadelake |
| return dot_16x1x16_uint8_int8_int32_cascadelake() |
| |
| |
| def dot_16x1x16_uint8_int8_int32_skylake(): |
| """ |
| Int8 dot product by every 4 elements using AVX512 Skylake instructions. |
| This function takes two arrays of uint8 and int8 datatype -- data[4] and |
| kernel[16][4] -- and computes a dot product of data[4] with every |
| 4 elements of kernels, resulting in output[16] of int32 datatype. |
| The pseudo code is as follows. |
| .. code-block:: c |
| void dot_16x1x16_uint8_int8_int32(uint8 data[4], int8 kernel[16][4], |
| int32 output[16]){ |
| for (int i = 0; i < 16; i++){ |
| output[i] = 0; |
| for (int k = 0; k < 4; k++){ |
| output[i] += data[k] * kernel[i][k] |
| } |
| } |
| } |
| |
| Physically, the kernel array sits in an AVX512 vector register and |
| the data[4] is broadcasted to another AVX512 vector register. This |
| function returns a TensorIntrin that can be used to tensorize |
| a schedule. |
| |
| Returns |
| ------- |
| intrin : TensorIntrin |
| The Skylake int8 TensorIntrin that can be used in tensorizing schedule |
| """ |
| |
| int32_lanes = 16 # 16 int32 lanes in AVX512 |
| num_int8_elements = 4 # 4 int8 elements in int32 |
| data = te.placeholder((num_int8_elements,), dtype="uint8", name="data") |
| kernel = te.placeholder((int32_lanes, num_int8_elements), dtype="int8", name="kernel") |
| k = te.reduce_axis((0, num_int8_elements), name="k") |
| C = te.compute( |
| (int32_lanes,), |
| lambda i: te.sum(data[k].astype("int32") * kernel[i, k].astype("int32"), axis=k), |
| name="C", |
| ) |
| |
| a_buffer = tvm.tir.decl_buffer( |
| data.shape, dtype="uint8", name="a_buffer", offset_factor=1, strides=[1] |
| ) |
| b_buffer = tvm.tir.decl_buffer( |
| kernel.shape, dtype="int8", name="b_buffer", offset_factor=1, strides=[te.var("ldw"), 1] |
| ) |
| |
| def _intrin_func(ins, outs): |
| def _instr(index): |
| ib = tvm.tir.ir_builder.create() |
| if index == 1: |
| ib.emit(outs[0].vstore(0, tvm.tir.const(0, "int32x16"))) |
| return ib.get() |
| |
| a_int8 = ins[0].vload([0], "uint8x4") |
| re_int32 = tvm.tir.call_intrin("int32", "tir.reinterpret", a_int8) |
| vec_ai32 = re_int32.astype("int32x16") |
| vec_a = tvm.tir.call_intrin("int8x64", "tir.reinterpret", vec_ai32) |
| vec_b = ins[1].vload([0, 0], "int8x64") |
| vec_one = tvm.tir.const(1, "int16x32") |
| pair_reduction = tvm.tir.call_llvm_pure_intrin( |
| "int16x32", |
| "llvm.x86.avx512.pmaddubs.w.512", |
| tvm.tir.const(0, "uint32"), |
| vec_a, |
| vec_b, |
| ) |
| quad_reduction = tvm.tir.call_llvm_pure_intrin( |
| "int32x16", |
| "llvm.x86.avx512.pmaddw.d.512", |
| tvm.tir.const(0, "uint32"), |
| pair_reduction, |
| vec_one, |
| ) |
| if index == 0: |
| ib.emit(outs[0].vstore(0, quad_reduction)) |
| else: |
| ib.emit(outs[0].vstore(0, quad_reduction + outs[0].vload([0], "int32x16"))) |
| return ib.get() |
| |
| # body, reset, update |
| return _instr(0), _instr(1), _instr(2) |
| |
| buffer_params = {"offset_factor": 1} |
| return te.decl_tensor_intrin( |
| C.op, |
| _intrin_func, |
| binds={data: a_buffer, kernel: b_buffer}, |
| default_buffer_params=buffer_params, |
| ) |
| |
| |
| def dot_16x1x16_uint8_int8_int16(): |
| """ |
| Int8 dot product by every 2 elements using AVX512 Skylake instructions. |
| This function takes two arrays of uint8 and int8 datatype -- data[2] and |
| kernel[4][32][2] -- and computes a dot product of data[2] with every |
| 2 elements of kernels, resulting in output[4][32] of int16 datatype. |
| The pseudo code is as follows. |
| .. code-block:: c |
| void dot_16x1x16_uint8_int8_int16(uint8 data[2], int8 kernel[32*4][2], |
| int16 output[32*4]){ |
| for (int i = 0; i< 4; i++){ |
| for (int j = 0; j < 32; j++){ |
| output[i][i] = 0; |
| for (int k = 0; k < 2; k++){ |
| output[i][j][k] += data[k] * kernel[i][j][k] |
| } |
| } |
| } |
| } |
| |
| Physically, the kernel array sits in four AVX512 vector registers and |
| the data[2] is broadcasted to another AVX512 vector register. This |
| function returns a TensorIntrin that can be used to tensorize |
| a schedule. |
| |
| Returns |
| ------- |
| intrin : TensorIntrin |
| The Skylake int8 TensorIntrin that can be used in tensorizing schedule |
| """ |
| |
| int16_lanes = 4 * 32 # 4*32 int32 lanes in 4 AVX512 vector registers |
| num_int8_elements = 2 # 2 int8 elements in int16 |
| data = te.placeholder((num_int8_elements,), dtype="uint8", name="data") |
| kernel = te.placeholder((int16_lanes, num_int8_elements), dtype="int8", name="kernel") |
| k = te.reduce_axis((0, num_int8_elements), name="k") |
| C = te.compute( |
| (int16_lanes,), |
| lambda i: te.sum(data[k].astype("int16") * kernel[i, k].astype("int16"), axis=k), |
| name="C", |
| ) |
| |
| a_buffer = tvm.tir.decl_buffer( |
| data.shape, dtype="uint8", name="a_buffer", offset_factor=1, strides=[1] |
| ) |
| b_buffer = tvm.tir.decl_buffer(kernel.shape, dtype="int8", name="b_buffer", offset_factor=1) |
| # strides=[te.var('ldw'), 1, 1]) |
| |
| def _intrin_func(ins, outs): |
| def _instr(index): |
| ib = tvm.tir.ir_builder.create() |
| if index == 1: |
| for i in range(4): |
| ib.emit(outs[0].vstore([i * 32], tvm.tir.const(0, "int16x32"))) |
| return ib.get() |
| |
| a_int8 = ins[0].vload([0], "uint8x2") |
| re_int16 = tvm.tir.call_intrin("int16", "tir.reinterpret", a_int8) |
| vec_ai16 = re_int16.astype("int16x32") |
| vec_a = tvm.tir.call_intrin("int8x64", "tir.reinterpret", vec_ai16) |
| |
| for i in range(4): |
| vec_b = ins[1].vload([i * 32, 0], "int8x64") |
| pair_reduction = tvm.tir.call_llvm_pure_intrin( |
| "int16x32", |
| "llvm.x86.avx512.pmaddubs.w.512", |
| tvm.tir.const(0, "uint32"), |
| vec_a, |
| vec_b, |
| ) |
| if index == 0: |
| ib.emit(outs[0].vstore([i * 32], pair_reduction)) |
| else: |
| ib.emit( |
| outs[0].vstore( |
| [i * 32], pair_reduction + outs[0].vload([i * 32], "int16x32") |
| ) |
| ) |
| return ib.get() |
| |
| # body, reset, update |
| return _instr(0), _instr(1), _instr(2) |
| |
| buffer_params = {"offset_factor": 1} |
| return te.decl_tensor_intrin( |
| C.op, |
| _intrin_func, |
| binds={data: a_buffer, kernel: b_buffer}, |
| default_buffer_params=buffer_params, |
| ) |
| |
| |
| def dot_16x1x16_uint8_int8_int32_cascadelake(): |
| """ |
| Int8 dot product by every 4 elements using AVX512VNNI Cascade Lake instructions. |
| This function takes two arrays of uint8 and int8 datatype -- data[4] and |
| kernel[16][4] -- and computes a dot product of data[4] with every |
| 4 elements of kernels, resulting in output[16] of int32 datatype. |
| The pseudo code is as follows. |
| .. code-block:: c |
| void dot_16x1x16_uint8_int8_int32_cascadelake(uint8 data[4], int8 kernel[16][4], |
| int32 output[16]){ |
| for (int i = 0; i < 16; i++){ |
| output[i] = 0; |
| for (int k = 0; k < 4; k++){ |
| output[i] += data[k] * kernel[i][k] |
| } |
| } |
| } |
| |
| Physically, the kernel array sits in an AVX512 vector register and |
| the data[4] is broadcasted to another AVX512 vector register. This |
| function returns a TensorIntrin that can be used to tensorize |
| a schedule. |
| |
| Returns |
| ------- |
| intrin : TensorIntrin |
| The Cascade Lake int8 TensorIntrin that can be used in tensorizing schedule |
| """ |
| |
| int32_lanes = 16 # 16 int32 lanes in AVX512 |
| num_int8_elements = 4 # 4 int8 elements in int32 |
| data = te.placeholder((num_int8_elements,), dtype="uint8", name="data") |
| kernel = te.placeholder((int32_lanes, num_int8_elements), dtype="int8", name="kernel") |
| k = te.reduce_axis((0, num_int8_elements), name="k") |
| C = te.compute( |
| (int32_lanes,), |
| lambda i: te.sum(data[k].astype("int32") * kernel[i, k].astype("int32"), axis=k), |
| name="C", |
| ) |
| |
| a_buffer = tvm.tir.decl_buffer( |
| data.shape, dtype="uint8", name="a_buffer", offset_factor=1, strides=[1] |
| ) |
| b_buffer = tvm.tir.decl_buffer( |
| kernel.shape, dtype="int8", name="b_buffer", offset_factor=1, strides=[te.var("ldw"), 1] |
| ) |
| |
| def _intrin_func(ins, outs): |
| def _instr(index): |
| ib = tvm.tir.ir_builder.create() |
| if index == 1: |
| ib.emit(outs[0].vstore(0, tvm.tir.const(0, "int32x16"))) |
| return ib.get() |
| |
| a_int8 = ins[0].vload([0], "uint8x4") |
| re_int32 = tvm.tir.call_intrin("int32", "tir.reinterpret", a_int8) |
| vec_ai32 = re_int32.astype("int32x16") |
| vec_b = ins[1].vload([0, 0], "int8x64") |
| |
| vnni_inst_name = "llvm.x86.avx512.vpdpbusd.512" |
| llvm_id = tvm.target.codegen.llvm_lookup_intrinsic_id(vnni_inst_name) |
| |
| if llvm_id != 0: # VNNI is available for current LLVM version |
| vec_bi32 = tvm.tir.call_intrin("int32x16", "tir.reinterpret", vec_b) |
| vec_zero = tvm.tir.const(0, "int32x16") |
| quad_reduction = tvm.tir.call_llvm_pure_intrin( |
| "int32x16", |
| "llvm.x86.avx512.vpdpbusd.512", |
| tvm.tir.const(0, "uint32"), |
| vec_zero, |
| vec_ai32, |
| vec_bi32, |
| ) |
| else: # Fall back to the normal AVX512 |
| vec_a = tvm.tir.call_intrin("int8x64", "tir.reinterpret", vec_ai32) |
| vec_one = tvm.tir.const(1, "int16x32") |
| pair_reduction = tvm.tir.call_llvm_pure_intrin( |
| "int16x32", |
| "llvm.x86.avx512.pmaddubs.w.512", |
| tvm.tir.const(0, "uint32"), |
| vec_a, |
| vec_b, |
| ) |
| quad_reduction = tvm.tir.call_llvm_pure_intrin( |
| "int32x16", |
| "llvm.x86.avx512.pmaddw.d.512", |
| tvm.tir.const(0, "uint32"), |
| pair_reduction, |
| vec_one, |
| ) |
| |
| if index == 0: |
| ib.emit(outs[0].vstore(0, quad_reduction)) |
| else: |
| ib.emit(outs[0].vstore(0, quad_reduction + outs[0].vload([0], "int32x16"))) |
| return ib.get() |
| |
| # body, reset, update |
| return _instr(0), _instr(1), _instr(2) |
| |
| buffer_params = {"offset_factor": 1} |
| return te.decl_tensor_intrin( |
| C.op, |
| _intrin_func, |
| binds={data: a_buffer, kernel: b_buffer}, |
| default_buffer_params=buffer_params, |
| ) |
