diff --git a/onnxruntime/core/codegen/common/utils.h b/onnxruntime/core/codegen/common/utils.h index 008613c1c9..4e3e5c1fed 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/common/common.h" #include "core/providers/nuphar/common/nuphar_settings.h" #include #include @@ -28,4 +29,18 @@ struct TargetFeature { TargetFeature GetTargetInfo(const codegen::CodeGenSettings& setttings); +// GCD (Greatest Common Divisor) +template +T GCD(T a, T b) { + ORT_ENFORCE(a >= 0); + ORT_ENFORCE(b >= 0); + if (a < b) std::swap(a, b); + if (b == 0) return a; + while (a % b != 0) { + a = a % b; + std::swap(a, b); + } + return b; +} + } // namespace onnxruntime diff --git a/onnxruntime/core/codegen/passes/scheduler/schedule_utils.cc b/onnxruntime/core/codegen/passes/scheduler/schedule_utils.cc index 8f14852356..8e86f13eaa 100644 --- a/onnxruntime/core/codegen/passes/scheduler/schedule_utils.cc +++ b/onnxruntime/core/codegen/passes/scheduler/schedule_utils.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/codegen/passes/scheduler/schedule_utils.h" namespace onnxruntime { @@ -91,7 +92,11 @@ bool TryVectorization( auto axis = compute_op->axis; tvm::IterVar x = axis[rank - 1]; if ((*tail_dim) > natural_vector_size) { - if ((*tail_dim) % natural_vector_size == 0) { + if ((*tail_dim) % natural_vector_size != 0) { + natural_vector_size = GCD(natural_vector_size, (*tail_dim)); + } + + if (natural_vector_size > 1) { tvm::IterVar xi, xo; ctx.schedule[tensor->op].split(x, static_cast(natural_vector_size), &xo, &xi); ctx.schedule[tensor->op].vectorize(xi); diff --git a/onnxruntime/core/providers/nuphar/compiler/nuphar_schedule_builder.cc b/onnxruntime/core/providers/nuphar/compiler/nuphar_schedule_builder.cc index 57432e2b61..3f54d50a56 100644 --- a/onnxruntime/core/providers/nuphar/compiler/nuphar_schedule_builder.cc +++ b/onnxruntime/core/providers/nuphar/compiler/nuphar_schedule_builder.cc @@ -8,6 +8,7 @@ #include "core/codegen/passes/scheduler/tvm_schedule_builder.h" #include "core/providers/nuphar/common/analysis/subgraph_codegen_stats.h" +#include "core/providers/nuphar/compiler/x86/x86_target_info.h" // TODO change name space namespace onnxruntime { @@ -34,8 +35,9 @@ static void Traverse(const tvm::Tensor& tensor, Promote(ctx_codegen.GetGraphStats())->IsOutputNode(node); if (is_real_output) { - // TODO change it to the value from Target - int64_t natural_vector_size = 16; + CodeGenTargetX86* target = dynamic_cast(ctx_codegen.GetCodeGenHandle()->codegen_target); + ORT_ENFORCE(target != nullptr); + int64_t natural_vector_size = target->NaturalVectorWidth(tensor->dtype.bits()); TryVectorization(tensor, natural_vector_size, ctx_schedule); // to x86 InsertRootScheduleAndClosure(tensor, ctx_schedule); diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/logsoftmax.cc b/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/logsoftmax.cc index aef32e3d3c..aee2c5f78f 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/logsoftmax.cc +++ b/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/logsoftmax.cc @@ -3,9 +3,10 @@ #include "core/providers/nuphar/compiler/x86/op_ir_creator/all_ops.h" -#include "core/providers/nuphar/mti_x86/math/logsoftmax.h" #include "core/framework/op_kernel_info.h" #include "core/providers/common.h" +#include "core/providers/nuphar/compiler/x86/x86_target_info.h" +#include "core/providers/nuphar/mti_x86/math/logsoftmax.h" namespace onnxruntime { namespace nuphar { @@ -14,16 +15,20 @@ namespace nuphar { Status NUPHAR_TVM_X86_OP_IR_CREATOR_CLASS(LogSoftmax)::Evaluate( const tvm::Array& inputs, const Node& node, - tvm_codegen::CodeGenContext&, + tvm_codegen::CodeGenContext& ctx_codegen, tvm::Array& outputs) { ProtoHelperNodeContext ctx(node); OpNodeProtoHelper info(&ctx); int64_t axis_i64; ORT_RETURN_IF_ERROR(info.GetAttr("axis", &axis_i64)); - axis_i64 = HandleNegativeAxis(axis_i64, gsl::narrow_cast(inputs[0]->shape.size())); - tvm::Tensor Y = nuphar::LogSoftmax(inputs[0], axis_i64); + + CodeGenTargetX86* target = dynamic_cast(ctx_codegen.GetCodeGenHandle()->codegen_target); + ORT_ENFORCE(target != nullptr); + int64_t natural_vector_size = target->NaturalVectorWidth(inputs[0]->dtype.bits()); + + tvm::Tensor Y = nuphar::LogSoftmax(inputs[0], axis_i64, natural_vector_size); outputs.push_back(Y); return Status::OK(); } diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/softmax.cc b/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/softmax.cc index 28efe2cf1c..a3e1ea5077 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/softmax.cc +++ b/onnxruntime/core/providers/nuphar/compiler/x86/op_ir_creator/math/softmax.cc @@ -3,9 +3,10 @@ #include "core/providers/nuphar/compiler/x86/op_ir_creator/all_ops.h" -#include "core/providers/nuphar/mti_x86/math/softmax.h" #include "core/framework/op_kernel_info.h" #include "core/providers/common.h" +#include "core/providers/nuphar/compiler/x86/x86_target_info.h" +#include "core/providers/nuphar/mti_x86/math/softmax.h" namespace onnxruntime { namespace nuphar { @@ -14,16 +15,20 @@ namespace nuphar { Status NUPHAR_TVM_X86_OP_IR_CREATOR_CLASS(Softmax)::Evaluate( const tvm::Array& inputs, const Node& node, - tvm_codegen::CodeGenContext&, + tvm_codegen::CodeGenContext& ctx_codegen, tvm::Array& outputs) { ProtoHelperNodeContext ctx(node); OpNodeProtoHelper info(&ctx); int64_t axis_i64; ORT_RETURN_IF_ERROR(info.GetAttr("axis", &axis_i64)); - axis_i64 = HandleNegativeAxis(axis_i64, gsl::narrow_cast(inputs[0]->shape.size())); - tvm::Tensor Y = Softmax(inputs[0], axis_i64); + + CodeGenTargetX86* target = dynamic_cast(ctx_codegen.GetCodeGenHandle()->codegen_target); + ORT_ENFORCE(target != nullptr); + int64_t natural_vector_size = target->NaturalVectorWidth(inputs[0]->dtype.bits()); + + tvm::Tensor Y = Softmax(inputs[0], axis_i64, natural_vector_size); outputs.push_back(Y); return Status::OK(); } diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/analysis_schedule.cc b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/analysis_schedule.cc index 485aa328d7..14ce7e33aa 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/analysis_schedule.cc +++ b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/analysis_schedule.cc @@ -12,9 +12,9 @@ namespace nuphar { bool TVM_SCHEDULER_CLASS(True, NupharX86UseCount)::Evaluate( const tvm::Tensor& tensor, const Node*, - tvm_codegen::CodeGenContext&, + tvm_codegen::CodeGenContext& ctx_codegen, tvm_codegen::ScheduleContext& ctx_sched) { - bool status_vec = TryVectorizationX86(tensor, ctx_sched); + bool status_vec = TryVectorizationX86(tensor, ctx_codegen, ctx_sched); bool status_r_and_c = tvm_codegen::InsertRootScheduleAndClosure(tensor, ctx_sched); return status_vec || status_r_and_c; } diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/nuphar_scheduler.h b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/nuphar_scheduler.h index 766d251235..a41642e7e3 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/nuphar_scheduler.h +++ b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/nuphar_scheduler.h @@ -31,11 +31,13 @@ DECLARE_TVM_SCHEDULER_CLASS(True, NupharX86PartialResult) // utilities bool TryVectorizationX86( const tvm::Tensor& tensor, - tvm_codegen::ScheduleContext& ctx); + tvm_codegen::CodeGenContext& ctx_codegen, + tvm_codegen::ScheduleContext& ctx_sched); bool InputRootScheduleWithVectorizationX86( const tvm::Tensor& tensor, - tvm_codegen::ScheduleContext& ctx); + tvm_codegen::CodeGenContext& ctx_codegen, + tvm_codegen::ScheduleContext& ctx_sched); } // namespace nuphar } // namespace onnxruntime diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/ort_type_schedule.cc b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/ort_type_schedule.cc index 2acdd99582..d32e35796a 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/ort_type_schedule.cc +++ b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/ort_type_schedule.cc @@ -8,6 +8,7 @@ #include "core/codegen/passes/scheduler/schedule_utils.h" #include "core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_ll_extern.h" #include "core/providers/nuphar/compiler/x86/scheduler/tensorize/intrin_gemv_ll_ir.h" +#include "core/providers/nuphar/compiler/x86/x86_target_info.h" #include "core/framework/op_kernel_info.h" #include @@ -16,21 +17,24 @@ namespace nuphar { bool TryVectorizationX86( const tvm::Tensor& tensor, - tvm_codegen::ScheduleContext& ctx) { - // TODO change it to the value from Target - int64_t natural_vector_size = 16; + tvm_codegen::CodeGenContext& ctx_codegen, + tvm_codegen::ScheduleContext& ctx_sched) { + CodeGenTargetX86* target = dynamic_cast(ctx_codegen.GetCodeGenHandle()->codegen_target); + ORT_ENFORCE(target != nullptr); + int64_t natural_vector_size = target->NaturalVectorWidth(tensor->dtype.bits()); - return TryVectorization(tensor, natural_vector_size, ctx); + return TryVectorization(tensor, natural_vector_size, ctx_sched); } bool InputRootScheduleWithVectorizationX86( const tvm::Tensor& tensor, - tvm_codegen::ScheduleContext& ctx) { + tvm_codegen::CodeGenContext& ctx_codegen, + tvm_codegen::ScheduleContext& ctx_sched) { bool status = false; for (auto& t : tensor->op->InputTensors()) { if (t->op->InputTensors().size() > 0) { - bool status_vec = TryVectorizationX86(t, ctx); - bool status_root = InsertRootSchedule(t, ctx); + bool status_vec = TryVectorizationX86(t, ctx_codegen, ctx_sched); + bool status_root = InsertRootSchedule(t, ctx_sched); status = status || status_root || status_vec; } } @@ -40,13 +44,13 @@ bool InputRootScheduleWithVectorizationX86( bool TVM_SCHEDULER_CLASS(Softmax, NupharX86OrtOpType)::Evaluate( const tvm::Tensor& tensor, const Node*, - tvm_codegen::CodeGenContext&, + tvm_codegen::CodeGenContext& ctx_codegen, tvm_codegen::ScheduleContext& ctx_sched) { bool status_softmax_itself = TryInlineSchedule(tensor, ctx_sched); // compute root the exp since it is reused more than once auto& tensor_exp = tensor->op->InputTensors()[0]; - bool status_vec = TryVectorizationX86(tensor_exp, ctx_sched); + bool status_vec = TryVectorizationX86(tensor_exp, ctx_codegen, ctx_sched); bool status_root = InsertRootSchedule(tensor_exp, ctx_sched); return status_softmax_itself || status_vec || status_root; } @@ -54,14 +58,14 @@ bool TVM_SCHEDULER_CLASS(Softmax, NupharX86OrtOpType)::Evaluate( bool TVM_SCHEDULER_CLASS(Split, NupharX86OrtOpType)::Evaluate( const tvm::Tensor& tensor, const Node*, - tvm_codegen::CodeGenContext&, + tvm_codegen::CodeGenContext& ctx_codegen, tvm_codegen::ScheduleContext& ctx_sched) { auto& tensor_split_input = tensor->op->InputTensors()[0]; // force inline for split since to avoid extra copy bool status_split_itself = TryInlineSchedule(tensor, ctx_sched); // add root for split's inputs to avoid inline of the inputs - bool status_vec = TryVectorizationX86(tensor_split_input, ctx_sched); + bool status_vec = TryVectorizationX86(tensor_split_input, ctx_codegen, ctx_sched); bool status_input_root = InsertRootSchedule(tensor_split_input, ctx_sched); return status_split_itself || status_vec || status_input_root; } 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 34278804d4..a0b8a8a4b9 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize_schedule.cc +++ b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tensorize_schedule.cc @@ -24,10 +24,11 @@ constexpr int bits_per_byte = 8; static Status TensorizeGEMVInteger16(const tvm::Tensor& tensor, const int64_t input_dim, - tvm_codegen::ScheduleContext& ctx) { + tvm_codegen::CodeGenContext& ctx_codegen, + tvm_codegen::ScheduleContext& ctx_sched) { // schedule for imatmul inputs - InsertRootScheduleAndClosure(tensor, ctx); - InputRootScheduleWithVectorizationX86(tensor, ctx); + InsertRootScheduleAndClosure(tensor, ctx_sched); + InputRootScheduleWithVectorizationX86(tensor, ctx_codegen, ctx_sched); // decide kernel shape std::vector kernel_shape; @@ -46,11 +47,11 @@ static Status TensorizeGEMVInteger16(const tvm::Tensor& tensor, auto y = xy[1]; auto z = compute_op->reduce_axis[0]; tvm::IterVar yo, yi; - ctx.schedule[tensor->op].split(y, shape[0], &yo, &yi); + ctx_sched.schedule[tensor->op].split(y, shape[0], &yo, &yi); tvm::IterVar zo, zi; - ctx.schedule[tensor->op].split(z, shape[1], &zo, &zi); - ctx.schedule[tensor->op].reorder({x, yo, zo, yi, zi}); - ctx.schedule[tensor->op].tensorize(yi, igemv16bit.CreateTensorIntrin()); + 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, igemv16bit.CreateTensorIntrin()); return Status::OK(); } @@ -58,10 +59,11 @@ static Status TensorizeGEMVInteger16(const tvm::Tensor& tensor, // TODO: refactor below function static Status TensorizeGEMVInteger(const tvm::Tensor& tensor, const int64_t input_dim, - tvm_codegen::ScheduleContext& ctx) { + tvm_codegen::CodeGenContext& ctx_codegen, + tvm_codegen::ScheduleContext& ctx_sched) { // schedule for imatmul inputs - InsertRootScheduleAndClosure(tensor, ctx); - InputRootScheduleWithVectorizationX86(tensor, ctx); + InsertRootScheduleAndClosure(tensor, ctx_sched); + InputRootScheduleWithVectorizationX86(tensor, ctx_codegen, ctx_sched); // decide kernel shape std::vector kernel_shape; @@ -82,24 +84,25 @@ static Status TensorizeGEMVInteger(const tvm::Tensor& tensor, auto y = xy[1]; auto z = compute_op->reduce_axis[0]; tvm::IterVar yo, yi; - ctx.schedule[tensor->op].split(y, shape[0], &yo, &yi); + ctx_sched.schedule[tensor->op].split(y, shape[0], &yo, &yi); tvm::IterVar zo, zi; - ctx.schedule[tensor->op].split(z, shape[1], &zo, &zi); - ctx.schedule[tensor->op].reorder({x, yo, zo, yi, zi}); - ctx.schedule[tensor->op].tensorize(yi, igemv8bit.CreateTensorIntrin()); + 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::ScheduleContext& ctx, + tvm_codegen::CodeGenContext& ctx_codegen, + tvm_codegen::ScheduleContext& ctx_sched, bool tensorize, const std::string& target_str) { // Schedule tensor and inputs as root - bool status_imatmul = InsertRootScheduleAndClosure(tensor, ctx); + bool status_imatmul = InsertRootScheduleAndClosure(tensor, ctx_sched); if (status_imatmul == false) return Status::OK(); - InputRootScheduleWithVectorizationX86(tensor, ctx); + InputRootScheduleWithVectorizationX86(tensor, ctx_codegen, ctx_sched); // Default tiling size // TODO: tuning tiling sizes later @@ -119,31 +122,31 @@ static Status TensorizeIGEMV(const tvm::Tensor& tensor, // no tiling need for IterVar x tvm::IterVar yo, yi; - ctx.schedule[tensor->op].split(y, kernel_shape[0], &yo, &yi); + ctx_sched.schedule[tensor->op].split(y, kernel_shape[0], &yo, &yi); tvm::IterVar zo, zi; - ctx.schedule[tensor->op].split(z, kernel_shape[1], &zo, &zi); - ctx.schedule[tensor->op].reorder({x, yo, zo, yi, 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.schedule[tensor->op].tensorize(yi, igemv8bit.CreateTensorIntrin()); + ctx_sched.schedule[tensor->op].tensorize(yi, igemv8bit.CreateTensorIntrin()); } return Status::OK(); } static Status TensorizeIGEMM(const tvm::Tensor& tensor, tvm_codegen::CodeGenContext& ctx_codegen, - tvm_codegen::ScheduleContext& ctx, + tvm_codegen::ScheduleContext& ctx_sched, 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); + bool status_imatmul = InsertRootScheduleAndClosure(tensor, ctx_sched); if (status_imatmul == false) return Status::OK(); - InputRootScheduleWithVectorizationX86(tensor, ctx); + InputRootScheduleWithVectorizationX86(tensor, ctx_codegen, ctx_sched); // Default tiling size int tensorize_batch = 4; @@ -212,23 +215,23 @@ static Status TensorizeIGEMM(const tvm::Tensor& tensor, auto y = xy[1]; auto z = compute_op->reduce_axis[0]; tvm::IterVar xo, xi; - ctx.schedule[tensor->op].split(x, kernel_shape[0], &xo, &xi); + ctx_sched.schedule[tensor->op].split(x, kernel_shape[0], &xo, &xi); tvm::IterVar yo, yi; - ctx.schedule[tensor->op].split(y, kernel_shape[1], &yo, &yi); + ctx_sched.schedule[tensor->op].split(y, kernel_shape[1], &yo, &yi); tvm::IterVar zo, zi; - ctx.schedule[tensor->op].split(z, kernel_shape[2], &zo, &zi); + ctx_sched.schedule[tensor->op].split(z, kernel_shape[2], &zo, &zi); // Loop nest permutation if (settings.HasOption(kNupharTensorize_IGEMM_Permute) && (settings.OptionMatches(kNupharTensorize_IGEMM_Permute, kNupharTensorize_IGEMM_Permute_All) || settings.OptionMatches(kNupharTensorize_IGEMM_Permute, kNupharTensorize_IGEMM_Permute_Outer))) { - ctx.schedule[tensor->op].reorder({yo, xo, zo, xi, yi, zi}); + ctx_sched.schedule[tensor->op].reorder({yo, xo, zo, xi, yi, zi}); } else { // Loop nest default order if (target_str == "avx") - ctx.schedule[tensor->op].reorder({yo, xo, zo, xi, yi, zi}); + ctx_sched.schedule[tensor->op].reorder({yo, xo, zo, xi, yi, zi}); else - ctx.schedule[tensor->op].reorder({xo, yo, zo, xi, yi, zi}); + ctx_sched.schedule[tensor->op].reorder({xo, yo, zo, xi, yi, zi}); } // Natural vector width @@ -269,7 +272,7 @@ static Status TensorizeIGEMM(const tvm::Tensor& tensor, igemm8bit.InsertTensorizeDimInfo("n", embed_meta); igemm8bit.InsertTensorizeDimInfo("k", input_meta); // Bind tensorization kernel - ctx.schedule[tensor->op].tensorize(xi, igemm8bit.CreateTensorIntrin()); + ctx_sched.schedule[tensor->op].tensorize(xi, igemm8bit.CreateTensorIntrin()); return Status::OK(); } @@ -322,7 +325,7 @@ static bool IMatMulTensorizeSchedule( bool status_tensorize = true; if (is8bit) { if (feature.hasAVX512) { // isAVX512 - status_tensorize = is_scalar ? TensorizeIGEMV(imatmul, ctx_sched, /*tensorize=*/false, "avx512-skylake").IsOK() + status_tensorize = is_scalar ? TensorizeIGEMV(imatmul, ctx_codegen, 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}, @@ -331,14 +334,14 @@ static bool IMatMulTensorizeSchedule( } 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() + 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(); } else if (feature.hasAVX) { // isAVX - status_tensorize = is_scalar ? TensorizeIGEMV(imatmul, ctx_sched, /*tensorize=*/false, "avx").IsOK() + status_tensorize = is_scalar ? TensorizeIGEMV(imatmul, ctx_codegen, 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}, @@ -353,7 +356,7 @@ static bool IMatMulTensorizeSchedule( // 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(); + status_tensorize = TensorizeGEMVInteger16(imatmul, *p_input_dim, ctx_codegen, ctx_sched).IsOK(); } else { ORT_NOT_IMPLEMENTED("Not supported target in 16bit Tensorization."); } diff --git a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tvm_rule_schedule.cc b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tvm_rule_schedule.cc index cd2811b0f8..646f22d7d7 100644 --- a/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tvm_rule_schedule.cc +++ b/onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tvm_rule_schedule.cc @@ -13,10 +13,10 @@ namespace nuphar { bool TVM_SCHEDULER_CLASS(Extern, NupharX86TVMRule)::Evaluate( const tvm::Tensor& tensor, const Node*, - tvm_codegen::CodeGenContext&, + tvm_codegen::CodeGenContext& ctx_codegen, tvm_codegen::ScheduleContext& ctx_sched) { bool status = InsertRootScheduleAndClosure(tensor, ctx_sched); - bool status_input = InputRootScheduleWithVectorizationX86(tensor, ctx_sched); + bool status_input = InputRootScheduleWithVectorizationX86(tensor, ctx_codegen, ctx_sched); return status || status_input; } diff --git a/onnxruntime/core/providers/nuphar/mti_x86/math/logsoftmax.cc b/onnxruntime/core/providers/nuphar/mti_x86/math/logsoftmax.cc index a4a57fdc9b..820aaee02e 100644 --- a/onnxruntime/core/providers/nuphar/mti_x86/math/logsoftmax.cc +++ b/onnxruntime/core/providers/nuphar/mti_x86/math/logsoftmax.cc @@ -8,8 +8,8 @@ namespace onnxruntime { namespace nuphar { -tvm::Tensor LogSoftmax(const tvm::Tensor& input, int64_t axis, const std::string& name) { - return internal::SoftmaxInternal(input, axis, name, /*logarithmic*/ true); +tvm::Tensor LogSoftmax(const tvm::Tensor& input, int64_t axis, int64_t vector_width, const std::string& name) { + return internal::SoftmaxInternal(input, axis, vector_width, name, /*logarithmic*/ true); } } // namespace nuphar diff --git a/onnxruntime/core/providers/nuphar/mti_x86/math/logsoftmax.h b/onnxruntime/core/providers/nuphar/mti_x86/math/logsoftmax.h index ea533ff4f2..9d5344b2d4 100644 --- a/onnxruntime/core/providers/nuphar/mti_x86/math/logsoftmax.h +++ b/onnxruntime/core/providers/nuphar/mti_x86/math/logsoftmax.h @@ -9,7 +9,7 @@ namespace onnxruntime { namespace nuphar { -tvm::Tensor LogSoftmax(const tvm::Tensor& input, int64_t axis, const std::string& name = "LogSoftmax"); +tvm::Tensor LogSoftmax(const tvm::Tensor& input, int64_t axis, int64_t vector_width, const std::string& name = "LogSoftmax"); } // namespace nuphar } // namespace onnxruntime diff --git a/onnxruntime/core/providers/nuphar/mti_x86/math/softmax.cc b/onnxruntime/core/providers/nuphar/mti_x86/math/softmax.cc index 7f2926f329..acb993713a 100644 --- a/onnxruntime/core/providers/nuphar/mti_x86/math/softmax.cc +++ b/onnxruntime/core/providers/nuphar/mti_x86/math/softmax.cc @@ -8,8 +8,8 @@ namespace onnxruntime { namespace nuphar { -tvm::Tensor Softmax(const tvm::Tensor& input, int64_t axis, const std::string& name) { - return internal::SoftmaxInternal(input, axis, name, /*logarithmic*/ false); +tvm::Tensor Softmax(const tvm::Tensor& input, int64_t axis, int64_t vector_width, const std::string& name) { + return internal::SoftmaxInternal(input, axis, vector_width, name, /*logarithmic*/ false); } } // namespace nuphar diff --git a/onnxruntime/core/providers/nuphar/mti_x86/math/softmax.h b/onnxruntime/core/providers/nuphar/mti_x86/math/softmax.h index ae794ad029..104b2d4ea4 100644 --- a/onnxruntime/core/providers/nuphar/mti_x86/math/softmax.h +++ b/onnxruntime/core/providers/nuphar/mti_x86/math/softmax.h @@ -9,7 +9,7 @@ namespace onnxruntime { namespace nuphar { -tvm::Tensor Softmax(const tvm::Tensor& input, int64_t axis, const std::string& name = "Softmax"); +tvm::Tensor Softmax(const tvm::Tensor& input, int64_t axis, int64_t vector_width, const std::string& name = "Softmax"); } // namespace nuphar } // namespace onnxruntime diff --git a/onnxruntime/core/providers/nuphar/mti_x86/math/softmax_internal.cc b/onnxruntime/core/providers/nuphar/mti_x86/math/softmax_internal.cc index 1e3c3822c2..78ecbf48d9 100644 --- a/onnxruntime/core/providers/nuphar/mti_x86/math/softmax_internal.cc +++ b/onnxruntime/core/providers/nuphar/mti_x86/math/softmax_internal.cc @@ -32,16 +32,15 @@ static int32_t FuseDim(int64_t axis) { tvm::Tensor SoftmaxInternal(const tvm::Tensor& input, int64_t axis, + int64_t vector_width, const std::string& name, bool logarithmic) { std ::vector reduce_axis = ReduceAxes(axis, input->shape.size()); - // TODO use natural vector size check by type later - int32_t vectorization_width = 16; int32_t fuse_dim = FuseDim(axis); - auto max_element = ReduceMax(input, reduce_axis, true, vectorization_width, false, fuse_dim, name + "_reduce_max"); + auto max_element = ReduceMax(input, reduce_axis, true, vector_width, false, fuse_dim, name + "_reduce_max"); auto x_shift = tvm_codegen::Sub(input, max_element, name + "_sub"); auto exp_x = nuphar::Exp(x_shift, name + "_exp"); - auto exp_x_sum = ReduceSum(exp_x, reduce_axis, true, vectorization_width, false, fuse_dim, name + "_reduce_sum"); + auto exp_x_sum = ReduceSum(exp_x, reduce_axis, true, vector_width, false, fuse_dim, name + "_reduce_sum"); if (logarithmic) { auto log_exp_x_sum = nuphar::Log(exp_x_sum, name + "_log"); diff --git a/onnxruntime/core/providers/nuphar/mti_x86/math/softmax_internal.h b/onnxruntime/core/providers/nuphar/mti_x86/math/softmax_internal.h index ef0202526a..a983a01da5 100644 --- a/onnxruntime/core/providers/nuphar/mti_x86/math/softmax_internal.h +++ b/onnxruntime/core/providers/nuphar/mti_x86/math/softmax_internal.h @@ -10,7 +10,7 @@ namespace onnxruntime { namespace nuphar { namespace internal { -tvm::Tensor SoftmaxInternal(const tvm::Tensor& input, int64_t axis, const std::string& name, bool logarithmic); +tvm::Tensor SoftmaxInternal(const tvm::Tensor& input, int64_t axis, int64_t vector_width, const std::string& name, bool logarithmic); } // namespace internal } // namespace nuphar