node_executor_exatensor.hpp 3.96 KB
Newer Older
1
/** ExaTN:: Tensor Runtime: Tensor graph node executor: Exatensor
2
REVISION: 2022/01/17
3

4
5
Copyright (C) 2018-2022 Dmitry Lyakh, Tiffany Mintz, Alex McCaskey
Copyright (C) 2018-2022 Oak Ridge National Laboratory (UT-Battelle)
6
7
8
9
10
11
12
13
14
15

Rationale:

**/

#ifndef EXATN_RUNTIME_EXATENSOR_NODE_EXECUTOR_HPP_
#define EXATN_RUNTIME_EXATENSOR_NODE_EXECUTOR_HPP_

#include "tensor_node_executor.hpp"

16
17
#include "talshxx.hpp"

18
19
#include <atomic>

20
21
22
23
24
25
26
namespace exatn {
namespace runtime {

class ExatensorNodeExecutor : public TensorNodeExecutor {

public:

27
28
29
30
31
  ExatensorNodeExecutor() = default;
  ExatensorNodeExecutor(const ExatensorNodeExecutor &) = delete;
  ExatensorNodeExecutor & operator=(const ExatensorNodeExecutor &) = delete;
  ExatensorNodeExecutor(ExatensorNodeExecutor &&) noexcept = delete;
  ExatensorNodeExecutor & operator=(ExatensorNodeExecutor &&) noexcept = delete;
32
  virtual ~ExatensorNodeExecutor() = default;
33

34
  void initialize(const ParamConf & parameters) override;
35

36
37
  void activateDryRun(bool dry_run) override;

38
39
  void activateFastMath() override;

40
  std::size_t getMemoryBufferSize() const override;
41

42
43
  std::size_t getMemoryUsage(std::size_t * free_mem) const override;

44
45
  double getTotalFlopCount() const override;

46
47
48
49
50
51
  int execute(numerics::TensorOpCreate & op,
              TensorOpExecHandle * exec_handle) override;
  int execute(numerics::TensorOpDestroy & op,
              TensorOpExecHandle * exec_handle) override;
  int execute(numerics::TensorOpTransform & op,
              TensorOpExecHandle * exec_handle) override;
52
53
54
55
  int execute(numerics::TensorOpSlice & op,
              TensorOpExecHandle * exec_handle) override;
  int execute(numerics::TensorOpInsert & op,
              TensorOpExecHandle * exec_handle) override;
56
57
58
59
  int execute(numerics::TensorOpAdd & op,
              TensorOpExecHandle * exec_handle) override;
  int execute(numerics::TensorOpContract & op,
              TensorOpExecHandle * exec_handle) override;
Dmitry I. Lyakh's avatar
Dmitry I. Lyakh committed
60
61
62
63
64
65
66
67
  int execute(numerics::TensorOpDecomposeSVD3 & op,
              TensorOpExecHandle * exec_handle) override;
  int execute(numerics::TensorOpDecomposeSVD2 & op,
              TensorOpExecHandle * exec_handle) override;
  int execute(numerics::TensorOpOrthogonalizeSVD & op,
              TensorOpExecHandle * exec_handle) override;
  int execute(numerics::TensorOpOrthogonalizeMGS & op,
              TensorOpExecHandle * exec_handle) override;
68
69
70
71
  int execute(numerics::TensorOpFetch & op,
              TensorOpExecHandle * exec_handle) override;
  int execute(numerics::TensorOpUpload & op,
              TensorOpExecHandle * exec_handle) override;
72
73
74
75
  int execute(numerics::TensorOpBroadcast & op,
              TensorOpExecHandle * exec_handle) override;
  int execute(numerics::TensorOpAllreduce & op,
              TensorOpExecHandle * exec_handle) override;
76

77
78
  bool sync(TensorOpExecHandle op_handle,
            int * error_code,
79
80
            bool wait = true) override;

81
  bool sync() override;
82

83
84
  bool discard(TensorOpExecHandle op_handle) override;

85
86
  bool prefetch(const numerics::TensorOperation & op) override;

87
88
  void clearCache() override;

89
90
91
  std::shared_ptr<talsh::Tensor> getLocalTensor(const numerics::Tensor & tensor,
                 const std::vector<std::pair<DimOffset,DimExtent>> & slice_spec) override;

92
93
94
95
  void * getTensorImage(const numerics::Tensor & tensor,
                        int device_kind,
                        int device_id,
                        std::size_t * size = nullptr) const override {return nullptr;}
96

97
98
99
100
  const std::string name() const override {return "exatensor-node-executor";}
  const std::string description() const override {return "ExaTENSOR tensor graph node executor";}
  std::shared_ptr<TensorNodeExecutor> clone() override {return std::make_shared<ExatensorNodeExecutor>();}

101
102
protected:
 //`ExaTENSOR executor state
103
 /** Size of the distributed Host memory buffer provided by ExaTENSOR in bytes **/
104
 std::atomic<std::size_t> exatensor_host_mem_buffer_size_;
105
106
107
108
109
110
};

} //namespace runtime
} //namespace exatn

#endif //EXATN_RUNTIME_EXATENSOR_NODE_EXECUTOR_HPP_