srkreddy1238 commented on code in PR #15058:
URL: https://github.com/apache/tvm/pull/15058#discussion_r1349641612
##########
src/relay/backend/token_allocator.cc:
##########
@@ -31,22 +31,39 @@
namespace tvm {
namespace relay {
+constexpr auto Is2DStorage = runtime::IsTextureStorage;
-size_t TokenAllocator1D::GetMemorySize(StorageToken* prototype) {
+/*
+ * Mixed mode memory allocator
+ */
+size_t TokenAllocatorMixed::GetMemorySize(StorageToken* prototype) {
TensorType ttype = prototype->ttype;
ICHECK(ttype.defined());
size_t size = 1;
- for (IndexExpr dim : ttype->shape) {
- const int64_t* pval = tir::as_const_int(dim);
- ICHECK(pval != nullptr) << "Cannot allocate memory symbolic tensor shape "
<< ttype->shape;
- ICHECK_GE(*pval, 0) << "Cannot allocate memory for tensor with negative
shape" << *pval;
- size *= static_cast<size_t>(pval[0]);
+ if (relay::Is2DStorage(prototype->virtual_device->memory_scope)) {
+ size = GetSize2D(prototype);
+ } else {
+ for (IndexExpr dim : ttype->shape) {
+ const int64_t* pval = tir::as_const_int(dim);
+ ICHECK(pval != nullptr) << "Cannot allocate memory symbolic tensor shape
" << ttype->shape;
+ ICHECK_GE(*pval, 0) << "Cannot allocate memory for tensor with negative
shape" << *pval;
+ 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);
return size;
}
-StorageToken* TokenAllocator1D::Request(StorageToken* prototype) {
+bool IsAdreno(StorageToken* tok) {
Review Comment:
Improved the API with first option now. I am working on introducing better
type for MemoryScope rather just String. Will revisit this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
