@@ -308,7 +308,7 @@ class StoragePlanRewriter : public IRMutator {
308308 // for global memory it is nullptr, means beginning of everything.
309309 const Node* attach_scope_{nullptr };
310310 // The constant size of the buffer in bits, only used if it is constant
311- size_t const_nbits{0 };
311+ uint64_t const_nbits{0 };
312312 // The storage scope.
313313 StorageScope scope;
314314 // Allocs that shares this entry.
@@ -323,7 +323,7 @@ class StoragePlanRewriter : public IRMutator {
323323 Type elem_type;
324324 // This is non-zero if this allocate is folded into another one
325325 // the address becomes alloc_var + sizeof(elem_type) * elem_offset;
326- size_t elem_offset{0 };
326+ uint64_t elem_offset{0 };
327327 };
328328 // Remap the index
329329 Expr RemapIndex (Type dtype, Expr index, StorageEntry* e) {
@@ -410,7 +410,7 @@ class StoragePlanRewriter : public IRMutator {
410410 if (info.defined ()) {
411411 align = (info->max_simd_bits + e->elem_type .bits () - 1 ) / e->elem_type .bits ();
412412 }
413- size_t total_elem = e->const_nbits / e->elem_type .bits ();
413+ uint64_t total_elem = e->const_nbits / e->elem_type .bits ();
414414 if (total_elem % align != 0 ) {
415415 total_elem += align - (total_elem % align);
416416 }
@@ -426,7 +426,8 @@ class StoragePlanRewriter : public IRMutator {
426426 total_elem += align - (total_elem % align);
427427 }
428428 }
429- Expr alloc_size = make_const (e->allocs [0 ]->extents [0 ].type (), total_elem);
429+ Expr alloc_size = make_const (e->allocs [0 ]->extents [0 ].type (),
430+ total_elem);
430431 e->new_alloc = Allocate::make (
431432 e->alloc_var , e->elem_type , {alloc_size}, const_true (),
432433 Evaluate::make (0 ));
@@ -503,8 +504,8 @@ class StoragePlanRewriter : public IRMutator {
503504 const StorageScope& scope) {
504505 // skip plan for local variable,
505506 // compiler can do a better job with register allocation.
506- const size_t match_range = 16 ;
507- size_t const_nbits = static_cast <size_t >(
507+ const uint64_t match_range = 16 ;
508+ uint64_t const_nbits = static_cast <uint64_t >(
508509 op->constant_allocation_size () * op->type .bits () * op->type .lanes ());
509510 if (scope.rank > 1 || op->type .is_handle ()) {
510511 return NewAlloc (op, scope, const_nbits);
@@ -575,7 +576,7 @@ class StoragePlanRewriter : public IRMutator {
575576 // The allocation assign map
576577 std::unordered_map<const Variable*, StorageEntry*> alloc_map_;
577578 // constant size free map.
578- std::multimap<size_t , StorageEntry*> const_free_map_;
579+ std::multimap<uint64_t , StorageEntry*> const_free_map_;
579580 // symbolic free list, for non constant items.
580581 std::list<StorageEntry*> sym_free_list_;
581582 // The allocations
0 commit comments