graph_runtime.cc 17.6 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
/*!
 * \file graph_runtime.cc
 */
23
#include <tvm/runtime/device_api.h>
Zhi committed
24
#include <tvm/runtime/ndarray.h>
25 26
#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/registry.h>
Zhi committed
27 28
#include <tvm/runtime/serializer.h>

29 30
#include <algorithm>
#include <functional>
31
#include <memory>
Zhi committed
32
#include <numeric>
33
#include <string>
34
#include <unordered_set>
35
#include <utility>
36
#include <vector>
37

38 39
#include "graph_runtime.h"

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
void GraphRuntime::Run() {
  // setup the array and requirements.
  for (size_t i = 0; i < op_execs_.size(); ++i) {
56
    if (op_execs_[i]) op_execs_[i]();
57
  }
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76
}
/*!
 * \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) {
  std::istringstream is(graph_json);
  dmlc::JSONReader reader(&is);
  this->Load(&reader);
  module_ = module;
  ctxs_ = ctxs;
  this->SetupStorage();
  this->SetupOpExecs();
77 78 79 80 81
  for (size_t i = 0; i < input_nodes_.size(); i++) {
    const uint32_t nid = input_nodes_[i];
    std::string& name = nodes_[nid].name;
    input_map_[name] = i;
  }
82 83 84 85 86 87 88
}
/*!
 * \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) {
89 90 91
  auto it = input_map_.find(name);
  if (it != input_map_.end()) {
    return it->second;
92
  }
93 94 95 96 97 98 99 100 101 102 103 104 105 106
  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);
}
/*!
107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126
 * \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
127 128
  for (DLTensor* t : input_dltensors_[eid]) {
    t->data = data_ref->data;
129 130 131
  }
}
/*!
132 133 134 135 136 137 138 139 140 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
 * \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]);
169

170 171 172 173 174
  // 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]);
175
  }
176

177 178
  data_entry_[eid].CopyTo(data_out);
}
179

180 181 182 183 184 185 186 187
/*!
 * \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);
}
188 189 190 191 192 193 194 195 196 197 198 199 200 201

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
202
  strm->Read(&sz);
203 204 205 206
  size_t size = static_cast<size_t>(sz);
  CHECK(size == names.size())
      << "Invalid parameters file format";
  for (size_t i = 0; i < size; ++i) {
207 208
    int in_idx = GetInputIndex(names[i]);
    CHECK_GE(in_idx, 0) << "Found param for non-existent input: " << names[i];
209 210
    uint32_t eid = this->entry_id(input_nodes_[in_idx], 0);
    CHECK_LT(eid, data_entry_.size());
211 212 213 214 215

    // 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);
216 217 218
  }
}

219
void GraphRuntime::ShareParams(const GraphRuntime& other, dmlc::Stream* strm) {
220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240
    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);
241 242
    const DLTensor* tmp = data_entry_[eid].operator->();
    data_alignment_[eid] = details::GetDataAlignment(*tmp);
243 244 245 246
  }
  this->SetupOpExecs();
}

247 248
void GraphRuntime::SetupStorage() {
  // Grab saved optimization plan from graph.
249
  std::vector<DLDataType> vtype;
250
  for (const std::string& s_type : attrs_.dltype) {
251
    vtype.push_back(tvm::runtime::String2DLDataType(s_type));
252
  }
253

Zhi committed
254 255
  // Size and device type of each storage pool entry.
  std::vector<PoolEntry> pool_entry;
256 257 258
  // Find the maximum space size.
  for (size_t i = 0; i < attrs_.shape.size(); ++i) {
    int storage_id = attrs_.storage_id[i];
Zhi committed
259 260 261 262 263
    // 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];
    }
264 265 266 267 268 269 270
    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;
271 272
    CHECK(bits % 8U ==  0U || bits ==1U);
    size_t bytes = ((bits + 7U) / 8U) * size;
273

Zhi committed
274 275 276 277 278 279 280
    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";
281
    }
Zhi committed
282 283
    pool_entry[sid].size = std::max(pool_entry[sid].size, bytes);
    pool_entry[sid].device_type = device_type;
284
  }
Zhi committed
285

286
  // Allocate the space.
Zhi committed
287
  for (const auto& pit : pool_entry) {
288
    std::vector<int64_t> shape;
Zhi committed
289 290 291 292 293 294 295 296 297
    // 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(
298
        NDArray::Empty(shape, DLDataType{kDLFloat, 32, 1}, ctx));
299
  }
Zhi committed
300 301 302 303 304

  // 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());
305
  data_alignment_.resize(num_node_entries());
306 307
  for (size_t i = 0; i < data_entry_.size(); ++i) {
    int storage_id = attrs_.storage_id[i];
308
    CHECK_LT(static_cast<size_t>(storage_id), storage_pool_.size());
Zhi committed
309 310
    data_entry_[i] =
        storage_pool_[storage_id].CreateView(attrs_.shape[i], vtype[i]);
311 312
    const DLTensor* tmp = data_entry_[i].operator->();
    data_alignment_[i] = details::GetDataAlignment(*tmp);
313 314 315 316
  }
}

void GraphRuntime::SetupOpExecs() {
317
  op_execs_.resize(this->GetNumOfNodes());
318 319 320 321 322 323 324
  input_dltensors_.resize(num_node_entries());
  std::unordered_set<uint32_t> input_node_eids;
  for (size_t i = 0; i < input_nodes_.size(); i++) {
    uint32_t nid = input_nodes_[i];
    input_node_eids.insert(entry_id(nid, 0));
  }

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

340 341
    std::shared_ptr<OpArgs> op_args = nullptr;
    std::tie(op_execs_[nid], op_args) =
342
        CreateTVMOp(inode.param, args, inode.inputs.size());
343 344 345 346 347 348 349

    for (size_t i = 0; i < inode.inputs.size(); i++) {
      uint32_t eid = this->entry_id(inode.inputs[i]);
      // check if op input is model input
      if (input_node_eids.count(eid) > 0) {
        input_dltensors_[eid].push_back(
            static_cast<DLTensor*>(op_args->arg_values[i].v_handle));
350 351
      }
    }
352 353 354
  }
}

355
std::pair<std::function<void()>, std::shared_ptr<GraphRuntime::OpArgs> > GraphRuntime::CreateTVMOp(
356 357 358
    const TVMOpParam& param,
    const std::vector<DLTensor>& args,
    size_t num_inputs) {
359
  std::shared_ptr<GraphRuntime::OpArgs> arg_ptr = std::make_shared<GraphRuntime::OpArgs>();
360
  // setup address.
361
  arg_ptr->args = args;
362 363 364 365 366
  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;
367
    DLTensor* t = &arg_ptr->args[i];
368 369
    v.v_handle = t;
    arg_ptr->arg_values.push_back(v);
370
    arg_ptr->arg_tcodes.push_back(kTVMDLTensorHandle);
371 372 373 374 375 376 377
    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
378

379
  if (param.func_name == "__nop") {
380
    return {[](){}, arg_ptr};
Zhi committed
381 382 383 384 385 386 387 388
  } 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));
    };
389
    return {fexec, arg_ptr};
390
  }
Zhi committed
391 392 393

  // Get compiled function from the module that contains both host and device
  // code.
394
  tvm::runtime::PackedFunc pf = module_.GetFunction(param.func_name, true);
395
  CHECK(pf != nullptr) << "no such function in module: " << param.func_name;
Zhi committed
396 397

  auto fexec = [arg_ptr, pf]() {
398 399 400 401 402 403
    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);
  };
404
  return {fexec, arg_ptr};
405 406 407 408
}

PackedFunc GraphRuntime::GetFunction(
    const std::string& name,
409
    const ObjectPtr<Object>& sptr_to_self) {
Zhi committed
410
  // Return member functions during query.
411 412
  if (name == "set_input") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
413
        if (args[0].type_code() == kTVMStr) {
414 415
          int in_idx = this->GetInputIndex(args[0]);
          if (in_idx >= 0) this->SetInput(in_idx, args[1]);
416 417 418 419
        } else {
          this->SetInput(args[0], args[1]);
        }
      });
420 421
  } else if (name == "set_input_zero_copy") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
422
      if (args[0].type_code() == kTVMStr) {
423 424 425 426 427 428
        int in_idx = this->GetInputIndex(args[0]);
        if (in_idx >= 0) this->SetInputZeroCopy(in_idx, args[1]);
      } else {
        this->SetInputZeroCopy(args[0], args[1]);
      }
    });
429 430
  } else if (name == "get_output") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
431 432 433 434 435 436
      if (args.num_args == 2) {
        this->CopyOutputTo(args[0], args[1]);
      } else {
        *rv = this->GetOutput(args[0]);
      }
    });
437 438
  } else if (name == "get_input") {
    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
439
        int in_idx = 0;
440
        if (args[0].type_code() == kTVMStr) {
441
          in_idx = this->GetInputIndex(args[0]);
442
        } else {
443
          in_idx = args[0];
444
        }
445 446 447 448 449 450
        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();
451
      });
452 453 454 455 456 457 458 459
  } 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());
      });
460 461 462 463 464 465 466 467
  } 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);
      });
468 469 470 471 472
  } else {
    return PackedFunc();
  }
}

Zhi committed
473 474 475
Module GraphRuntimeCreate(const std::string& sym_json,
                          const tvm::runtime::Module& m,
                          const std::vector<TVMContext>& ctxs) {
476
  auto exec = make_object<GraphRuntime>();
Zhi committed
477
  exec->Init(sym_json, m, ctxs);
478 479 480
  return Module(exec);
}

Zhi committed
481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499
// 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.
500
TVM_REGISTER_GLOBAL("tvm.graph_runtime.create")
Zhi committed
501 502 503 504 505 506 507
  .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);
508 509 510
  });
}  // namespace runtime
}  // namespace tvm