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

Add parallel for tensorized gemm (#2517)

Browse source 
* add parallel for tensorize gemm

* add option to control parallel

* change to a more clean way to control
This commit is contained in:
baowenlei 2019-12-02 22:05:46 -08:00 committed by KeDengMS
parent c1be615c45
commit 25c260fdef
1 changed files with 18 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -221,17 +221,25 @@ static Status TensorizeIGEMM(const tvm::Tensor& tensor,
tvm::IterVar zo, zi;
ctx_sched.schedule[tensor->op].split(z, kernel_shape[2], &zo, &zi);
tvm::Array<tvm::IterVar> fused_axis;
// 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_sched.schedule[tensor->op].reorder({yo, xo, zo, xi, yi, zi});
fused_axis.push_back(yo);
fused_axis.push_back(xo);
} else {
// Loop nest default order
if (target_str == "avx")
if (target_str == "avx") {
ctx_sched.schedule[tensor->op].reorder({yo, xo, zo, xi, yi, zi});
else
fused_axis.push_back(yo);
fused_axis.push_back(xo);
} else {
ctx_sched.schedule[tensor->op].reorder({xo, yo, zo, xi, yi, zi});
fused_axis.push_back(yo);
fused_axis.push_back(xo);
}
}
// Natural vector width
@ -274,6 +282,14 @@ static Status TensorizeIGEMM(const tvm::Tensor& tensor,
// Bind tensorization kernel
ctx_sched.schedule[tensor->op].tensorize(xi, igemm8bit.CreateTensorIntrin());
tvm::IterVar parallel_axis;
ctx_sched.schedule[tensor->op].fuse(fused_axis, &parallel_axis);
int64_t workloads_threshold = Promote<NupharCodeGenCtx>(&ctx_codegen)->GetCodeGenHandle()->parallel_min_workloads;
if (workloads_threshold > 0) {
ctx_sched.schedule[tensor->op].parallel(parallel_axis);
}
return Status::OK();
}