[tvm] branch main updated: [PROFILER, VM] Fix timer device type for reshape_tensor (#9518)

[areusch]

This is an automated email from the ASF dual-hosted git repository.

areusch pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new d66a40a  [PROFILER,VM] Fix timer device type for reshape_tensor (#9518)
d66a40a is described below

commit d66a40a613c0e7291f01e6cd41e090b0cf525e0c
Author: Tristan Konolige <tkonol...@octoml.ai>
AuthorDate: Thu Nov 18 09:07:44 2021 -0800

    [PROFILER,VM] Fix timer device type for reshape_tensor (#9518)
    
    * [PROFILER,VM] Fix timer device type for reshape_tensor
    
    The index of the host device was changed in the VM's `devices_` array,
    but we forgot to update this index in the profiler.
    
    * formatting
---
 src/runtime/vm/profiler/vm.cc                     |  2 +-
 tests/python/unittest/test_runtime_vm_profiler.py | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/runtime/vm/profiler/vm.cc b/src/runtime/vm/profiler/vm.cc
index e5afb0e..fe27595 100644
--- a/src/runtime/vm/profiler/vm.cc
+++ b/src/runtime/vm/profiler/vm.cc
@@ -107,7 +107,7 @@ void VirtualMachineDebug::OpStartHook(Instruction instr) {
       Device dst_dev = GetDevice(instr.device_copy.dst_device_index);
       prof_.operator*().StartCall("VM::DeviceCopy", dst_dev, {});
     } else if (instr.op == Opcode::ReshapeTensor) {
-      prof_.operator*().StartCall("VM::ReshapeTensor", devices_[1], {});
+      prof_.operator*().StartCall("VM::ReshapeTensor", 
devices_[exec_->host_device_index], {});
     } else if (instr.op == Opcode::AllocTensor) {
       auto shape = std::vector<int64_t>(instr.alloc_tensor.ndim);
 
diff --git a/tests/python/unittest/test_runtime_vm_profiler.py 
b/tests/python/unittest/test_runtime_vm_profiler.py
index 0499a3e..45bce02 100644
--- a/tests/python/unittest/test_runtime_vm_profiler.py
+++ b/tests/python/unittest/test_runtime_vm_profiler.py
@@ -38,6 +38,21 @@ def test_basic(dev, target):
     assert "softmax" in str(res)
 
 
+def test_vm_reshape_and_copy():
+    target = "llvm"
+    dev = tvm.gpu()
+    x_np = np.random.uniform(size=(8, 16)).astype("float32")
+    x = relay.var("x", shape=(8, 16), dtype="float32")
+    y = relay.reshape(x, [-1, 4, 8])
+    mod = tvm.IRModule()
+    mod["main"] = relay.Function([x], y)
+    with tvm.transform.PassContext(opt_level=3):
+        exec = relay.vm.compile(mod, "llvm")
+    assert "reshape_tensor" in exec.bytecode
+    vm = profiler_vm.VirtualMachineProfiler(exec, dev)
+    vm.profile(tvm.nd.array(x_np))
+
+
 if __name__ == "__main__":
     import sys
     import pytest