InfiniOps/src/moore/add/kernel.h at 154c00969c2e5cc18700220df7e16f855e5bea8f · InfiniTensor/InfiniOps · 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
#ifndef INFINI_OPS_MOORE_ADD_KERNEL_H_
#define INFINI_OPS_MOORE_ADD_KERNEL_H_

#include <utility>

#include "cuda/add/kernel.h"

namespace infini::ops {

namespace add {

struct MooreBackend {
  using stream_t = musaStream_t;

  static constexpr auto malloc = [](auto&&... args) {
    return musaMalloc(std::forward<decltype(args)>(args)...);
  };

  static constexpr auto memcpy = [](auto&&... args) {
    return musaMemcpy(std::forward<decltype(args)>(args)...);
  };

  static constexpr auto free = [](auto&&... args) {
    return musaFree(std::forward<decltype(args)>(args)...);
  };

  static constexpr auto memcpyH2D = musaMemcpyHostToDevice;
};

}  // namespace add

template <>
class Operator<Add, Device::Type::kMoore> : public CudaAdd<add::MooreBackend> {
 public:
  using CudaAdd<add::MooreBackend>::CudaAdd;
};

}  // namespace infini::ops

#endif