InfiniOps/src/moore/swiglu/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
41
#ifndef INFINI_OPS_MOORE_SWIGLU_KERNEL_H_
#define INFINI_OPS_MOORE_SWIGLU_KERNEL_H_

#include <utility>

#include "cuda/swiglu/kernel.h"

namespace infini::ops {

namespace swiglu {

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 swiglu

template <>
class Operator<Swiglu, Device::Type::kMoore>
    : public CudaSwiglu<swiglu::MooreBackend> {
 public:
  using CudaSwiglu<swiglu::MooreBackend>::CudaSwiglu;
};

}  // namespace infini::ops

#endif