From 715e3657239bc6b5d92f434e71b935677dfcfe27 Mon Sep 17 00:00:00 2001 From: baowenlei <40643946+baowenlei@users.noreply.github.com> Date: Wed, 18 Dec 2019 21:02:52 -0800 Subject: [PATCH] Enable int8 GEMV tensorization (#2696) Enable int8 GEMV tensorization and improve performance. --- .../math/quantize/matmul_integer.cc | 13 +-- .../scheduler/tensorize/intrin_gemm_8bit.cc | 6 +- .../scheduler/tensorize/intrin_gemm_8bit.h | 12 +- .../scheduler/tensorize/intrin_gemv_8bit.cc | 105 ------------------ .../scheduler/tensorize/intrin_gemv_8bit.h | 20 ---- .../x86/scheduler/tensorize_schedule.cc | 75 +++---------- .../providers/cpu/math/matmul_integer_test.cc | 8 +- 7 files changed, 37 insertions(+), 202 deletions(-) delete mode 100644 onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.cc delete mode 100644 onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.h diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/quantize/matmul_integer.cc b/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/quantize/matmul_integer.cc index 21e1c95f2d..a4f5c75ae6 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/quantize/matmul_integer.cc +++ b/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/quantize/matmul_integer.cc @@ -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]; diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemm_8bit.cc b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemm_8bit.cc index 0b738d4738..0a2242dfb1 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemm_8bit.cc +++ b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemm_8bit.cc @@ -101,8 +101,12 @@ void TensorizeIntGemm8bit::TensorizeReduceKernel(std::vector& 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, diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemm_8bit.h b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemm_8bit.h index 4859a94d74..9130458eb8 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemm_8bit.h +++ b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemm_8bit.h @@ -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); diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.cc b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.cc deleted file mode 100644 index 5fe99d4e2e..0000000000 --- a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.cc +++ /dev/null @@ -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 -#include -#include - -namespace onnxruntime { -namespace nuphar { - -TensorizeIntGemv8bit::TensorizeIntGemv8bit(const std::string& name, const std::vector& 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 diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.h b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.h deleted file mode 100644 index cfa069c038..0000000000 --- a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_8bit.h +++ /dev/null @@ -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& vshape); - - virtual ~TensorizeIntGemv8bit() = default; - - tvm::TensorIntrin CreateTensorIntrin() override; -}; - -} // namespace nuphar -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize_schedule.cc b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize_schedule.cc index 9f3078addc..3495cd85c0 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize_schedule.cc +++ b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize_schedule.cc @@ -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 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(); - 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}, diff --git a/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc b/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc index 5a9ca45b57..775c09123a 100644 --- a/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc +++ b/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc @@ -45,8 +45,9 @@ TEST(MatmulIntegerOpTest, MatMulInteger_WithZero_ZeroPoint) { template std::vector ToVector(const int* value, int size) { std::vector data(size); - for (int i = 0; i < size; i++) + for (int i = 0; i < size; i++) { data[i] = static_cast(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