leandron commented on a change in pull request #6337:
URL: https://github.com/apache/incubator-tvm/pull/6337#discussion_r477165445
##########
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:
Is `body` here expected to have some value, or maybe should be replaced
with `func.body`? or even removed in case it is fully covered by the previous
`if` statement?
Also, I think it would be good to have a test case to cover this statement,
if there is not one already.
##########
File path: python/tvm/relay/analysis/analysis.py
##########
@@ -28,6 +28,21 @@
from .feature import Feature
+def context_analysis(mod, default_context):
+ """Analyze the device context information of each IR node in a Relay
+ program.
+
+ Parameters
+ ----------
+ expr : tvm.IRModule
+ The input module.
Review comment:
I guess here `expr` is actually `mod` ?
----------------------------------------------------------------
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]
