diff --git a/onnxruntime/core/codegen/common/utils.cc b/onnxruntime/core/codegen/common/utils.cc index dfa3abf06a..39b871b0c6 100644 --- a/onnxruntime/core/codegen/common/utils.cc +++ b/onnxruntime/core/codegen/common/utils.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/codegen/common/utils.h" +#include "core/common/cpuid_info.h" #include "core/common/make_unique.h" #include @@ -60,8 +61,42 @@ void GetStrides(const int64_t* shape, int ndim, std::vector& strides) { strides.resize(ndim); strides[ndim - 1] = 1; for (int64_t i = ndim - 2; i >= 0; i--) { - strides[i] = strides[i+1] * shape[i+1]; + strides[i] = strides[i + 1] * shape[i + 1]; } } +// Common utils to get target option +TargetFeature GetTargetInfo(const codegen::CodeGenSettings& settings) { + TargetFeature feature; + + std::string target_str = ""; + if (settings.HasOption(nuphar::kNupharCodeGenTarget) && settings.HasOption(nuphar::kNupharCachePath)) { + target_str = settings.GetOptionValue(nuphar::kNupharCodeGenTarget); + } + + bool isAVX = false; + bool isAVX2 = false; + bool isAVX512 = false; + if (target_str == "avx") { + isAVX = true; + } else if (target_str == "avx2") { + isAVX = true; + isAVX2 = true; + } else if (target_str == "avx512") { + isAVX = true; + isAVX2 = true; + isAVX512 = true; + } else { + isAVX = CPUIDInfo::GetCPUIDInfo().HasAVX(); + isAVX2 = CPUIDInfo::GetCPUIDInfo().HasAVX2(); + isAVX512 = CPUIDInfo::GetCPUIDInfo().HasAVX512Skylake(); + } + + feature.hasAVX = isAVX; + feature.hasAVX2 = isAVX2; + feature.hasAVX512 = isAVX512; + + return feature; +} + } // namespace onnxruntime diff --git a/onnxruntime/core/codegen/common/utils.h b/onnxruntime/core/codegen/common/utils.h index d85df1d01e..008613c1c9 100644 --- a/onnxruntime/core/codegen/common/utils.h +++ b/onnxruntime/core/codegen/common/utils.h @@ -2,6 +2,7 @@ // Licensed under the MIT License. #pragma once +#include "core/providers/nuphar/common/nuphar_settings.h" #include #include #include @@ -19,4 +20,12 @@ int64_t TotalSize(const std::vector& shape); void GetStrides(const int64_t* shape, int ndim, std::vector& strides); +struct TargetFeature { + bool hasAVX; + bool hasAVX2; + bool hasAVX512; +}; + +TargetFeature GetTargetInfo(const codegen::CodeGenSettings& setttings); + } // namespace onnxruntime diff --git a/onnxruntime/core/common/cpuid_info.cc b/onnxruntime/core/common/cpuid_info.cc index 7986cd6186..e8595849c0 100644 --- a/onnxruntime/core/common/cpuid_info.cc +++ b/onnxruntime/core/common/cpuid_info.cc @@ -54,7 +54,9 @@ CPUIDInfo::CPUIDInfo() noexcept { const int AVX_MASK = 0x6; const int AVX512_MASK = 0xE6; int value = XGETBV(); - has_avx_ = (data[2] & (1 << 28)) && ((value & AVX_MASK) == AVX_MASK); + bool has_sse2 = (data[3] & (1 << 26)); + bool has_ssse3 = (data[2] & (1 << 9)); + has_avx_ = has_sse2 && has_ssse3 && (data[2] & (1 << 28)) && ((value & AVX_MASK) == AVX_MASK); bool has_avx512 = (value & AVX512_MASK) == AVX512_MASK; has_f16c_ = has_avx_ && (data[2] & (1 << 29)) && (data[3] & (1 << 26)); diff --git a/onnxruntime/core/providers/nuphar/common/nuphar_settings.h b/onnxruntime/core/providers/nuphar/common/nuphar_settings.h index f3deee31eb..8e836b7e84 100644 --- a/onnxruntime/core/providers/nuphar/common/nuphar_settings.h +++ b/onnxruntime/core/providers/nuphar/common/nuphar_settings.h @@ -42,7 +42,7 @@ constexpr static const char* kNupharFastMath_ShortPolynormial = "short_polynormi constexpr static const char* kNupharFastActivation = "nuphar_fast_activation"; // fast activation constexpr static const char* kNupharActivations_DeepCpu = "deep_cpu_activation"; -// Option to control nuphar code generation target (avx2 or avx512) +// Option to control nuphar code generation target (avx / avx2 / avx512) constexpr static const char* kNupharCodeGenTarget = "nuphar_codegen_target"; // cache version number (MAJOR.MINOR.PATCH) following https://semver.org/ diff --git a/onnxruntime/core/providers/nuphar/compiler/codegen_manager.cc b/onnxruntime/core/providers/nuphar/compiler/codegen_manager.cc index c7aa47277e..582ada8f3b 100644 --- a/onnxruntime/core/providers/nuphar/compiler/codegen_manager.cc +++ b/onnxruntime/core/providers/nuphar/compiler/codegen_manager.cc @@ -98,14 +98,22 @@ static void RegisterAllNupharSchedulers(tvm_codegen::TVMScheduleRegistry* sched_ // 3. Create Weight layout instances // BEGIN: Nuphar Weight Layouts classes static void RegisterAllNupharWeightLayouts(tvm_codegen::WeightLayoutRegistry* layout_registry) { + // AVX512 layout_registry->Register( std::move(onnxruntime::make_unique(ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT8, 64))); layout_registry->Register( std::move(onnxruntime::make_unique(ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT16, 64))); + // AVX2 layout_registry->Register( std::move(onnxruntime::make_unique(ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT8, 32))); layout_registry->Register( std::move(onnxruntime::make_unique(ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT16, 32))); + // AVX + layout_registry->Register( + std::move(onnxruntime::make_unique(ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT8, 16))); + layout_registry->Register( + std::move(onnxruntime::make_unique(ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_INT16, 16))); + layout_registry->Register( std::move(onnxruntime::make_unique(ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_FLOAT, 8))); layout_registry->Register( 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 ba901ce25d..21e1c95f2d 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 @@ -3,6 +3,7 @@ #include "core/providers/nuphar/compiler/x86/op_ir_creator/all_ops.h" +#include "core/codegen/common/utils.h" #include "core/codegen/mti/math/binary_ops.h" #include "core/codegen/mti/math/matmul_ops.h" #include "core/codegen/mti/mti_tvm_utils.h" @@ -49,10 +50,11 @@ tvm::Tensor IMatMulTensorize(const tvm::Tensor& A, A_reshape = A; } else { A_reshape = tvm_codegen::Reshape(A, {batchseq_dim, input_dim}, name + "_reshape_X"); - if (input_dim != input_padded) { - tvm::Expr pad_value = tvm::make_const(A->dtype, 0); - A_reshape = tvm_codegen::PadLastDim(A_reshape, vector_width, pad_value); - } + } + + if (input_dim != input_padded) { + tvm::Expr pad_value = tvm::make_const(A->dtype, 0); + A_reshape = tvm_codegen::PadLastDim(A_reshape, vector_width, pad_value); } tvm::Tensor Y = tvm::compute( @@ -113,6 +115,8 @@ static Status EvaluateMatMulInteger( // Enviornment variables option const codegen::CodeGenSettings& settings = codegen::CodeGenSettings::Instance(); + TargetFeature feature = GetTargetInfo(settings); + bool force_mkl = false; if (settings.HasOption(kNupharIMatMulForceMkl)) { force_mkl = true; @@ -122,13 +126,9 @@ static Status EvaluateMatMulInteger( force_no_tensorize = true; } - // Target instruction option - bool isAVX2 = CPUIDInfo::GetCPUIDInfo().HasAVX2(); - bool isAVX512 = CPUIDInfo::GetCPUIDInfo().HasAVX512Skylake(); - // 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 && (isAVX512 || (isAVX2 && !isGEMV)); + bool use_tensorization = !force_mkl && !force_no_tensorize && (feature.hasAVX512 || (feature.hasAVX2 && !isGEMV) || (!feature.hasAVX2 && feature.hasAVX)); // Model input option auto B_NodeArg = node.InputDefs()[1]; @@ -139,19 +139,19 @@ static Status EvaluateMatMulInteger( //TODO: change to use MLAS when no layout could apply tvm::Tensor B_marshalled = tvm_codegen::Transpose(B, {1, 0}); - bool use_extern_MKL = (force_mkl || !isAVX2); + bool use_extern_MKL = (force_mkl || !feature.hasAVX2); tvm::Tensor output_tensor = use_extern_MKL ? IMatMulExternMKL(A, B_marshalled, output_shape, input_dim, embed_dim, name + "_IMatMulExternMKL") : IMatMulExternAVX2(A, B_marshalled, output_shape, input_dim, embed_dim, name + "_IMatMulExternAVX2"); + outputs.push_back(output_tensor); } else if (use_tensorization) { // vector width determined from target hardware - // AVX2: vector width 32 = 256bits / 8bit; 16 = 256 bits / 16bits; - // AVX512: vector width 64 = 512bits / 8bit; 32 = 512 bits / 16bits; - int vector_width = 32; + // AVX: vector width 16 = 128 bits / 8 bits; + // AVX2: vector width 32 = 256 bits / 8 bits; + // AVX512: vector width 64 = 512 bits / 8 bits; CodeGenTargetX86* target = dynamic_cast(ctx_codegen.GetCodeGenHandle()->codegen_target); - if (target != nullptr) { - vector_width = target->NaturalVectorWidth(B->dtype.bits()) / 2; - } + ORT_ENFORCE(target != nullptr, "CodeGen target unknown: not AVX/AVX2/AVX512 !"); + int vector_width = target->NaturalVectorWidth(B->dtype.bits()) / 2; // TVM has known issue when handling tensorization of matmul: [1x1] = [1xK]x[Kx1] // and this case is not likely happen in real model @@ -192,7 +192,7 @@ static Status EvaluateMatMulInteger( auto layout_key = tvm_codegen::WeightLayoutTranspose2D::GetKey(TensorProtoDataType(B_NodeArg)); tvm::Tensor B_marshalled = ctx_nuphar->ApplyWeightLayout(layout_key, B_name, B, true); - bool use_extern_AVX2 = (!force_mkl && isAVX2); + bool use_extern_AVX2 = (!force_mkl && feature.hasAVX2); tvm::Tensor output_tensor = use_extern_AVX2 ? IMatMulExternAVX2(A, B_marshalled, output_shape, input_dim, embed_dim, name + "_IMatMulExternAVX2") : IMatMulExternMKL(A, B_marshalled, output_shape, input_dim, embed_dim, name + "_IMatMulExternMKL"); @@ -222,7 +222,13 @@ static Status EvaluateMatMulInteger16( bool is16bit = (A->dtype == HalideIR::type_of() && B->dtype == HalideIR::type_of()); - if (B->shape.size() == 2 && is16bit) { + const codegen::CodeGenSettings& settings = codegen::CodeGenSettings::Instance(); + TargetFeature feature = GetTargetInfo(settings); + + // 16bit on AVX fall back to 32bit + bool AVXonly = feature.hasAVX && !feature.hasAVX2; + + if (!AVXonly && (B->shape.size() == 2 && is16bit)) { const int64_t* p_input_dim = tvm::as_const_int(B->shape[0]); const int64_t* p_embed_dim = tvm::as_const_int(B->shape[1]); @@ -237,16 +243,11 @@ static Status EvaluateMatMulInteger16( } output_shape.push_back(tvm::Expr(gsl::narrow_cast(embed_dim))); - // Enviornment variable option - const codegen::CodeGenSettings& settings = codegen::CodeGenSettings::Instance(); bool force_mkl = false; if (settings.HasOption(kNupharIMatMulForceMkl)) { force_mkl = true; } - // Target instruction option - bool isAVX2 = CPUIDInfo::GetCPUIDInfo().HasAVX2(); - // Model input option auto B_NodeArg = node.InputDefs()[1]; const std::string& B_name = B_NodeArg->Name(); @@ -256,7 +257,7 @@ static Status EvaluateMatMulInteger16( //TODO: change to use MLAS when no layout could apply tvm::Tensor B_marshalled = tvm_codegen::Transpose(B, {1, 0}); - bool use_extern_MKL = (force_mkl || !isAVX2); + bool use_extern_MKL = (force_mkl || !feature.hasAVX2); tvm::Tensor output_tensor = use_extern_MKL ? IMatMul16ExternMKL(A, B_marshalled, output_shape, input_dim, embed_dim, node.Name() + "_IMatMulExternMKL") : IMatMul16ExternAVX2(A, B_marshalled, output_shape, input_dim, embed_dim, node.Name() + "_IMatMulExternAVX2"); outputs.push_back(output_tensor); @@ -264,7 +265,7 @@ static Status EvaluateMatMulInteger16( auto layout_key = tvm_codegen::WeightLayoutTranspose2D::GetKey(TensorProtoDataType(B_NodeArg)); tvm::Tensor B_marshalled = ctx_nuphar->ApplyWeightLayout(layout_key, B_name, B, true); - bool use_extern_AVX2 = (!force_mkl && isAVX2); + bool use_extern_AVX2 = (!force_mkl && feature.hasAVX2); tvm::Tensor output_tensor = use_extern_AVX2 ? IMatMul16ExternAVX2(A, B_marshalled, output_shape, input_dim, embed_dim, node.Name() + "_IMatMulExternAVX2") : IMatMul16ExternMKL(A, B_marshalled, output_shape, input_dim, embed_dim, node.Name() + "_IMatMulExternMKL"); outputs.push_back(output_tensor); 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 c1fb063467..0b738d4738 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 @@ -26,6 +26,8 @@ tvm::Expr TensorizeIntGemm8bit::CreatePredicateMask(int tail_size) { mask_lanes = 16; } else if (tensorize_target_ == "avx2") { mask_lanes = 8; + } else if (tensorize_target_ == "avx") { + mask_lanes = 4; } else { ORT_NOT_IMPLEMENTED("Tensorization only support avx2/avx512-skylake currently!"); } @@ -221,6 +223,12 @@ tvm::TensorIntrin TensorizeIntGemm8bit::CreateTensorIntrin() { HalideIR::Int(8, 32), HalideIR::Int(16, 16), HalideIR::Int(32, 8), "llvm.x86.avx2.pmadd.ub.sw", "llvm.x86.avx2.pmadd.wd"); tensorize_targets_meta_.emplace(tensorize_target_, avx2_info); + } else if (tensorize_target_ == "avx") { + TensorizeTargetInfo avx_info(HalideIR::UInt(1, 4), HalideIR::UInt(8, 16), + HalideIR::Int(8, 16), HalideIR::Int(16, 8), HalideIR::Int(32, 4), + "llvm.x86.ssse3.pmadd.ub.sw.128", "llvm.x86.sse2.pmadd.wd"); + + tensorize_targets_meta_.emplace(tensorize_target_, avx_info); } else { ORT_NOT_IMPLEMENTED("Tensorization only support avx2/avx512-skylake currently!"); } 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 8c6c3b4b42..34278804d4 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize_schedule.cc +++ b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize_schedule.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include "core/codegen/common/utils.h" #include "core/common/cpuid_info.h" #include "core/framework/op_kernel_info.h" #include "core/providers/nuphar/common/analysis/subgraph_codegen_stats.h" @@ -54,6 +55,7 @@ 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::ScheduleContext& ctx) { @@ -89,31 +91,54 @@ static Status TensorizeGEMVInteger(const tvm::Tensor& tensor, return Status::OK(); } -static Status TensorizeReduction(const tvm::Tensor& tensor, - tvm_codegen::ScheduleContext& ctx, - const std::string& target_str) { +static Status TensorizeIGEMV(const tvm::Tensor& tensor, + tvm_codegen::ScheduleContext& ctx, + bool tensorize, + const std::string& target_str) { // Schedule tensor and inputs as root bool status_imatmul = InsertRootScheduleAndClosure(tensor, ctx); if (status_imatmul == false) return Status::OK(); - InputRootScheduleWithVectorizationX86(tensor, ctx); - // Loop tiling - int reduce_tile_size = (target_str == "avx512-skylake") ? 1024 : 512; + + // 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; + + // Tensorize kernel shape + std::vector kernel_shape; + kernel_shape.push_back(tensorize_embed); + kernel_shape.push_back(tensorize_input); + 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]; + + // no tiling need for IterVar x + tvm::IterVar yo, yi; + ctx.schedule[tensor->op].split(y, kernel_shape[0], &yo, &yi); tvm::IterVar zo, zi; - ctx.schedule[tensor->op].split(z, reduce_tile_size, &zo, &zi); + ctx.schedule[tensor->op].split(z, kernel_shape[1], &zo, &zi); + ctx.schedule[tensor->op].reorder({x, yo, zo, yi, zi}); + + if (tensorize) { + // TODO: refine tensorize gemv class + TensorizeIntGemv8bit igemv8bit("igemv8bit", kernel_shape); + ctx.schedule[tensor->op].tensorize(yi, igemv8bit.CreateTensorIntrin()); + } return Status::OK(); } -static Status TensorizeGEMMInteger(const tvm::Tensor& tensor, - tvm_codegen::CodeGenContext& ctx_codegen, - tvm_codegen::ScheduleContext& ctx, - tvm::Expr batchseq_expr, - const std::vector embed_dim_vec, - const std::vector input_dim_vec, - const std::string& target_str) { +static Status TensorizeIGEMM(const tvm::Tensor& tensor, + tvm_codegen::CodeGenContext& ctx_codegen, + tvm_codegen::ScheduleContext& ctx, + tvm::Expr batchseq_expr, + const std::vector embed_dim_vec, + const std::vector input_dim_vec, + const std::string& target_str) { // Schedule tensor and inputs as root bool status_imatmul = InsertRootScheduleAndClosure(tensor, ctx); if (status_imatmul == false) @@ -125,13 +150,17 @@ static Status TensorizeGEMMInteger(const tvm::Tensor& tensor, int tensorize_embed = 8; int tensorize_input = 32; if (target_str == "avx512-skylake") { - tensorize_batch = 4; + tensorize_batch = 8; tensorize_embed = 16; tensorize_input = 64; } else if (target_str == "avx2") { tensorize_batch = 8; tensorize_embed = 16; tensorize_input = 32; + } else if (target_str == "avx") { + tensorize_batch = 2; + tensorize_embed = 4; + tensorize_input = 8; } codegen::CodeGenSettings& settings = codegen::CodeGenSettings::Instance(); @@ -161,6 +190,8 @@ static Status TensorizeGEMMInteger(const tvm::Tensor& tensor, embed_min = 16; } else if (target_str == "avx2") { embed_min = 8; + } else if (target_str == "avx") { + embed_min = 4; } tensorize_embed = (tensorize_embed % embed_min != 0) ? ((tensorize_embed + embed_min - 1) / embed_min) * embed_min : tensorize_embed; @@ -194,18 +225,21 @@ static Status TensorizeGEMMInteger(const tvm::Tensor& tensor, ctx.schedule[tensor->op].reorder({yo, xo, zo, xi, yi, zi}); } else { // Loop nest default order - ctx.schedule[tensor->op].reorder({xo, yo, zo, xi, yi, zi}); + if (target_str == "avx") + ctx.schedule[tensor->op].reorder({yo, xo, zo, xi, yi, zi}); + else + ctx.schedule[tensor->op].reorder({xo, yo, zo, xi, yi, zi}); } // Natural vector width - // AVX2: vector width 32 = 256bits / 8bit; 16 = 256 bits / 16bits; - // AVX512: vector width 64 = 512bits / 8bit; 32 = 512 bits / 16bits; - int vector_width = 32; - int tensor_bits = tensor->op->InputTensors()[1]->dtype.bits(); + // AVX: vector width 16 = 128 bits / 8 bits; 8 = 128 bits / 16bits; + // AVX2: vector width 32 = 256 bits / 8 bits; 16 = 256 bits / 16bits; + // AVX512: vector width 64 = 512 bits / 8 bits; 32 = 512 bits / 16bits; CodeGenTargetX86* target = dynamic_cast(ctx_codegen.GetCodeGenHandle()->codegen_target); - if (target != nullptr) { - vector_width = target->NaturalVectorWidth(tensor_bits) / 2; - } + ORT_ENFORCE(target != nullptr, "CodeGen target unknown: not AVX/AVX2/AVX512 !"); + int tensor_bits = tensor->op->InputTensors()[1]->dtype.bits(); + int vector_width = target->NaturalVectorWidth(tensor_bits) / 2; + // Layout shape int layout_tile_row = (sizeof(int32_t) * bits_per_byte) / tensor_bits; int layout_tile_col = ((vector_width * bits_per_byte) / tensor_bits) / layout_tile_row; @@ -282,32 +316,46 @@ static bool IMatMulTensorizeSchedule( // so add option to fall back to a general reduction bool is_scalar = isGEMV && (*p_embed_dim == 1); + codegen::CodeGenSettings& settings = codegen::CodeGenSettings::Instance(); + TargetFeature feature = GetTargetInfo(settings); + bool status_tensorize = true; if (is8bit) { - if (CPUIDInfo::GetCPUIDInfo().HasAVX512Skylake()) { - status_tensorize = is_scalar ? TensorizeReduction(imatmul, ctx_sched, "avx512-skylake").IsOK() - : TensorizeGEMMInteger(imatmul, ctx_codegen, ctx_sched, batchseq_expr, - {*p_embed_dim, *p_embed_dim_padded}, - {*p_input_dim, *p_input_dim_padded}, - "avx512-skylake") + if (feature.hasAVX512) { // isAVX512 + status_tensorize = is_scalar ? TensorizeIGEMV(imatmul, ctx_sched, /*tensorize=*/false, "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 (CPUIDInfo::GetCPUIDInfo().HasAVX2()) { + } 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_sched).IsOK() - : TensorizeGEMMInteger(imatmul, ctx_codegen, ctx_sched, batchseq_expr, - {*p_embed_dim, *p_embed_dim_padded}, - {*p_input_dim, *p_input_dim_padded}, - "avx2") + : 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_sched, /*tensorize=*/false, "avx").IsOK() + : TensorizeIGEMM(imatmul, ctx_codegen, ctx_sched, batchseq_expr, + {*p_embed_dim, *p_embed_dim_padded}, + {*p_input_dim, *p_input_dim_padded}, + "avx") + .IsOK(); + } else { + ORT_NOT_IMPLEMENTED("Not supported target in 8bit Tensorization, should be one of avx/avx2/avx512."); } } else { // 16bit // TODO: add 16bit tensorize GEMV/GEMM for AVX512 - if (CPUIDInfo::GetCPUIDInfo().HasAVX2()) { + if (feature.hasAVX2) { //isAVX2 // TODO: add 16bit tensorize GEMM for AVX2 ORT_ENFORCE(isGEMV, "16bit GEMM is not supported!"); // TODO: release 16bit tensorize GEMV for AVX2 status_tensorize = TensorizeGEMVInteger16(imatmul, *p_input_dim, ctx_sched).IsOK(); + } else { + ORT_NOT_IMPLEMENTED("Not supported target in 16bit Tensorization."); } } @@ -333,6 +381,7 @@ bool TVM_SCHEDULER_CLASS(MatMulInteger, NupharX86Tensorize)::Evaluate( return status_reshape || status_tensorize; } +// TODO: enable 16 bit tensorization bool TVM_SCHEDULER_CLASS(MatMulInteger16, NupharX86Tensorize)::Evaluate( const tvm::Tensor& tensor, const Node* node, diff --git a/onnxruntime/core/providers/nuphar/nuphar_execution_provider.cc b/onnxruntime/core/providers/nuphar/nuphar_execution_provider.cc index 03f6b61b84..b36d526828 100644 --- a/onnxruntime/core/providers/nuphar/nuphar_execution_provider.cc +++ b/onnxruntime/core/providers/nuphar/nuphar_execution_provider.cc @@ -57,9 +57,9 @@ NupharExecutionProvider::NupharExecutionProvider(const NupharExecutionProviderIn target_str = default_nuphar_target_str; } + const auto& cpu_id_info = CPUIDInfo::GetCPUIDInfo(); if (target_str == llvm_target_str) { // auto detect from CPU ID - const auto& cpu_id_info = CPUIDInfo::GetCPUIDInfo(); if (cpu_id_info.HasAVX512f()) { codegen_target_ = CodeGenTarget_AVX512(); } else if (cpu_id_info.HasAVX2()) { @@ -81,9 +81,21 @@ NupharExecutionProvider::NupharExecutionProvider(const NupharExecutionProviderIn ORT_NOT_IMPLEMENTED("Not supported target, should be one of stackvm/llvm/avx/avx2/avx512."); } - CreateTVMTarget(); + if (settings.HasOption(nuphar::kNupharCodeGenTarget)) { + if ((target_str == "avx512" && !cpu_id_info.HasAVX512f()) || + (target_str == "avx2" && !cpu_id_info.HasAVX2()) || + (target_str == "avx" && !cpu_id_info.HasAVX())) { + LOGS_DEFAULT(WARNING) << "NUPHAR_CODEGEN_TARGET is not compatible with host machine." + "Target code will be generated, but exectuion will fail!"; + } + // For CPU, use target as host since the tvm_host_target_ is the one used to generate code in TVM + tvm_target_ = tvm::Target::create(codegen_target_->GetTargetName()); + tvm_host_target_ = tvm::Target::create(codegen_target_->GetTargetName()); + } else { + CreateTVMTarget(); + tvm_host_target_ = tvm::Target::create(GetCurrentHostTargetString()); + } - tvm_host_target_ = tvm::Target::create(GetCurrentHostTargetString()); tvm_ctx_.device_type = static_cast(tvm_target_->device_type); tvm_ctx_.device_id = 0; // use the default device id for CPU allocator diff --git a/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc b/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc index ad4c4d969a..5a9ca45b57 100644 --- a/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc +++ b/onnxruntime/test/providers/cpu/math/matmul_integer_test.cc @@ -72,19 +72,21 @@ void RunMatMulIntegerU8S8Test(const int M, const int N, const int K) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kNGraphExecutionProvider}); // currently nGraph provider does not support gemm_u8s8 } -TEST(MatmulIntegerOpTest, MatMulInteger_Uint8_Int8) { - // GEMV +TEST(MatmulIntegerOpTest, MatMulInteger_Uint8_Int8_Scalar) { RunMatMulIntegerU8S8Test(1, 1, 32); RunMatMulIntegerU8S8Test(1, 1, 260); RunMatMulIntegerU8S8Test(1, 1, 288); +} + +TEST(MatmulIntegerOpTest, MatMulInteger_Uint8_Int8_GEMV) { RunMatMulIntegerU8S8Test(1, 2, 16); RunMatMulIntegerU8S8Test(1, 2, 64); - // GEMM +} + +TEST(MatmulIntegerOpTest, MatMulInteger_Uint8_Int8_GEMM) { RunMatMulIntegerU8S8Test(2, 2, 40); RunMatMulIntegerU8S8Test(2, 48, 33); RunMatMulIntegerU8S8Test(2, 51, 40); - RunMatMulIntegerU8S8Test(6, 10, 34); - RunMatMulIntegerU8S8Test(8, 16, 64); } } // namespace test