Commit 5cb729ec by Haichen Shen Committed by Tianqi Chen

[BUGFIX] Seg fault in memory planing for symbolic shape (#2317)

parent 97ca4031
...@@ -196,7 +196,7 @@ class StorageAllocator : public StorageAllocaBaseVisitor { ...@@ -196,7 +196,7 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
protected: protected:
using StorageAllocaBaseVisitor::VisitExpr_; using StorageAllocaBaseVisitor::VisitExpr_;
// override create token by getting token as prototype requirements. // override create token by getting token as prototype requirements.
void CreateToken(const ExprNode* op, bool can_realloc) final { void CreateToken(const ExprNode* op, bool can_realloc) final {
CHECK(!token_map_.count(op)); CHECK(!token_map_.count(op));
auto it = prototype_.find(op); auto it = prototype_.find(op);
CHECK(it != prototype_.end()); CHECK(it != prototype_.end());
...@@ -253,12 +253,12 @@ class StorageAllocator : public StorageAllocaBaseVisitor { ...@@ -253,12 +253,12 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
size_t size = 1; size_t size = 1;
for (IndexExpr dim : ttype->shape) { for (IndexExpr dim : ttype->shape) {
const int64_t* pval = as_const_int(dim); const int64_t* pval = as_const_int(dim);
CHECK_GE(*pval, 0) <<
"can not allocate memory for tensor with negative shape" <<
*pval;
CHECK(pval != nullptr) CHECK(pval != nullptr)
<< "Cannot allocate memory symbolic tensor shape " << "Cannot allocate memory symbolic tensor shape "
<< ttype->shape; << ttype->shape;
CHECK_GE(*pval, 0)
<< "Cannot allocate memory for tensor with negative shape"
<< *pval;
size *= static_cast<size_t>(pval[0]); size *= static_cast<size_t>(pval[0]);
} }
size *= DivRoundUp(ttype->dtype.bits() * ttype->dtype.lanes(), 8); size *= DivRoundUp(ttype->dtype.bits() * ttype->dtype.lanes(), 8);
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment