3131
3232namespace tvm {
3333namespace relay {
34+ constexpr auto Is2DStorage = runtime::IsTextureStorage;
3435
35- size_t TokenAllocator1D::GetMemorySize (StorageToken* prototype) {
36+ /*
37+ * Mixed mode memory allocator
38+ */
39+ size_t TokenAllocatorMixed::GetMemorySize (StorageToken* prototype) {
3640 TensorType ttype = prototype->ttype ;
3741 ICHECK (ttype.defined ());
3842 size_t size = 1 ;
39- for (IndexExpr dim : ttype->shape ) {
40- const int64_t * pval = tir::as_const_int (dim);
41- ICHECK (pval != nullptr ) << " Cannot allocate memory symbolic tensor shape " shape ;
42- ICHECK_GE (*pval, 0 ) << " Cannot allocate memory for tensor with negative shape"
43- size *= static_cast <size_t >(pval[0 ]);
43+ if (relay::Is2DStorage (prototype->virtual_device ->memory_scope )) {
44+ size = GetSize2D (prototype);
45+ } else {
46+ for (IndexExpr dim : ttype->shape ) {
47+ const int64_t * pval = tir::as_const_int (dim);
48+ ICHECK (pval != nullptr ) << " Cannot allocate memory symbolic tensor shape " shape ;
49+ ICHECK_GE (*pval, 0 ) << " Cannot allocate memory for tensor with negative shape"
50+ size *= static_cast <size_t >(pval[0 ]);
51+ }
52+ size *= DivRoundUp (ttype->dtype .bits () * ttype->dtype .lanes (), 8 );
4453 }
45- size *= DivRoundUp (ttype->dtype .bits () * ttype->dtype .lanes (), 8 );
4654 return size;
4755}
4856
49- StorageToken* TokenAllocator1D ::Request
57+ StorageToken* TokenAllocatorMixed ::Request
5058 // calculate the size;
5159 size_t size = GetMemorySize (prototype);
5260 // search memory block in [size / match_range_, size * match_range_)
@@ -56,145 +64,77 @@ StorageToken* TokenAllocator1D::Request(StorageToken* prototype) {
5664 auto begin = free_.lower_bound (size / match_range_);
5765 auto mid = free_.lower_bound (size);
5866 auto end = free_.upper_bound (size * match_range_);
67+ auto prototype_keys = prototype->virtual_device ->target ->GetKeys ();
68+ bool is_prototype_adreno =
69+ std::find (prototype_keys.begin (), prototype_keys.end (), " adreno" end ();
5970 // search for memory blocks larger than requested
6071 for (auto it = mid; it != end; ++it) {
6172 StorageToken* tok = it->second ;
62- if (!tok->is_compatible (*prototype)) continue ;
63- ICHECK_EQ (tok->ref_counter , 0 );
64- // Use exect matching strategy
65- tok->max_bytes = std::max (size, tok->max_bytes );
66- tok->ref_counter = prototype->ref_counter ;
67- // find a exact match, erase from map and return
68- free_.erase (it);
69- return tok;
73+ // TODO(Siva): We need a additional ways of comparing VirtualDevice
74+ auto tok_keys = tok->virtual_device ->target ->GetKeys ();
75+ bool is_tok_adreno = std::find (tok_keys.begin (), tok_keys.end (), " adreno" end ();
76+
77+ if (tok->is_compatible (*prototype) || (is_prototype_adreno && is_tok_adreno)) {
78+ ICHECK_EQ (tok->ref_counter , 0 );
79+ // Use exect matching strategy
80+ tok->max_bytes = std::max (size, tok->max_bytes );
81+ tok->ref_counter = prototype->ref_counter ;
82+ // find a exact match, erase from map and return
83+ free_.erase (it);
84+ return tok;
85+ }
7086 }
7187 // then search for memory blocks smaller than requested space
7288 for (auto it = mid; it != begin;) {
7389 --it;
7490 StorageToken* tok = it->second ;
75- if (!tok->is_compatible (*prototype)) continue ;
76- ICHECK_EQ (tok->ref_counter , 0 );
77- // Use exect matching strategy
78- tok->max_bytes = std::max (size, tok->max_bytes );
79- tok->ref_counter = prototype->ref_counter ;
80- // erase from map and return
81- free_.erase (it);
82- return tok;
91+ auto tok_keys = tok->virtual_device ->target ->GetKeys ();
92+ bool is_tok_adreno = std::find (tok_keys.begin (), tok_keys.end (), " adreno" end ();
93+ if (tok->is_compatible (*prototype) || (is_prototype_adreno && is_tok_adreno)) {
94+ ICHECK_EQ (tok->ref_counter , 0 );
95+ // Use exect matching strategy
96+ tok->max_bytes = std::max (size, tok->max_bytes );
97+ tok->ref_counter = prototype->ref_counter ;
98+ // erase from map and return
99+ free_.erase (it);
100+ return tok;
101+ }
83102 }
84103 return nullptr ;
85104}
86105
87- StorageToken* TokenAllocator1D ::Allocint64_t storage_id) {
106+ StorageToken* TokenAllocatorMixed ::Allocint64_t storage_id) {
88107 size_t size = GetMemorySize (prototype);
89108 prototype->max_bytes = size;
90109 prototype->storage_id = storage_id;
91110 data_.push_back (prototype);
92111 return prototype;
93112}
94113
95- void TokenAllocator1D ::CheckForRelease
114+ void TokenAllocatorMixed ::CheckForRelease
96115 ICHECK_GE (tok->storage_id , 0 );
97116 ICHECK_GE (tok->ref_counter , 0 );
98117 if (tok->ref_counter == 0 ) {
99118 free_.insert ({tok->max_bytes , tok});
100119 }
101120}
102121
103- StorageToken* TokenAllocator2D::Request (StorageToken* prototype) {
104- auto shape = GetSize2D (prototype);
105- const int64_t max_ratio = 5 ;
106- int64_t min_added_size_x = std::numeric_limits<int64_t >::max ();
107- int64_t min_added_size_y = std::numeric_limits<int64_t >::max ();
108- int64_t min_wasted_size_x = std::numeric_limits<int64_t >::max ();
109- int64_t min_wasted_size_y = std::numeric_limits<int64_t >::max ();
110- int64_t best_storage_id = -1 ;
111- MemBlock new_mem;
112- for (int64_t free_id : free_list_) {
113- MemBlock& cached = blocks_[free_id];
114- // Can only reuse texture 2d blocks of the same type
115- if (cached.token_ ->ttype ->dtype != prototype->ttype ->dtype ) {
116- continue ;
117- }
118- // Can only reuse texture 2d blocks of the same scope
119- // Because reusing textures with different memory scope may lead to
120- // accuracy issues, because the data will be packed in a different way for
121- // different memory scopes.
122- if (cached.token_ ->virtual_device ->memory_scope != prototype->virtual_device ->memory_scope ) {
123- continue ;
124- }
125- // avoid reusing too small and too big textures
126- if (shape.width / cached.x_ > max_ratio || cached.x_ / shape.width > max_ratio ||
127- shape.height / cached.y_ > max_ratio || cached.y_ / shape.height > max_ratio) {
128- continue ;
129- }
130- int64_t new_width = std::max (cached.x_ , shape.width );
131- int64_t new_height = std::max (cached.y_ , shape.height );
132- int64_t added_size_x = new_width - cached.x_ ;
133- int64_t added_size_y = new_height - cached.y_ ;
134- int64_t wasted_size_x = new_width - shape.width ;
135- int64_t wasted_size_y = new_height - shape.height ;
136- // Prioritize minimization of added size first, then minimize
137- // wasted size among blocks which would not require expansion
138- if ((min_added_size_x > 0 && added_size_x < min_added_size_x) ||
139- (min_added_size_y > 0 && added_size_y < min_added_size_y) ||
140- (min_added_size_x == added_size_x && wasted_size_x < min_wasted_size_x) ||
141- (min_added_size_y == added_size_y && wasted_size_y < min_wasted_size_y)) {
142- min_added_size_x = added_size_x;
143- min_added_size_y = added_size_y;
144- min_wasted_size_x = wasted_size_x;
145- min_wasted_size_y = wasted_size_y;
146- best_storage_id = free_id;
147- new_mem.x_ = new_width;
148- new_mem.y_ = new_height;
149- }
150- }
151-
152- if (min_added_size_x == 0 && min_added_size_y == 0 ) {
153- // use existing block
154- free_list_.erase (best_storage_id);
155- blocks_[best_storage_id].token_ ->ref_counter += prototype->ref_counter ;
156- return blocks_[best_storage_id].token_ ;
157- } else if (min_added_size_x <= shape.width || min_added_size_y <= shape.height ) {
158- // Reset the reference counter of the now live token
159- free_list_.erase (best_storage_id);
160- new_mem.token_ = prototype;
161- new_mem.token_ ->ref_counter += 1 ;
162- new_mem.token_ ->storage_id = best_storage_id;
163- blocks_[best_storage_id] = new_mem;
164- return new_mem.token_ ;
165- }
166- return nullptr ;
167- }
168-
169- StorageToken* TokenAllocator2D::Alloc (StorageToken* prototype, int64_t storage_id) {
170- auto shape = GetSize2D (prototype);
171- MemBlock block;
172- block.x_ = shape.width ;
173- block.y_ = shape.height ;
174- prototype->storage_id = storage_id;
175- block.token_ = prototype;
176- blocks_[prototype->storage_id ] = block;
177- return prototype;
178- }
179-
180- void TokenAllocator2D::CheckForRelease (StorageToken* tok) {
181- ICHECK_GE (tok->storage_id , 0 );
182- ICHECK_GE (tok->ref_counter , 0 );
183- if (tok->ref_counter == 0 ) {
184- free_list_.insert (tok->storage_id );
185- }
186- }
187-
188- runtime::Texture2DShape<int64_t > TokenAllocator2D::GetSize2D (StorageToken* prototype) {
122+ size_t TokenAllocatorMixed::GetSize2D (StorageToken* prototype) {
189123 TensorType ttype = prototype->ttype ;
190124 ICHECK (ttype.defined ());
191- size_t axis = runtime::DefaultTextureLayoutSeparator (ttype->shape .size (),
192- prototype->virtual_device ->memory_scope );
193125 struct Shape {
194126 const Array<PrimExpr>& shape;
195127 int64_t operator [](size_t i) const { return *tir::as_const_int (shape[i]); }
128+ int size () { return this ->shape .size (); }
196129 };
197- return runtime::ApplyTexture2DFlattening<int64_t >(Shape{ttype->shape }, ttype->shape .size (), axis);
130+ auto shape = Shape{ttype->shape };
131+ int image_row_align =
132+ prototype->virtual_device ->target ->GetAttr <Integer>(" image_base_address_alignment"
133+ .value_or (Integer (64 ))
134+ ->value ;
135+ return runtime::GetTextureMemorySize<Shape>(shape, ttype->dtype .bits (), ttype->dtype .lanes (),
136+ prototype->virtual_device ->memory_scope ,
137+ image_row_align);
198138}
199139
200140} // namespace relay
0 commit comments