saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-22 19:23:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Enable int8 GEMV tensorization (#2696)

Browse source 
Enable int8 GEMV tensorization and improve performance.
This commit is contained in:
baowenlei 2019-12-18 21:02:52 -08:00 committed by GitHub
parent 971bc439b5
commit 715e365723
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 37 additions and 202 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/quantize/matmul_integer.cc
View file

@ -126,9 +126,8 @@ static Status EvaluateMatMulInteger(
force_no_tensorize = true;
}
// Tensorization: AVX2: 8bit GEMM AVX512: 8bit GEMV and GEMM
bool isGEMV = (p_batch_seq_dim != nullptr && *p_batch_seq_dim == 1);
bool use_tensorization = !force_mkl && !force_no_tensorize && (feature.hasAVX512 || (feature.hasAVX2 && !isGEMV) || (!feature.hasAVX2 && feature.hasAVX));
// Tensorization: AVX2: 8bit GEMV/GEMM AVX512: 8bit GEMV/GEMM
bool use_tensorization = !force_mkl && !force_no_tensorize && (feature.hasAVX512 || feature.hasAVX2 || feature.hasAVX);
// Model input option
auto B_NodeArg = node.InputDefs()[1];
@ -156,14 +155,14 @@ static Status EvaluateMatMulInteger(
// TVM has known issue when handling tensorization of matmul: [1x1] = [1xK]x[Kx1]
// and this case is not likely happen in real model
// so add option to fall back to a general reduction
bool isScalar = isGEMV && (embed_dim == 1);
bool isScalar = (p_batch_seq_dim != nullptr && *p_batch_seq_dim == 1) && (embed_dim == 1);
// Tensorization has two layout options: 1) Transpose or 2) Tiling
auto layout_key = !isScalar ? tvm_codegen::WeightLayoutTiling2D::GetKey(TensorProtoDataType(B_NodeArg), vector_width)
: tvm_codegen::WeightLayoutTranspose2D::GetKey(TensorProtoDataType(B_NodeArg));
auto layout_key = isScalar ? tvm_codegen::WeightLayoutTranspose2D::GetKey(TensorProtoDataType(B_NodeArg)) : tvm_codegen::WeightLayoutTiling2D::GetKey(TensorProtoDataType(B_NodeArg), vector_width);
tvm::Tensor B_marshalled = ctx_nuphar->ApplyWeightLayout(layout_key, B_name, B, true);
tvm::Tensor output_tensor = IMatMulTensorize(A, B_marshalled, batchseq_dim, input_dim, embed_dim, vector_width, name + "_IMatMulTensorizeAVX512");
tvm::Tensor output_tensor = IMatMulTensorize(A, B_marshalled, batchseq_dim, input_dim, embed_dim, vector_width, name + "_IMatMulTensorize");
// Post processing output tensor
tvm::Expr embed_padded = B_marshalled->shape[0];

6
onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemm_8bit.cc
View file

@ -101,8 +101,12 @@ void TensorizeIntGemm8bit::TensorizeReduceKernel(std::vector<tvm::Stmt>& inits,
tvm::Expr c_i32v_pred = _0_i32v;
for (int inner_k = 0; inner_k < tensorize_dim_k.tile_size / tensorize_dim_k.layout_size; inner_k++) {
// guard reducion dim k
tvm::Expr guard_reduce = (tensorize_dim_k.dim_iter * tensorize_dim_k.tile_size) + inner_k * tensorize_dim_k.layout_size < tensorize_dim_k.dim_size;
// buffer a regular load
auto a_u8x4 = a_buf.vload({inner_m, inner_k * tensorize_dim_k.layout_size}, u8x4);
auto a_u8x4 = tvm::ir::Select::make(guard_reduce,
a_buf.vload({inner_m, inner_k * tensorize_dim_k.layout_size}, u8x4), tvm::make_const(u8x4, 0));
auto a_u8v = ExpandScalarUInt8(a_u8x4);
// buffer b regular load
auto b_i8v = b_buf.vload({inner_n + inner_k % tensorize_dim_n.layout_size,

12
onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemm_8bit.h
View file

@ -19,14 +19,14 @@ struct TensorizeDimMeta {
tvm::Expr tail_cond;
tvm::Expr load_offset;
TensorizeDimMeta(tvm::Expr dimIter, tvm::Expr dimSize, int tileSize, int layoutSize, bool hasTail, int tailSize, tvm::Expr tailCond)
: dim_iter(dimIter), dim_size(dimSize), tile_size(tileSize), layout_size(layoutSize), has_tail(hasTail), tail_size(tailSize), tail_cond(tailCond) {}
TensorizeDimMeta(tvm::Expr dim_iter, tvm::Expr dim_size, int tile_size, int layout_size, bool has_tail, int tail_size, tvm::Expr tail_cond)
: dim_iter(dim_iter), dim_size(dim_size), tile_size(tile_size), layout_size(layout_size), has_tail(has_tail), tail_size(tail_size), tail_cond(tail_cond) {}
TensorizeDimMeta(tvm::Expr dimIter, tvm::Expr dimSize, int tileSize, bool hasTail, tvm::Expr tailCond)
: dim_iter(dimIter), dim_size(dimSize), tile_size(tileSize), has_tail(hasTail), tail_cond(tailCond) {}
TensorizeDimMeta(tvm::Expr dim_iter, tvm::Expr dim_size, int tile_size, bool has_tail, tvm::Expr tail_cond)
: dim_iter(dim_iter), dim_size(dim_size), tile_size(tile_size), has_tail(has_tail), tail_cond(tail_cond) {}
TensorizeDimMeta(tvm::Expr dimIter, int tileSize, int layoutSize, tvm::Expr offset)
: dim_iter(dimIter), tile_size(tileSize), layout_size(layoutSize), load_offset(offset) {}
TensorizeDimMeta(tvm::Expr dim_iter, tvm::Expr dim_size, int tile_size, int layout_size, tvm::Expr load_offset)
: dim_iter(dim_iter), dim_size(dim_size), tile_size(tile_size), layout_size(layout_size), load_offset(load_offset) {}
int64_t DimSizeValue() {
const int64_t* p_dim_size = tvm::as_const_int(dim_size);

105
onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.cc
View file

@ -1,105 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.h"
#include "core/providers/nuphar/compiler/x86/scheduler/tensorize/tensorize_utilities.h"
#include <tvm/buffer.h>
#include <tvm/codegen.h>
#include <tvm/ir.h>
namespace onnxruntime {
namespace nuphar {
TensorizeIntGemv8bit::TensorizeIntGemv8bit(const std::string& name, const std::vector<int32_t>& vshape)
: TensorizeBase(name, "TensorizeIntGemv8bit_Parameter", {vshape[0], vshape[1]}) {
}
tvm::TensorIntrin TensorizeIntGemv8bit::CreateTensorIntrin() {
tvm::Expr m(shape_[0]);
tvm::Expr l(shape_[1]);
auto a = tvm::placeholder({l}, HalideIR::UInt(8));
auto b = tvm::placeholder({m, l}, HalideIR::Int(8));
auto k = tvm::reduce_axis({0, l});
auto c = tvm::compute({m}, [&](tvm::Var i) {
return tvm::sum(tvm::cast(HalideIR::Int(32), a(k)) * tvm::cast(HalideIR::Int(32), b(i, k)), {k});
});
auto a_buf = tvm::BufferNode::make(
tvm::Var("a", tvm::Handle()),
a->dtype,
a->shape,
/*strides*/ {1},
tvm::Var("a_offset"),
"a",
"",
0,
/*offset_factor*/ 1);
auto b_buf = tvm::BufferNode::make(
tvm::Var("b", tvm::Handle()),
b->dtype,
b->shape,
/*strides*/ {tvm::Var("s1"), 1},
tvm::Var("b_offset"),
"b",
"",
0,
/*offset_factor*/ 1);
auto c_buf = tvm::BufferNode::make(
tvm::Var("c", tvm::Handle()),
c->dtype,
c->shape,
/*strides*/ {1},
tvm::Var("c_offset"),
"c",
"",
0,
/*offset_factor*/ 1);
int h_unroll = shape_[1] / 32;
auto sum_int32x8 = tvm::make_const(HalideIR::Int(32, 8), 0);
auto one = tvm::make_const(HalideIR::Int(16, 16), 1);
for (int i = 0; i < h_unroll; ++i) {
auto a_uint8x32 = a_buf.vload({i * 32}, HalideIR::UInt(8, 32));
auto b_int8x32 = b_buf.vload({0, i * 32}, HalideIR::Int(8, 32));
auto axb_int16x16 = tvm_codegen::LLVMIntrinsic(HalideIR::Int(16, 16),
"llvm.x86.avx2.pmadd.ub.sw",
{a_uint8x32, b_int8x32});
auto axb_int32x8 = tvm_codegen::LLVMIntrinsic(HalideIR::Int(32, 8),
"llvm.x86.avx2.pmadd.wd",
{axb_int16x16, one});
sum_int32x8 += axb_int32x8;
}
sum_int32x8 = tvm_codegen::LLVMIntrinsic(HalideIR::Int(32, 8),
"llvm.x86.avx2.phadd.d",
{sum_int32x8, sum_int32x8});
sum_int32x8 = tvm_codegen::LLVMIntrinsic(HalideIR::Int(32, 8),
"llvm.x86.avx2.phadd.d",
{sum_int32x8, sum_int32x8});
auto sum_int32x4_l = tvm_codegen::VectorLow(sum_int32x8);
auto sum_int32x4_h = tvm_codegen::VectorHigh(sum_int32x8);
auto sum_int32x4 = sum_int32x4_l + sum_int32x4_h;
auto sum_int32x1 = tvm_codegen::ExtractElement(sum_int32x4, 0);
auto reset = c_buf.vstore({0}, tvm::make_const(HalideIR::Int(32, 1), 0));
auto body = c_buf.vstore({0}, sum_int32x1);
auto update = c_buf.vstore({0}, sum_int32x1 + c_buf.vload({0}, HalideIR::Int(32, 1)));
return tvm::TensorIntrinNode::make(
"intrin_gemv_8bit",
c->op,
{a, b},
{a_buf, b_buf, c_buf},
body,
reset,
update);
}
} // namespace nuphar
} // namespace onnxruntime

20
onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.h
View file

@ -1,20 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/providers/nuphar/compiler/x86/scheduler/tensorize/tensorize_base.h"
namespace onnxruntime {
namespace nuphar {
class TensorizeIntGemv8bit : public tvm_codegen::TensorizeBase {
public:
TensorizeIntGemv8bit(const std::string& name, const std::vector<int32_t>& vshape);
virtual ~TensorizeIntGemv8bit() = default;
tvm::TensorIntrin CreateTensorIntrin() override;
};
} // namespace nuphar
} // namespace onnxruntime

75
onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize_schedule.cc
View file

@ -9,7 +9,6 @@
#include "core/providers/nuphar/compiler/nuphar_codegen_ctx.h"
#include "core/providers/nuphar/compiler/x86/scheduler/nuphar_scheduler.h"
#include "core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_16bit.h"
#include "core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.h"
#include "core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemm_8bit.h"
#include "core/providers/nuphar/compiler/x86/x86_target_info.h"
#include "core/codegen/passes/scheduler/schedule_utils.h"
@ -56,48 +55,10 @@ static Status TensorizeGEMVInteger16(const tvm::Tensor& tensor,
return Status::OK();
}
// TODO: refactor below function
static Status TensorizeGEMVInteger(const tvm::Tensor& tensor,
const int64_t input_dim,
tvm_codegen::CodeGenContext& ctx_codegen,
tvm_codegen::ScheduleContext& ctx_sched) {
// schedule for imatmul inputs
InsertRootScheduleAndClosure(tensor, ctx_sched);
InputRootScheduleWithVectorizationX86(tensor, ctx_codegen, ctx_sched);
// decide kernel shape
std::vector<int32_t> kernel_shape;
kernel_shape.push_back(1);
if (input_dim <= 256) {
kernel_shape.push_back(input_dim);
} else if (input_dim % 64 == 0) {
kernel_shape.push_back(64);
} else {
kernel_shape.push_back(32);
}
TensorizeIntGemv8bit igemv8bit("igemv8bit", kernel_shape);
auto shape = igemv8bit.Shape();
auto compute_op = tensor->op.as<tvm::ComputeOpNode>();
auto xy = compute_op->axis;
auto x = xy[0];
auto y = xy[1];
auto z = compute_op->reduce_axis[0];
tvm::IterVar yo, yi;
ctx_sched.schedule[tensor->op].split(y, shape[0], &yo, &yi);
tvm::IterVar zo, zi;
ctx_sched.schedule[tensor->op].split(z, shape[1], &zo, &zi);
ctx_sched.schedule[tensor->op].reorder({x, yo, zo, yi, zi});
ctx_sched.schedule[tensor->op].tensorize(yi, igemv8bit.CreateTensorIntrin());
return Status::OK();
}
static Status TensorizeIGEMV(const tvm::Tensor& tensor,
tvm_codegen::CodeGenContext& ctx_codegen,
tvm_codegen::ScheduleContext& ctx_sched,
bool tensorize,
const std::string& target_str) {
static Status TensorizeIScalar(const tvm::Tensor& tensor,
tvm_codegen::CodeGenContext& ctx_codegen,
tvm_codegen::ScheduleContext& ctx_sched,
const std::string& target_str) {
// Schedule tensor and inputs as root
bool status_imatmul = InsertRootScheduleAndClosure(tensor, ctx_sched);
if (status_imatmul == false)
@ -105,7 +66,6 @@ static Status TensorizeIGEMV(const tvm::Tensor& tensor,
InputRootScheduleWithVectorizationX86(tensor, ctx_codegen, ctx_sched);
// Default tiling size
// TODO: tuning tiling sizes later
int tensorize_embed = 1;
int tensorize_input = (target_str == "avx512-skylake") ? 1024 : (target_str == "avx2") ? 512 : 256;
@ -119,19 +79,12 @@ static Status TensorizeIGEMV(const tvm::Tensor& tensor,
auto x = xy[0];
auto y = xy[1];
auto z = compute_op->reduce_axis[0];
// no tiling need for IterVar x
tvm::IterVar yo, yi;
ctx_sched.schedule[tensor->op].split(y, kernel_shape[0], &yo, &yi);
tvm::IterVar zo, zi;
ctx_sched.schedule[tensor->op].split(z, kernel_shape[1], &zo, &zi);
ctx_sched.schedule[tensor->op].reorder({x, yo, zo, yi, zi});
if (tensorize) {
// TODO: refine tensorize gemv class
TensorizeIntGemv8bit igemv8bit("igemv8bit", kernel_shape);
ctx_sched.schedule[tensor->op].tensorize(yi, igemv8bit.CreateTensorIntrin());
}
return Status::OK();
}
@ -273,7 +226,7 @@ static Status TensorizeIGEMM(const tvm::Tensor& tensor,
tvm::Expr input_iter(zo);
tvm::Expr load_shift = tvm::ir::Simplify(kernel_shape[2] / layout_tile_row * input_dim_padded - kernel_shape[2]);
tvm::Expr load_offset = tvm::ir::Simplify((input_iter % (std::max(1, vector_width / kernel_shape[2]))) * load_shift);
TensorizeDimMeta input_meta(input_iter, kernel_shape[2], layout_tile_row, load_offset);
TensorizeDimMeta input_meta(input_iter, input_dim_padded, kernel_shape[2], layout_tile_row, load_offset);
TensorizeIntGemm8bit igemm8bit("igemm8bit", kernel_shape, target_str);
igemm8bit.InsertTensorizeDimInfo("m", batchseq_meta);
@ -341,23 +294,21 @@ static bool IMatMulTensorizeSchedule(
bool status_tensorize = true;
if (is8bit) {
if (feature.hasAVX512) { // isAVX512
status_tensorize = is_scalar ? TensorizeIGEMV(imatmul, ctx_codegen, ctx_sched, /*tensorize=*/false, "avx512-skylake").IsOK()
status_tensorize = is_scalar ? TensorizeIScalar(imatmul, ctx_codegen, ctx_sched, "avx512-skylake").IsOK()
: TensorizeIGEMM(imatmul, ctx_codegen, ctx_sched, batchseq_expr,
{*p_embed_dim, *p_embed_dim_padded},
{*p_input_dim, *p_input_dim_padded},
"avx512-skylake")
.IsOK();
} else if (feature.hasAVX2) { // isAVX2
ORT_ENFORCE(!is_scalar, "scalar AVX2 is not supported!");
// TODO: release 8bit tensorize GEMV for AVX2
status_tensorize = isGEMV ? TensorizeGEMVInteger(imatmul, *p_input_dim, ctx_codegen, ctx_sched).IsOK()
: TensorizeIGEMM(imatmul, ctx_codegen, ctx_sched, batchseq_expr,
{*p_embed_dim, *p_embed_dim_padded},
{*p_input_dim, *p_input_dim_padded},
"avx2")
.IsOK();
status_tensorize = is_scalar ? TensorizeIScalar(imatmul, ctx_codegen, ctx_sched, "avx2").IsOK()
: TensorizeIGEMM(imatmul, ctx_codegen, ctx_sched, batchseq_expr,
{*p_embed_dim, *p_embed_dim_padded},
{*p_input_dim, *p_input_dim_padded},
"avx2")
.IsOK();
} else if (feature.hasAVX) { // isAVX
status_tensorize = is_scalar ? TensorizeIGEMV(imatmul, ctx_codegen, ctx_sched, /*tensorize=*/false, "avx").IsOK()
status_tensorize = is_scalar ? TensorizeIScalar(imatmul, ctx_codegen, ctx_sched, "avx").IsOK()
: TensorizeIGEMM(imatmul, ctx_codegen, ctx_sched, batchseq_expr,
{*p_embed_dim, *p_embed_dim_padded},
{*p_input_dim, *p_input_dim_padded},

8
onnxruntime/test/providers/cpu/math/matmul_integer_test.cc
View file

@ -45,8 +45,9 @@ TEST(MatmulIntegerOpTest, MatMulInteger_WithZero_ZeroPoint) {
template <typename T>
std::vector<T> ToVector(const int* value, int size) {
std::vector<T> data(size);
for (int i = 0; i < size; i++)
for (int i = 0; i < size; i++) {
data[i] = static_cast<T>(value[i]);
}
return data;
}
@ -81,12 +82,17 @@ TEST(MatmulIntegerOpTest, MatMulInteger_Uint8_Int8_Scalar) {
TEST(MatmulIntegerOpTest, MatMulInteger_Uint8_Int8_GEMV) {
RunMatMulIntegerU8S8Test(1, 2, 16);
RunMatMulIntegerU8S8Test(1, 2, 64);
RunMatMulIntegerU8S8Test(1, 8, 36);
RunMatMulIntegerU8S8Test(1, 8, 68);
RunMatMulIntegerU8S8Test(1, 8, 400);
RunMatMulIntegerU8S8Test(1, 512, 1024);
}
TEST(MatmulIntegerOpTest, MatMulInteger_Uint8_Int8_GEMM) {
RunMatMulIntegerU8S8Test(2, 2, 40);
RunMatMulIntegerU8S8Test(2, 48, 33);
RunMatMulIntegerU8S8Test(2, 51, 40);
RunMatMulIntegerU8S8Test(4, 8, 68);
}
} // namespace test