saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Adjust codegen vectorization width from target (#2439)

Browse source 
* Adjust codegen vectorization width from target
This commit is contained in:
baowenlei 2019-11-20 13:28:15 -08:00 committed by GitHub
parent a5daa6faea
commit 9b7b5e2c27
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
16 changed files with 114 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
onnxruntime/core/codegen/common/utils.h
View file

@ -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 <cassert>
#include <memory>
@ -28,4 +29,18 @@ struct TargetFeature {
TargetFeature GetTargetInfo(const codegen::CodeGenSettings& setttings);
// GCD (Greatest Common Divisor)
template <typename T>
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

7
onnxruntime/core/codegen/passes/scheduler/schedule_utils.cc
View file

@ -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<int64_t>(natural_vector_size, (*tail_dim));
}
if (natural_vector_size > 1) {
tvm::IterVar xi, xo;
ctx.schedule[tensor->op].split(x, static_cast<int32_t>(natural_vector_size), &xo, &xi);
ctx.schedule[tensor->op].vectorize(xi);

6
onnxruntime/core/providers/nuphar/compiler/nuphar_schedule_builder.cc
View file

@ -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<CodeGenUnitStats>(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<CodeGenTargetX86*>(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);

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

@ -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<tvm::Tensor>& inputs,
const Node& node,
tvm_codegen::CodeGenContext&,
tvm_codegen::CodeGenContext& ctx_codegen,
tvm::Array<tvm::Tensor>& outputs) {
ProtoHelperNodeContext ctx(node);
OpNodeProtoHelper<ProtoHelperNodeContext> info(&ctx);
int64_t axis_i64;
ORT_RETURN_IF_ERROR(info.GetAttr<int64_t>("axis", &axis_i64));
axis_i64 = HandleNegativeAxis(axis_i64, gsl::narrow_cast<int64_t>(inputs[0]->shape.size()));
tvm::Tensor Y = nuphar::LogSoftmax(inputs[0], axis_i64);
CodeGenTargetX86* target = dynamic_cast<CodeGenTargetX86*>(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();
}

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

@ -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<tvm::Tensor>& inputs,
const Node& node,
tvm_codegen::CodeGenContext&,
tvm_codegen::CodeGenContext& ctx_codegen,
tvm::Array<tvm::Tensor>& outputs) {
ProtoHelperNodeContext ctx(node);
OpNodeProtoHelper<ProtoHelperNodeContext> info(&ctx);
int64_t axis_i64;
ORT_RETURN_IF_ERROR(info.GetAttr<int64_t>("axis", &axis_i64));
axis_i64 = HandleNegativeAxis(axis_i64, gsl::narrow_cast<int64_t>(inputs[0]->shape.size()));
tvm::Tensor Y = Softmax(inputs[0], axis_i64);
CodeGenTargetX86* target = dynamic_cast<CodeGenTargetX86*>(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();
}

4
onnxruntime/core/providers/nuphar/compiler/x86/scheduler/analysis_schedule.cc
View file

@ -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;
}

6
onnxruntime/core/providers/nuphar/compiler/x86/scheduler/nuphar_scheduler.h
View file

@ -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

26
onnxruntime/core/providers/nuphar/compiler/x86/scheduler/ort_type_schedule.cc
View file

@ -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 <tvm/tvm.h>
@ -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<CodeGenTargetX86*>(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;
}

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

@ -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<int32_t> 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<int32_t> 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<int64_t> embed_dim_vec,
const std::vector<int64_t> 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.");
}

4
onnxruntime/core/providers/nuphar/compiler/x86/scheduler/tvm_rule_schedule.cc
View file

@ -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;
}

4
onnxruntime/core/providers/nuphar/mti_x86/math/logsoftmax.cc
View file

@ -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

2
onnxruntime/core/providers/nuphar/mti_x86/math/logsoftmax.h
View file

@ -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

4
onnxruntime/core/providers/nuphar/mti_x86/math/softmax.cc
View file

@ -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

2
onnxruntime/core/providers/nuphar/mti_x86/math/softmax.h
View file

@ -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

7
onnxruntime/core/providers/nuphar/mti_x86/math/softmax_internal.cc
View file

@ -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<int64_t> 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");

2
onnxruntime/core/providers/nuphar/mti_x86/math/softmax_internal.h
View file

@ -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