executorch/extension/module/ptn_module.h at main · pytorch/executorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

#pragma once

#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <string_view>
#include <unordered_set>
#include <variant>
#include <vector>

#include <executorch/runtime/core/data_loader.h>
#include <executorch/runtime/core/error.h>
#include <executorch/runtime/core/evalue.h>
#include <executorch/runtime/core/result.h>
#include <executorch/runtime/executor/method_meta.h>
#include <executorch/runtime/executor/program.h>

#ifndef ET_PTN_MODULE_NAMESPACE
#ifdef USE_ATEN_LIB
#define ET_PTN_MODULE_NAMESPACE native_module::aten
#else // !USE_ATEN_LIB
#define ET_PTN_MODULE_NAMESPACE native_module
#endif // USE_ATEN_LIB
#endif // ET_PTN_MODULE_NAMESPACE

namespace executorch::extension::ET_PTN_MODULE_NAMESPACE::internal {

using ET_RUNTIME_NAMESPACE::MethodMeta;
using ET_RUNTIME_NAMESPACE::Program;

class PtnModule {
 public:
  virtual ~PtnModule();

  virtual runtime::Result<size_t> num_methods() const = 0;
  virtual runtime::Result<std::unordered_set<std::string>> method_names()
      const = 0;
  virtual runtime::Result<MethodMeta> method_meta(
      const std::string& method_name) = 0;
  virtual runtime::Error load_method(const std::string& method_name) = 0;
  virtual bool unload_method(const std::string& method_name) = 0;
  virtual bool is_method_loaded(const std::string& method_name) const = 0;
  virtual runtime::Result<std::vector<runtime::EValue>> execute(
      const std::string& method_name,
      const std::vector<runtime::EValue>& inputs) = 0;
  virtual runtime::Error set_input(
      const std::string& method_name,
      const runtime::EValue& input,
      size_t index) = 0;
  virtual runtime::Error set_inputs(
      const std::string& method_name,
      const std::vector<runtime::EValue>& inputs) = 0;
  virtual runtime::Error set_output(
      const std::string& method_name,
      runtime::EValue output,
      size_t index) = 0;
  virtual runtime::Error set_outputs(
      const std::string& method_name,
      const std::vector<runtime::EValue>& outputs) = 0;
  virtual runtime::Result<std::vector<runtime::EValue>> get_outputs(
      const std::string& method_name) = 0;
  virtual runtime::Result<runtime::EValue> get_output(
      const std::string& method_name,
      size_t index) = 0;
};

struct PtnFileSource {
  enum class Mode : uint8_t {
    Read,
    Mmap,
  };

  std::string_view path;
  Mode mode;
};

using PtnSource =
    std::variant<std::reference_wrapper<runtime::DataLoader>, PtnFileSource>;

struct PtnHooks {
  runtime::Result<std::unique_ptr<PtnModule>> (*load)(
      const PtnSource& source,
      Program::Verification verification) = nullptr;
};

// This registration keeps PTN linkage optional: Module depends only on this
// hook contract, while a linked PTN provider supplies the implementation.
runtime::Error register_ptn_hooks(const PtnHooks& hooks);
const PtnHooks* get_ptn_hooks();

} // namespace executorch::extension::ET_PTN_MODULE_NAMESPACE::internal

namespace executorch::extension::ET_PTN_MODULE_NAMESPACE {

runtime::Result<std::unique_ptr<internal::PtnModule>> load_ptn(
    const internal::PtnSource& source,
    internal::Program::Verification verification);

} // namespace executorch::extension::ET_PTN_MODULE_NAMESPACE