module.cc 6.31 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
/*!
 *  Copyright (c) 2017 by Contributors
 * \file module.cc
23
 * \brief TVM module system
24 25 26 27 28
 */
#include <tvm/runtime/module.h>
#include <tvm/runtime/registry.h>
#include <tvm/runtime/packed_func.h>
#include <unordered_set>
29
#include <cstring>
nhynes committed
30
#ifndef _LIBCPP_SGX_CONFIG
31
#include "file_util.h"
nhynes committed
32
#endif
33 34 35 36 37

namespace tvm {
namespace runtime {

void Module::Import(Module other) {
38 39 40 41
  // specially handle rpc
  if (!std::strcmp((*this)->type_key(), "rpc")) {
    static const PackedFunc* fimport_ = nullptr;
    if (fimport_ == nullptr) {
42
      fimport_ = runtime::Registry::Get("rpc._ImportRemoteModule");
43 44 45 46 47
      CHECK(fimport_ != nullptr);
    }
    (*fimport_)(*this, other);
    return;
  }
48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67
  // cyclic detection.
  std::unordered_set<const ModuleNode*> visited{other.node_.get()};
  std::vector<const ModuleNode*> stack{other.node_.get()};
  while (!stack.empty()) {
    const ModuleNode* n = stack.back();
    stack.pop_back();
    for (const Module& m : n->imports_) {
      const ModuleNode* next = m.node_.get();
      if (visited.count(next)) continue;
      visited.insert(next);
      stack.push_back(next);
    }
  }
  CHECK(!visited.count(node_.get()))
      << "Cyclic dependency detected during import";
  node_->imports_.emplace_back(std::move(other));
}

Module Module::LoadFromFile(const std::string& file_name,
                            const std::string& format) {
nhynes committed
68
#ifndef _LIBCPP_SGX_CONFIG
69 70 71 72 73 74
  std::string fmt = GetFileFormat(file_name, format);
  CHECK(fmt.length() != 0)
      << "Cannot deduce format of file " << file_name;
  if (fmt == "dll" || fmt == "dylib" || fmt == "dso") {
    fmt = "so";
  }
75
  std::string load_f_name = "module.loadfile_" + fmt;
76 77 78 79 80 81
  const PackedFunc* f = Registry::Get(load_f_name);
  CHECK(f != nullptr)
      << "Loader of " << format << "("
      << load_f_name << ") is not presented.";
  Module m = (*f)(file_name, format);
  return m;
nhynes committed
82 83 84
#else
  LOG(FATAL) << "SGX does not support LoadFromFile";
#endif
85 86
}

87 88 89 90 91 92 93 94 95 96 97 98 99 100
void ModuleNode::SaveToFile(const std::string& file_name,
                            const std::string& format) {
  LOG(FATAL) << "Module[" << type_key() << "] does not support SaveToFile";
}

void ModuleNode::SaveToBinary(dmlc::Stream* stream) {
  LOG(FATAL) << "Module[" << type_key() << "] does not support SaveToBinary";
}

std::string ModuleNode::GetSource(const std::string& format) {
  LOG(FATAL) << "Module[" << type_key() << "] does not support GetSource";
  return "";
}

101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121
const PackedFunc* ModuleNode::GetFuncFromEnv(const std::string& name) {
  auto it = import_cache_.find(name);
  if (it != import_cache_.end()) return it->second.get();
  PackedFunc pf;
  for (Module& m : this->imports_) {
    pf = m.GetFunction(name, false);
    if (pf != nullptr) break;
  }
  if (pf == nullptr) {
    const PackedFunc* f = Registry::Get(name);
    CHECK(f != nullptr)
        << "Cannot find function " << name
        << " in the imported modules or global registry";
    return f;
  } else {
    std::unique_ptr<PackedFunc> f(new PackedFunc(pf));
    import_cache_[name] = std::move(f);
    return import_cache_.at(name).get();
  }
}

122
bool RuntimeEnabled(const std::string& target) {
123
  std::string f_name;
124 125 126
  if (target == "cpu") {
    return true;
  } else if (target == "cuda" || target == "gpu") {
127
    f_name = "device_api.gpu";
128
  } else if (target == "cl" || target == "opencl" || target == "sdaccel") {
129
    f_name = "device_api.opencl";
130 131
  } else if (target == "gl" || target == "opengl") {
    f_name = "device_api.opengl";
132 133
  } else if (target == "mtl" || target == "metal") {
    f_name = "device_api.metal";
134 135
  } else if (target == "vulkan") {
    f_name = "device_api.vulkan";
136 137
  } else if (target == "stackvm") {
    f_name = "codegen.build_stackvm";
138 139
  } else if (target == "rpc") {
    f_name = "device_api.rpc";
140 141
  } else if (target == "vpi" || target == "verilog") {
    f_name = "device_api.vpi";
142 143
  } else if (target == "micro_dev") {
    f_name = "device_api.micro_dev";
144
  } else if (target.length() >= 5 && target.substr(0, 5) == "nvptx") {
145
    f_name = "device_api.gpu";
146
  } else if (target.length() >= 4 && target.substr(0, 4) == "rocm") {
147
    f_name = "device_api.rocm";
148 149 150 151
  } else if (target.length() >= 4 && target.substr(0, 4) == "llvm") {
    const PackedFunc* pf = runtime::Registry::Get("codegen.llvm_target_enabled");
    if (pf == nullptr) return false;
    return (*pf)(target);
152 153 154
  } else {
    LOG(FATAL) << "Unknown optional runtime " << target;
  }
155
  return runtime::Registry::Get(f_name) != nullptr;
156 157
}

158
TVM_REGISTER_GLOBAL("module._Enabled")
159 160 161 162
.set_body([](TVMArgs args, TVMRetValue *ret) {
    *ret = RuntimeEnabled(args[0]);
    });

163
TVM_REGISTER_GLOBAL("module._GetSource")
164 165 166 167
.set_body([](TVMArgs args, TVMRetValue *ret) {
    *ret = args[0].operator Module()->GetSource(args[1]);
    });

168
TVM_REGISTER_GLOBAL("module._ImportsSize")
169 170 171 172 173
.set_body([](TVMArgs args, TVMRetValue *ret) {
    *ret = static_cast<int64_t>(
        args[0].operator Module()->imports().size());
    });

174
TVM_REGISTER_GLOBAL("module._GetImport")
175 176 177 178 179
.set_body([](TVMArgs args, TVMRetValue *ret) {
    *ret = args[0].operator Module()->
        imports().at(args[1].operator int());
    });

180
TVM_REGISTER_GLOBAL("module._GetTypeKey")
181
.set_body([](TVMArgs args, TVMRetValue *ret) {
182
    *ret = std::string(args[0].operator Module()->type_key());
183 184
    });

185
TVM_REGISTER_GLOBAL("module._LoadFromFile")
186 187 188 189
.set_body([](TVMArgs args, TVMRetValue *ret) {
    *ret = Module::LoadFromFile(args[0], args[1]);
    });

190
TVM_REGISTER_GLOBAL("module._SaveToFile")
191 192 193 194 195 196
.set_body([](TVMArgs args, TVMRetValue *ret) {
    args[0].operator Module()->
        SaveToFile(args[1], args[2]);
    });
}  // namespace runtime
}  // namespace tvm