graph_runtime.cc 18.3 KB
Newer Older
1 2 3 4 5 6 7 8
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
9
 *
10
 *   http://www.apache.org/licenses/LICENSE-2.0
11
 *
12 13 14 15 16 17 18 19
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

20 21 22 23
/*!
 *  Copyright (c) 2017 by Contributors
 * \file graph_runtime.cc
 */
Zhi committed
24 25
#include "graph_runtime.h"

26
#include <tvm/runtime/device_api.h>
Zhi committed
27
#include <tvm/runtime/ndarray.h>
28 29
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>
Zhi committed
30 31
#include <tvm/runtime/serializer.h>

32 33
#include <algorithm>
#include <functional>
Zhi committed
34 35
#include <numeric>
#include <vector>
36
#include <string>
37 38
#include <memory>
#include <utility>
39 40 41

namespace tvm {
namespace runtime {
42 43 44 45 46 47 48
namespace details {
inline size_t GetDataAlignment(const DLTensor& arr) {
  size_t align = (arr.dtype.bits / 8) * arr.dtype.lanes;
  if (align < kAllocAlignment) return kAllocAlignment;
  return align;
}
}  // namespace details
49 50

/*!
51
 * \brief Run all the operations one by one.
52
 */
53 54 55 56
void GraphRuntime::Run() {
  // setup the array and requirements.
  for (size_t i = 0; i < op_execs_.size(); ++i) {
    if (op_execs_[i]) op_execs_[i]();
57
  }
58 59 60 61 62 63 64 65 66 67 68 69
}
/*!
 * \brief Initialize the graph executor with graph and context.
 * \param graph_json The execution graph.
 * \param module The module containing the compiled functions for the host
 * processor.
 * \param ctxs The context of the host and devices where graph nodes will be
 * executed on.
 */
void GraphRuntime::Init(const std::string& graph_json,
                        tvm::runtime::Module module,
                        const std::vector<TVMContext>& ctxs) {
nhynes committed
70
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
71
  std::istringstream is(graph_json);
nhynes committed
72
#else
73
  std::string is = graph_json;
nhynes committed
74
#endif
75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91
  dmlc::JSONReader reader(&is);
  this->Load(&reader);
  module_ = module;
  ctxs_ = ctxs;
  this->SetupStorage();
  this->SetupOpExecs();
}
/*!
 * \brief Get the input index given the name of input.
 * \param name The name of the input.
 * \return The index of input.
 */
int GraphRuntime::GetInputIndex(const std::string& name) {
  for (size_t i = 0; i< input_nodes_.size(); ++i) {
    uint32_t nid = input_nodes_[i];
    if (nodes_[nid].name == name) {
      return static_cast<int>(i);
92 93
    }
  }
94 95 96 97 98 99 100 101 102 103 104 105 106 107
  LOG(WARNING) << "Warning: cannot find \"" << name << "\" among input";
  return -1;
}
/*!
 * \brief set index-th input to the graph.
 * \param index The input index.
 * \param data_in The input data.
 */
void GraphRuntime::SetInput(int index, DLTensor* data_in) {
  CHECK_LT(static_cast<size_t>(index), input_nodes_.size());
  uint32_t eid = this->entry_id(input_nodes_[index], 0);
  data_entry_[eid].CopyFrom(data_in);
}
/*!
108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
 * \brief set index-th input to the graph without copying the data.
 * \param index The input index.
 * \param data_ref The input data that is referred.
 */
void GraphRuntime::SetInputZeroCopy(int index, DLTensor* data_ref) {
  CHECK_LT(static_cast<size_t>(index), input_nodes_.size());
  uint32_t eid = this->entry_id(input_nodes_[index], 0);
  const DLTensor* old_t = data_entry_[eid].operator->();

  // check the consistency of input
  CHECK_EQ(data_alignment_[eid], details::GetDataAlignment(*data_ref));
  CHECK_EQ(reinterpret_cast<size_t>(data_ref->data) % kAllocAlignment, 0);
  CHECK_EQ(old_t->ndim, static_cast<size_t>(data_ref->ndim));
  CHECK_EQ(old_t->ctx.device_type, data_ref->ctx.device_type);
  CHECK_EQ(old_t->ctx.device_id, data_ref->ctx.device_id);
  for (auto i = 0; i < data_ref->ndim; ++i) {
    CHECK_EQ(old_t->shape[i], data_ref->shape[i]);
  }

  // Update the data pointer for each argument of each op
  for (auto& op_arg : op_args_) {
    if (op_arg) {
      const auto it = op_arg->input_entry_ids.find(eid);
      if (it != op_arg->input_entry_ids.end()) {
        for (const auto i : it->second) {
          DLTensor* t = static_cast<DLTensor*>(op_arg->arg_values[i].v_handle);
          t->data = data_ref->data;
        }
      }
    }
  }
}
/*!
141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177
 * \brief Get the number of outputs
 *
 * \return The number of outputs from graph.
 */
int GraphRuntime::NumOutputs() const {
  return outputs_.size();
}
/*!
 * \brief Return NDArray for given input index.
 * \param index The input index.
 *
 * \return NDArray corresponding to given input node index.
 */
NDArray GraphRuntime::GetInput(int index) const {
  CHECK_LT(static_cast<size_t>(index), input_nodes_.size());
  uint32_t eid = this->entry_id(input_nodes_[index], 0);
  return data_entry_[eid];
}
/*!
 * \brief Return NDArray for given output index.
 * \param index The output index.
 *
 * \return NDArray corresponding to given output node index.
 */
NDArray GraphRuntime::GetOutput(int index) const {
  CHECK_LT(static_cast<size_t>(index), outputs_.size());
  uint32_t eid = this->entry_id(outputs_[index]);
  return data_entry_[eid];
}
/*!
 * \brief Copy index-th output to data_out.
 * \param index The output index.
 * \param data_out the output data.
 */
void GraphRuntime::CopyOutputTo(int index, DLTensor* data_out) {
  CHECK_LT(static_cast<size_t>(index), outputs_.size());
  uint32_t eid = this->entry_id(outputs_[index]);
178

179 180 181 182 183
  // Check the shapes to avoid receiving in different dimension but same size.
  const NDArray& data = data_entry_[eid];
  CHECK_EQ(data->ndim, data_out->ndim);
  for (int32_t j = 0; j < data->ndim; ++j) {
    CHECK_EQ(data->shape[j], data_out->shape[j]);
184
  }
185

186 187
  data_entry_[eid].CopyTo(data_out);
}
188

189 190 191 192 193 194 195 196
/*!
 * \brief Load parameters from parameter blob.
 * \param param_blob A binary blob of parameter.
 */
void GraphRuntime::LoadParams(const std::string& param_blob) {
  dmlc::MemoryStringStream strm(const_cast<std::string*>(&param_blob));
  this->LoadParams(&strm);
}
197 198 199 200 201 202 203 204 205 206 207 208 209 210

void GraphRuntime::LoadParams(dmlc::Stream* strm) {
  uint64_t header, reserved;
  CHECK(strm->Read(&header))
      << "Invalid parameters file format";
  CHECK(header == kTVMNDArrayListMagic)
      << "Invalid parameters file format";
  CHECK(strm->Read(&reserved))
      << "Invalid parameters file format";

  std::vector<std::string> names;
  CHECK(strm->Read(&names))
      << "Invalid parameters file format";
  uint64_t sz;
tqchen committed
211
  strm->Read(&sz);
212 213 214 215
  size_t size = static_cast<size_t>(sz);
  CHECK(size == names.size())
      << "Invalid parameters file format";
  for (size_t i = 0; i < size; ++i) {
216 217
    int in_idx = GetInputIndex(names[i]);
    CHECK_GE(in_idx, 0) << "Found param for non-existent input: " << names[i];
218 219
    uint32_t eid = this->entry_id(input_nodes_[in_idx], 0);
    CHECK_LT(eid, data_entry_.size());
220 221 222 223 224

    // The data_entry is allocated on device, NDArray.load always load the array into CPU.
    NDArray temp;
    temp.Load(strm);
    data_entry_[eid].CopyFrom(temp);
225 226 227
  }
}

228
void GraphRuntime::ShareParams(const GraphRuntime& other, dmlc::Stream* strm) {
229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249
    uint64_t header, reserved;
    CHECK(strm->Read(&header))
      << "Invalid parameters file format";
    CHECK(header == kTVMNDArrayListMagic)
      << "Invalid parameters file format";
    CHECK(strm->Read(&reserved))
      << "Invalid parameters file format";
  std::vector<std::string> names;
  CHECK(strm->Read(&names)) << "Invalid parameters file format";
  uint64_t sz;
  strm->Read(&sz);
  size_t size = static_cast<size_t>(sz);
  CHECK(size == names.size()) << "Invalid parameters file format";
  for (size_t i = 0; i < size; ++i) {
    int in_idx = GetInputIndex(names[i]);
    CHECK_GE(in_idx, 0) << "Found param for non-existent input: " << names[i];
    uint32_t eid = this->entry_id(input_nodes_[in_idx], 0);
    CHECK_LT(eid, data_entry_.size());
    CHECK_EQ(data_entry_[eid].use_count(), 1);
    data_entry_[eid] = other.GetInput(GetInputIndex(names[i]));
    CHECK_GT(data_entry_[eid].use_count(), 1);
250 251
    const DLTensor* tmp = data_entry_[eid].operator->();
    data_alignment_[eid] = details::GetDataAlignment(*tmp);
252 253 254 255
  }
  this->SetupOpExecs();
}

256 257 258 259 260 261
void GraphRuntime::SetupStorage() {
  // Grab saved optimization plan from graph.
  std::vector<TVMType> vtype;
  for (const std::string& s_type : attrs_.dltype) {
    vtype.push_back(tvm::runtime::String2TVMType(s_type));
  }
Zhi committed
262 263 264

  // Size and device type of each storage pool entry.
  std::vector<PoolEntry> pool_entry;
265 266 267
  // Find the maximum space size.
  for (size_t i = 0; i < attrs_.shape.size(); ++i) {
    int storage_id = attrs_.storage_id[i];
Zhi committed
268 269 270 271 272
    // Use the fallback device if no device index is available.
    int device_type = static_cast<int>(ctxs_[0].device_type);
    if (!attrs_.device_index.empty()) {
      device_type = attrs_.device_index[i];
    }
273 274 275 276 277 278 279
    size_t size = 1;
    for (int64_t sz : attrs_.shape[i]) {
      size *= static_cast<size_t>(sz);
    }
    CHECK_GE(storage_id, 0) << "Do not support runtime shape op";
    DLDataType t = vtype[i];
    size_t bits = t.bits * t.lanes;
280 281
    CHECK(bits % 8U ==  0U || bits ==1U);
    size_t bytes = ((bits + 7U) / 8U) * size;
282

Zhi committed
283 284 285 286 287 288 289
    uint32_t sid = static_cast<uint32_t>(storage_id);
    if (sid >= pool_entry.size()) {
      pool_entry.resize(sid + 1, {0, -1});
    } else {
      CHECK(pool_entry[sid].device_type == -1 ||
            pool_entry[sid].device_type == device_type)
          << "The same pool entry cannot be assigned to multiple devices";
290
    }
Zhi committed
291 292
    pool_entry[sid].size = std::max(pool_entry[sid].size, bytes);
    pool_entry[sid].device_type = device_type;
293
  }
Zhi committed
294

295
  // Allocate the space.
Zhi committed
296
  for (const auto& pit : pool_entry) {
297
    std::vector<int64_t> shape;
Zhi committed
298 299 300 301 302 303 304 305 306 307
    // This for loop is very fast since there are usually only a couple of
    // devices available on the same hardware.
    const auto& cit =
        std::find_if(ctxs_.begin(), ctxs_.end(), [&pit](const TVMContext& c) {
          return pit.device_type == static_cast<int>(c.device_type);
        });
    TVMContext ctx = cit == ctxs_.end() ? ctxs_[0] : *cit;
    shape.push_back(static_cast<int64_t>(pit.size + 3) / 4);
    storage_pool_.push_back(
        NDArray::Empty(shape, DLDataType{kDLFloat, 32, 1}, ctx));
308
  }
Zhi committed
309 310 311 312 313

  // Assign the pooled entries. A unified memory pool is used to simplifiy
  // memory assignment for each node entry. The allocated memory on each device
  // is mapped to this pool.
  data_entry_.resize(num_node_entries());
314
  data_alignment_.resize(num_node_entries());
315 316
  for (size_t i = 0; i < data_entry_.size(); ++i) {
    int storage_id = attrs_.storage_id[i];
317
    CHECK_LT(static_cast<size_t>(storage_id), storage_pool_.size());
Zhi committed
318 319
    data_entry_[i] =
        storage_pool_[storage_id].CreateView(attrs_.shape[i], vtype[i]);
320 321
    const DLTensor* tmp = data_entry_[i].operator->();
    data_alignment_[i] = details::GetDataAlignment(*tmp);
322 323 324 325
  }
}

void GraphRuntime::SetupOpExecs() {
326
  op_execs_.resize(this->GetNumOfNodes());
327
  op_args_.resize(this->GetNumOfNodes());
328
  // setup the array and requirements.
329
  for (uint32_t nid = 0; nid < this->GetNumOfNodes(); ++nid) {
330 331 332
    const auto& inode = nodes_[nid];
    if (inode.op_type == "null") continue;
    std::vector<DLTensor> args;
333
    std::vector<uint32_t> input_entry_ids;
334
    for (const auto& e : inode.inputs) {
335 336 337
      uint32_t eid = this->entry_id(e);
      args.push_back(*(data_entry_[eid].operator->()));
      input_entry_ids.push_back(eid);
338 339 340
    }
    for (uint32_t index = 0; index < inode.param.num_outputs; ++index) {
      uint32_t eid = this->entry_id(nid, index);
341
      args.push_back(*(data_entry_[eid].operator->()));
342
    }
Zhi committed
343 344
    CHECK(inode.op_type == "tvm_op") << "Can only take tvm_op as op";

345 346 347 348 349 350 351 352 353 354 355 356
    std::tie(op_execs_[nid], op_args_[nid]) =
        CreateTVMOp(inode.param, args, inode.inputs.size());
    auto& entry_to_input_pos = op_args_[nid]->input_entry_ids;
    for (uint32_t i = 0; i < input_entry_ids.size(); ++i) {
      const auto eid = input_entry_ids[i];
      auto it = entry_to_input_pos.find(eid);
      if (it == entry_to_input_pos.end()) {
        entry_to_input_pos.emplace(eid, std::vector<uint32_t>{i});
      } else {
        it->second.push_back(i);
      }
    }
357 358 359
  }
}

360
std::pair<std::function<void()>, std::shared_ptr<GraphRuntime::OpArgs> > GraphRuntime::CreateTVMOp(
361 362 363
    const TVMOpParam& param,
    const std::vector<DLTensor>& args,
    size_t num_inputs) {
364
  std::shared_ptr<GraphRuntime::OpArgs> arg_ptr = std::make_shared<GraphRuntime::OpArgs>();
365
  // setup address.
366
  arg_ptr->args = args;
367 368 369 370 371
  if (param.flatten_data) {
    arg_ptr->shape_data.resize(arg_ptr->args.size());
  }
  for (size_t i = 0; i < arg_ptr->args.size(); ++i) {
    TVMValue v;
372
    DLTensor* t = &arg_ptr->args[i];
373 374 375 376 377 378 379 380 381 382
    v.v_handle = t;
    arg_ptr->arg_values.push_back(v);
    arg_ptr->arg_tcodes.push_back(kArrayHandle);
    if (param.flatten_data) {
      arg_ptr->shape_data[i] = std::accumulate(
          t->shape, t->shape + t->ndim, 1, std::multiplies<int64_t>());
      t->ndim = 1;
      t->shape = &(arg_ptr->shape_data[i]);
    }
  }
Zhi committed
383

384
  if (param.func_name == "__nop") {
385
    return {[](){}, arg_ptr};
Zhi committed
386 387 388 389 390 391 392 393
  } else if (param.func_name == "__copy") {
    // Perform cross device data copy.
    // Directly copy data from the input to the output.
    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));
    };
394
    return {fexec, arg_ptr};
395
  }
Zhi committed
396 397 398

  // Get compiled function from the module that contains both host and device
  // code.
399 400
  tvm::runtime::PackedFunc pf = module_.GetFunction(param.func_name, false);
  CHECK(pf != nullptr) << "no such function in module: " << param.func_name;
Zhi committed
401 402

  auto fexec = [arg_ptr, pf]() {
403 404 405 406 407 408
    TVMRetValue rv;
    TVMArgs targs(arg_ptr->arg_values.data(),
                  arg_ptr->arg_tcodes.data(),
                  static_cast<int>(arg_ptr->arg_values.size()));
    pf.CallPacked(targs, &rv);
  };
409
  return {fexec, arg_ptr};
410 411 412 413 414
}

PackedFunc GraphRuntime::GetFunction(
    const std::string& name,
    const std::shared_ptr<ModuleNode>& sptr_to_self) {
Zhi committed
415
  // Return member functions during query.
416 417 418
  if (name == "set_input") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
        if (args[0].type_code() == kStr) {
419 420
          int in_idx = this->GetInputIndex(args[0]);
          if (in_idx >= 0) this->SetInput(in_idx, args[1]);
421 422 423 424
        } else {
          this->SetInput(args[0], args[1]);
        }
      });
425 426 427 428 429 430 431 432 433
  } else if (name == "set_input_zero_copy") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
      if (args[0].type_code() == kStr) {
        int in_idx = this->GetInputIndex(args[0]);
        if (in_idx >= 0) this->SetInputZeroCopy(in_idx, args[1]);
      } else {
        this->SetInputZeroCopy(args[0], args[1]);
      }
    });
434 435
  } else if (name == "get_output") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
436 437 438 439 440 441
      if (args.num_args == 2) {
        this->CopyOutputTo(args[0], args[1]);
      } else {
        *rv = this->GetOutput(args[0]);
      }
    });
442 443
  } else if (name == "get_input") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
444
        int in_idx = 0;
445
        if (args[0].type_code() == kStr) {
446
          in_idx = this->GetInputIndex(args[0]);
447
        } else {
448
          in_idx = args[0];
449
        }
450 451 452 453 454 455
        CHECK_GE(in_idx, 0);
        *rv = this->GetInput(in_idx);
      });
  } else if (name == "get_num_outputs") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
        *rv = this->NumOutputs();
456
      });
457 458 459 460 461 462 463 464
  } else if (name == "run") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
        this->Run();
      });
  } else if (name == "load_params") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
        this->LoadParams(args[0].operator std::string());
      });
465 466 467 468 469 470 471 472
  } else if (name == "share_params") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
        const auto& module = args[0].operator Module();
        CHECK_EQ(module.operator->()->type_key(), "GraphRuntime");
        const auto& param_blob = args[1].operator std::string();
        dmlc::MemoryStringStream strm(const_cast<std::string*>(&param_blob));
        this->ShareParams(dynamic_cast<const GraphRuntime&>(*module.operator->()), &strm);
      });
473 474 475 476 477
  } else {
    return PackedFunc();
  }
}

Zhi committed
478 479 480
Module GraphRuntimeCreate(const std::string& sym_json,
                          const tvm::runtime::Module& m,
                          const std::vector<TVMContext>& ctxs) {
481
  std::shared_ptr<GraphRuntime> exec = std::make_shared<GraphRuntime>();
Zhi committed
482
  exec->Init(sym_json, m, ctxs);
483 484 485
  return Module(exec);
}

Zhi committed
486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504
// Get all context for the host and other runtime devices.
std::vector<TVMContext> GetAllContext(const TVMArgs& args) {
  // Reserve the first item as the fallback device.
  std::vector<TVMContext> ret;
  TVMContext ctx;
  for (int i = 2; i < args.num_args; i += 2) {
    int dev_type = args[i];
    ctx.device_type = static_cast<DLDeviceType>(dev_type);
    ctx.device_id = args[i + 1];
    ret.push_back(ctx);
  }
  return ret;
}

// 4-argument version is currently reserved to keep support of calling
// from tvm4j and javascript, since they don't have heterogeneous
// execution support yet. For heterogenenous execution, at least 5 arguments will
// be passed in. The third one is the number of devices.
// Eventually, we will only probably pass TVMContext for all the languages.
505
TVM_REGISTER_GLOBAL("tvm.graph_runtime.create")
Zhi committed
506 507 508 509 510 511 512
  .set_body([](TVMArgs args, TVMRetValue* rv) {
    CHECK_GE(args.num_args, 4)
        << "The expected number of arguments for graph_runtime.create is "
           "at least 4, but it has "
        << args.num_args;
    const auto& contexts = GetAllContext(args);
    *rv = GraphRuntimeCreate(args[0], args[1], contexts);
513 514 515
  });

TVM_REGISTER_GLOBAL("tvm.graph_runtime.remote_create")
Zhi committed
516 517 518 519 520
  .set_body([](TVMArgs args, TVMRetValue* rv) {
    CHECK_GE(args.num_args, 4) << "The expected number of arguments for "
                                  "graph_runtime.remote_create is "
                                  "at least 4, but it has "
                               << args.num_args;
521
    void* mhandle = args[1];
Zhi committed
522 523 524
    const auto& contexts = GetAllContext(args);
    *rv = GraphRuntimeCreate(
        args[0], *static_cast<tvm::runtime::Module*>(mhandle), contexts);
525 526 527
  });
}  // namespace runtime
}  // namespace tvm