Which Compiler pass embeds "device_copy" calls between modules

yogeesh · August 16, 2023, 10:31am

hi ,
I have a question regarding Runtime module execution wrt to memory managment. Basically if there is a scenario where :

Module1 (CPU)  -----> Module2 (GPU)  -------> Module3 (CPU)

are getting executed , lets say i have done :

with tvm.transform.PassContext(opt_level=2):
        graph, lib, params = relay.build(mod, target="rocm" , target_host = "llvm", params=None)

and then ran the runtime module using:

lib.export_library("libx.so", fcompile=False)
    lib = runtime.load_module("libx.so")
    rt_mod = graph_runtime.create(graph, lib ,tvm.device("rocm" , 0))
    # rt_mod = debug_runtime.create(graph, lib ,tvm.device("rocm" , 0))
    execute_rt_mod(rt_mod)

, what i need to understand is where exactly the device_copy calls are embedded between Module1 (CPU) and module2 (GPU) and similary between Module2(GPU) and module3(CPU). thanks

gangmul12 · August 17, 2023, 4:42pm

Hi,

I’m not that familiar with TVM, but I just looked for it.

I think “CopyDataFromTo” actually invokes data copy.

github.com

apache/tvm/blob/98320ab5bfff4d6778e87bdb763ffbe1b08c4c35/src/runtime/rocm/rocm_device_api.cc#L145-L147




void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t to_offset, size_t size,
                    Device dev_from, Device dev_to, DLDataType type_hint,

In AOT/Graph executor, "CopyDataFromTo"s are called if name of the function is “__copy”

github.com

apache/tvm/blob/98320ab5bfff4d6778e87bdb763ffbe1b08c4c35/src/runtime/graph_executor/graph_executor.cc#L585-L595


} else if (param.func_name == "__copy") {
  // Perform cross device data copy.
  // Directly copy data from the input to the output.
  // TODO(mbs): device_copy cleanup.
  auto fexec = [arg_ptr]() {
    DLTensor* from = static_cast<DLTensor*>(arg_ptr->arg_values[0].v_handle);
    DLTensor* to = static_cast<DLTensor*>(arg_ptr->arg_values[1].v_handle);
    TVM_CCALL(TVMArrayCopyFromTo(from, to, nullptr));
  };
  return {fexec, arg_ptr};
}

function “__copy” is generated by codegen if DeviceCopyProps are defined

github.com

apache/tvm/blob/98320ab5bfff4d6778e87bdb763ffbe1b08c4c35/src/relay/backend/graph_executor_codegen.cc#L413-L423


DeviceCopyProps device_copy_props = GetDeviceCopyProps(call_node);
CallLoweredProps call_lowered_props = GetCallLoweredProps(call_node);
if (device_copy_props.body.defined()) {
  // The graph executor expects to see a normal call to the undefined @__copy function.
  // The source and destination device annotations are no longer needed since they have
  // been captured in the StorageInfos for both input and output.
  // TODO(mbs): device_copy cleanup
  func_name = "__copy";
  for (const auto& n : VisitExpr(device_copy_props.body)) {
    inputs.push_back(n);
  }

It seems that DeviceCopyProps are automatically defined if it is device_copy op

github.com

apache/tvm/blob/98320ab5bfff4d6778e87bdb763ffbe1b08c4c35/src/relay/op/memory/device_copy.cc#L94-L97


// Get device copy props for original device copy op
DeviceCopyProps GetDeviceCopyProps(const CallNode* call_node) {
  if (call_node->op == DeviceCopyOp()) {
    ICHECK_EQ(call_node->args.size(), 1) << "device_copy expects one argument";

or it can be injected through relay PlanDevices transform pass

github.com

apache/tvm/blob/98320ab5bfff4d6778e87bdb763ffbe1b08c4c35/src/relay/transforms/device_planner.cc#L1070-L1108


/* =============== Phase 4 =============== */
/*!
 * \brief Inserts missing "device_copy" CallNodes, and ensures the device type of every
 * sub-expression in a module can be easily recovered by a later transformation using simple
 * lexical scoping rules (e.g. for memory planning).
 *
 * - Discard any existing "on_device" CallNodes since their job is done. Similarly, discard
 *   any existing "device_copy" CallNodes which are no-ops.
 *
 * - The result virtual device for a function is stored in the function's virtual_device_ field
 *   and the virtual devices of the function's parameters are stored in the parameter's
 *   virtual_device_ field.
 *
 * - Additional "device_copy" CallNodes are inserted wherever there's a transition between
 *   storage device types. Since the DeviceAnalyzer phase succeeded this can only happen
 *   where the original program explicitly allowed a transition using an "on_device" CallNode.
 *   That is, we do not not try to 'fix' a program with inconsistent devices.
 *
 * - Additional "on_device" CallNodes are inserted so that a later transform can discover
 *   the device for an arbitrary sub-expression by looking only for the lexically enclosing

This file has been truncated. show original

github.com

apache/tvm/blob/98320ab5bfff4d6778e87bdb763ffbe1b08c4c35/src/relay/transforms/device_planner.cc#L1520-L1528


// This function is declared in the public <tvm/relay/transform.h>.
tvm::transform::Pass PlanDevices(CompilationConfig config) {
  std::vector<Pass> passes;
  passes.emplace_back(Rewrite());
  passes.emplace_back(Check(config));
  passes.emplace_back(InferType());
  passes.emplace_back(PlanDevicesCore(config));
  return tvm::transform::Sequential(passes, "PlanDevices");
}

and PlanDevices transform pass seems that it is one of the default pass in Relay optimze

github.com

apache/tvm/blob/98320ab5bfff4d6778e87bdb763ffbe1b08c4c35/src/relay/backend/build_module.cc#L344-L350




// Always plan devices so the remaining passes don't need to distinguish homogeneous vs
// hetrogenous execution.
pass_seqs.push_back(transform::PlanDevices(config_));


// Fuse the operations if it is needed.
pass_seqs.push_back(transform::FuseOps());

Hope this reply helps

yogeesh · August 18, 2023, 6:08am

hi @gangmul12 thanks for your reply, its helpful. also after brief code review , i fugred that PlanDevices the lowest pass which does device planning and inserts device_copy calls in the graph. In detail , Check(config) pass is responsible for creating a new copy node and inserting it .