zhiics commented on a change in pull request #6337:
URL: https://github.com/apache/incubator-tvm/pull/6337#discussion_r477431597
##########
File path: python/tvm/relay/transform/memory_alloc.py
##########
@@ -19,26 +19,45 @@
A pass for manifesting explicit memory allocations.
"""
import numpy as np
+
+from tvm.ir.transform import PassContext, module_pass
+from tvm import nd, container
+from ..function import Function
from ..expr_functor import ExprVisitor, ExprMutator
from ..scope_builder import ScopeBuilder
-from . import transform
from .. import op
from ... import DataType, register_func
from .. import ty, expr
from ..backend import compile_engine
from ..op.memory import flatten_tuple_type, from_tuple_type, to_tuple_type
from ...import cpu
from ..op.memory import alloc_storage
+from ..analysis import context_analysis as _context_analysis
+from ..._ffi.runtime_ctypes import TVMContext
def alloc_tensor(storage, shape, dtype='float32', assert_shape=None):
offset = expr.const(0, dtype="int64")
return op.memory.alloc_tensor(storage, offset, shape, dtype, assert_shape)
+
def is_primitive(call):
return hasattr(call, 'op') and hasattr(call.op, 'attrs') and \
hasattr(call.op.attrs, 'Primitive') and
int(call.op.attrs.Primitive) == 1
+def is_device_copy(func):
+ """
+ Check if the current relay expression is shape_of call. We can simply check
+ the body of it if it is a function becase the shape_of op is opaque.
+ """
+ if isinstance(func, Function):
+ body = func.body
+ return isinstance(body, expr.Call) and body.op == op.get("device_copy")
+ if isinstance(func, expr.Call):
+ return body.op == op.get("device_copy")
Review comment:
yeah, it should be func.body. It could be a fused primitive op or a
device copy callnode. We have a device copy test for context analysis pass but
not for memory alloc pass because it is usually automatically added
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
