colin
/
TIM-VX
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix crash when eletwise inputs are different rank (#665) 

Fix crash in AlignPermuteVectorForElmentWise() if inputs tensor have different rank

Type: Bug fix

Signed-off-by: Chen <jack.chen@verisilicon.com>
This commit is contained in:
chxin66 2023-12-06 17:15:15 +08:00 committed by GitHub
parent e013cf0a65
commit 720f0a485a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 52 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
src/tim/transform/ops/elementwise_layout_inference.h
View File

@ -42,6 +42,32 @@ class ElementWiseLayoutInfer : public OpLayoutInfer {
void OnInputs(
std::vector<std::shared_ptr<vx::Tensor>>& next_tensors) override {
auto in_0 = op_->impl()->InputsTensor()[0];
auto in_1 = op_->impl()->InputsTensor()[1];
std::shared_ptr<tim::vx::Tensor> short_tensor =
in_0->GetShape().size() > in_1->GetShape().size() ? in_1 : in_0;
std::shared_ptr<tim::vx::Tensor> long_tensor =
in_0->GetShape().size() < in_1->GetShape().size() ? in_1 : in_0;
if (in_0->GetSpec().attr_ != tim::vx::CONSTANT &&
in_1->GetSpec().attr_ != tim::vx::CONSTANT &&
in_0->GetShape().size() != in_1->GetShape().size()) {
auto pv_long = context_->GetPermuteVector(long_tensor);
auto pv_short = context_->GetPermuteVector(short_tensor);
auto rank_long = pv_long->Rank();
auto rank_short = pv_short->Rank();
auto expanded_pv = MakeShared(rank_long);
// if different size, expand short pv to long pv
for (uint32_t i = 0; i < rank_short; ++i) {
expanded_pv->At(i) = pv_short->At(i); // replace low dims with short pv
}
std::vector<uint32_t> expanded_shape(short_tensor->GetShape());
for (uint32_t i = 0; i < rank_long; ++i) {
if (i >= rank_short) expanded_shape.push_back(1);
}
short_tensor->GetSpec().SetShape(expanded_shape);
context_->SetPermuteVector(short_tensor, expanded_pv); // set new expand pv
}
auto required_pv = AlignPermuteVectorForElementWise();
auto elementwise = context_->infer_graph_->CreateOperation<OpType>();
for (const auto& i_src : op_->impl()->InputsTensor()) {
@ -63,6 +89,32 @@ class MultiplyLayoutInfer : public OpLayoutInfer {
void OnInputs(
std::vector<std::shared_ptr<vx::Tensor>>& next_tensors) override {
auto in_0 = op_->impl()->InputsTensor()[0];
auto in_1 = op_->impl()->InputsTensor()[1];
std::shared_ptr<tim::vx::Tensor> short_tensor =
in_0->GetShape().size() > in_1->GetShape().size() ? in_1 : in_0;
std::shared_ptr<tim::vx::Tensor> long_tensor =
in_0->GetShape().size() < in_1->GetShape().size() ? in_1 : in_0;
if (in_0->GetSpec().attr_ != tim::vx::CONSTANT &&
in_1->GetSpec().attr_ != tim::vx::CONSTANT &&
in_0->GetShape().size() != in_1->GetShape().size()) {
auto pv_long = context_->GetPermuteVector(long_tensor);
auto pv_short = context_->GetPermuteVector(short_tensor);
auto rank_long = pv_long->Rank();
auto rank_short = pv_short->Rank();
auto expanded_pv = MakeShared(rank_long);
// if different size, expand short pv to long pv
for (uint32_t i = 0; i < rank_short; ++i) {
expanded_pv->At(i) = pv_short->At(i); // replace low dims with short pv
}
std::vector<uint32_t> expanded_shape(short_tensor->GetShape());
for (uint32_t i = 0; i < rank_long; ++i) {
if (i >= rank_short) expanded_shape.push_back(1);
}
short_tensor->GetSpec().SetShape(expanded_shape);
context_->SetPermuteVector(short_tensor, expanded_pv); // set new expand pv
}
auto required_pv = AlignPermuteVectorForElementWise();
auto multiply =
context_->infer_graph_->CreateOperation<tim::vx::ops::Multiply>(