Skip to content

Commit a290b37

Browse files
wooway777wooway777
authored
feat: add cuda gelu infinilm (#675)
1 parent 4f6ef35 commit a290b37

8 files changed

Lines changed: 370 additions & 0 deletions

File tree

‎src/base/gelu_infinilm.h‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
#ifndef INFINI_OPS_BASE_GELU_INFINILM_H_
2+
#define INFINI_OPS_BASE_GELU_INFINILM_H_
3+
4+
#include <cassert>
5+
#include <string>
6+
7+
#include "operator.h"
8+
9+
namespace infini::ops {
10+
11+
class GeluInfinilm : public Operator<GeluInfinilm> {
12+
public:
13+
GeluInfinilm(const Tensor input, const std::string approximate, Tensor out)
14+
: input_shape_{input.shape()},
15+
input_strides_{input.strides()},
16+
input_type_{input.dtype()},
17+
out_shape_{out.shape()},
18+
out_strides_{out.strides()},
19+
out_type_{out.dtype()},
20+
approximate_{approximate},
21+
output_size_{out.numel()},
22+
ndim_{out.ndim()},
23+
is_input_contiguous_{input.IsContiguous()},
24+
is_out_contiguous_{out.IsContiguous()},
25+
device_index_{out.device().index()} {
26+
assert(input_shape_ == out_shape_ &&
27+
"`GeluInfinilm` input and output shapes must match");
28+
assert(input_type_ == out_type_ &&
29+
"`GeluInfinilm` input and output dtypes must match");
30+
assert((approximate.empty() || approximate == "none") &&
31+
"`GeluInfinilm` only supports exact approximation");
32+
assert(!out.HasBroadcastDim() &&
33+
"`GeluInfinilm` output must not have broadcasted dimensions");
34+
}
35+
36+
virtual void operator()(const Tensor input, const std::string approximate,
37+
Tensor out) const = 0;
38+
39+
protected:
40+
Tensor::Shape input_shape_;
41+
42+
Tensor::Strides input_strides_;
43+
44+
DataType input_type_;
45+
46+
Tensor::Shape out_shape_;
47+
48+
Tensor::Strides out_strides_;
49+
50+
DataType out_type_;
51+
52+
std::string approximate_{};
53+
54+
Tensor::Size output_size_{0};
55+
56+
Tensor::Size ndim_{0};
57+
58+
bool is_input_contiguous_{false};
59+
60+
bool is_out_contiguous_{false};
61+
62+
int device_index_{0};
63+
};
64+
65+
} // namespace infini::ops
66+
67+
#endif

‎src/native/cuda/iluvatar/ops/gelu_infinilm/kernel.h‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
#ifndef INFINI_OPS_ILUVATAR_GELU_INFINILM_KERNEL_H_
2+
#define INFINI_OPS_ILUVATAR_GELU_INFINILM_KERNEL_H_
3+
4+
#include <utility>
5+
6+
#include "native/cuda/iluvatar/caster.cuh"
7+
#include "native/cuda/iluvatar/runtime_.h"
8+
#include "native/cuda/ops/gelu_infinilm/kernel.h"
9+
10+
namespace infini::ops {
11+
12+
template <>
13+
class Operator<GeluInfinilm, Device::Type::kIluvatar>
14+
: public CudaGeluInfinilm<Runtime<Device::Type::kIluvatar>> {
15+
public:
16+
using CudaGeluInfinilm<Runtime<Device::Type::kIluvatar>>::CudaGeluInfinilm;
17+
};
18+
19+
} // namespace infini::ops
20+
21+
#endif

‎src/native/cuda/metax/ops/gelu_infinilm/kernel.h‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
#ifndef INFINI_OPS_METAX_GELU_INFINILM_KERNEL_H_
2+
#define INFINI_OPS_METAX_GELU_INFINILM_KERNEL_H_
3+
4+
#include <utility>
5+
6+
#include "native/cuda/metax/caster.cuh"
7+
#include "native/cuda/metax/runtime_.h"
8+
#include "native/cuda/ops/gelu_infinilm/kernel.h"
9+
10+
namespace infini::ops {
11+
12+
template <>
13+
class Operator<GeluInfinilm, Device::Type::kMetax>
14+
: public CudaGeluInfinilm<Runtime<Device::Type::kMetax>> {
15+
public:
16+
using CudaGeluInfinilm<Runtime<Device::Type::kMetax>>::CudaGeluInfinilm;
17+
};
18+
19+
} // namespace infini::ops
20+
21+
#endif

‎src/native/cuda/moore/ops/gelu_infinilm/kernel.h‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
#ifndef INFINI_OPS_MOORE_GELU_INFINILM_KERNEL_H_
2+
#define INFINI_OPS_MOORE_GELU_INFINILM_KERNEL_H_
3+
4+
#include <utility>
5+
6+
#include "native/cuda/moore/caster.cuh"
7+
#include "native/cuda/moore/polyfills.cuh"
8+
#include "native/cuda/moore/runtime_.h"
9+
#include "native/cuda/ops/gelu_infinilm/kernel.h"
10+
11+
namespace infini::ops {
12+
13+
template <>
14+
class Operator<GeluInfinilm, Device::Type::kMoore>
15+
: public CudaGeluInfinilm<Runtime<Device::Type::kMoore>> {
16+
public:
17+
using CudaGeluInfinilm<Runtime<Device::Type::kMoore>>::CudaGeluInfinilm;
18+
};
19+
20+
} // namespace infini::ops
21+
22+
#endif

‎src/native/cuda/nvidia/ops/gelu_infinilm/kernel.h‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
#ifndef INFINI_OPS_NVIDIA_GELU_INFINILM_KERNEL_H_
2+
#define INFINI_OPS_NVIDIA_GELU_INFINILM_KERNEL_H_
3+
4+
#include <utility>
5+
6+
#include "native/cuda/nvidia/caster.cuh"
7+
#include "native/cuda/nvidia/runtime_.h"
8+
#include "native/cuda/ops/gelu_infinilm/kernel.h"
9+
10+
namespace infini::ops {
11+
12+
template <>
13+
class Operator<GeluInfinilm, Device::Type::kNvidia>
14+
: public CudaGeluInfinilm<Runtime<Device::Type::kNvidia>> {
15+
public:
16+
using CudaGeluInfinilm<Runtime<Device::Type::kNvidia>>::CudaGeluInfinilm;
17+
};
18+
19+
} // namespace infini::ops
20+
21+
#endif

‎src/native/cuda/ops/gelu_infinilm/kernel.cuh‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
#ifndef INFINI_OPS_CUDA_GELU_INFINILM_KERNEL_CUH_
2+
#define INFINI_OPS_CUDA_GELU_INFINILM_KERNEL_CUH_
3+
4+
#include <cmath>
5+
#include <cstddef>
6+
7+
#include "native/cuda/caster.cuh"
8+
#include "native/cuda/kernel_commons.cuh"
9+
10+
namespace infini::ops {
11+
12+
namespace {
13+
14+
template <Device::Type kDev, typename T>
15+
__device__ __forceinline__ T GeluInfinilmExact(T x) {
16+
if constexpr (std::is_same_v<T, double>) {
17+
const double v = x;
18+
return 0.5 * v * (1.0 + erf(v * 0.70710678118654752440));
19+
} else {
20+
const float v = Caster<kDev>::template Cast<float>(x);
21+
const float y = 0.5f * v * (1.0f + erff(v * 0.70710678118654752440f));
22+
return Caster<kDev>::template Cast<T>(y);
23+
}
24+
}
25+
26+
} // namespace
27+
28+
template <Device::Type kDev, typename T, unsigned int block_size>
29+
__global__ void GeluInfinilmKernel(T* __restrict__ out,
30+
const T* __restrict__ input,
31+
const size_t* __restrict__ out_shape,
32+
const size_t* __restrict__ input_shape,
33+
const ptrdiff_t* __restrict__ out_strides,
34+
const ptrdiff_t* __restrict__ input_strides,
35+
size_t output_size, size_t ndim,
36+
bool out_contiguous, bool input_contiguous) {
37+
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
38+
39+
if (idx < output_size) {
40+
size_t out_idx =
41+
out_contiguous ? idx : IndexToOffset(idx, ndim, out_shape, out_strides);
42+
size_t input_idx =
43+
input_contiguous ? idx
44+
: IndexToOffset(idx, ndim, input_shape, input_strides);
45+
out[out_idx] = GeluInfinilmExact<kDev>(input[input_idx]);
46+
}
47+
}
48+
49+
} // namespace infini::ops
50+
51+
#endif

‎src/native/cuda/ops/gelu_infinilm/kernel.h‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
#ifndef INFINI_OPS_CUDA_GELU_INFINILM_KERNEL_H_
2+
#define INFINI_OPS_CUDA_GELU_INFINILM_KERNEL_H_
3+
4+
#include <algorithm>
5+
#include <cstddef>
6+
#include <cstring>
7+
#include <vector>
8+
9+
#include "base/gelu_infinilm.h"
10+
#include "common/generic_utils.h"
11+
#include "data_type.h"
12+
#include "dispatcher.h"
13+
#include "native/cuda/kernel_commons.cuh"
14+
#include "native/cuda/ops/gelu_infinilm/kernel.cuh"
15+
#include "native/cuda/runtime_utils.h"
16+
17+
namespace infini::ops {
18+
19+
template <typename Backend>
20+
class CudaGeluInfinilm : public GeluInfinilm {
21+
public:
22+
CudaGeluInfinilm(const Tensor input, const std::string approximate,
23+
Tensor out)
24+
: GeluInfinilm{input, approximate, out} {
25+
size_t shape_size = ndim_ * sizeof(*d_input_shape_);
26+
size_t strides_size = ndim_ * sizeof(*d_input_strides_);
27+
const size_t metadata_size = 2 * (shape_size + strides_size);
28+
std::vector<std::byte> metadata(metadata_size);
29+
30+
Backend::Malloc((void**)&d_metadata_, metadata_size);
31+
32+
size_t offset = 0;
33+
d_input_shape_ = reinterpret_cast<Tensor::Size*>(d_metadata_ + offset);
34+
std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size);
35+
offset += shape_size;
36+
37+
d_out_shape_ = reinterpret_cast<Tensor::Size*>(d_metadata_ + offset);
38+
std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size);
39+
offset += shape_size;
40+
41+
d_input_strides_ = reinterpret_cast<Tensor::Stride*>(d_metadata_ + offset);
42+
std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size);
43+
offset += strides_size;
44+
45+
d_out_strides_ = reinterpret_cast<Tensor::Stride*>(d_metadata_ + offset);
46+
std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size);
47+
48+
Backend::Memcpy(d_metadata_, metadata.data(), metadata_size,
49+
Backend::MemcpyHostToDevice);
50+
}
51+
52+
~CudaGeluInfinilm() { Backend::Free(d_metadata_); }
53+
54+
void operator()(const Tensor input, const std::string approximate,
55+
Tensor out) const override {
56+
(void)approximate;
57+
auto cuda_stream =
58+
static_cast<typename Backend::Stream>(stream_ ? stream_ : 0);
59+
int block_size = std::min(
60+
RuntimeUtils<Backend::kDeviceType>::GetOptimalBlockSize(), 1024);
61+
dim3 block(std::min(static_cast<Tensor::Size>(block_size), output_size_));
62+
dim3 grid(utils::CeilDiv(output_size_, block.x));
63+
64+
DispatchFunc<AllFloatTypes, List<128, 256, 512, 1024>>(
65+
{static_cast<int64_t>(out_type_), block_size},
66+
[&](auto list_tag) {
67+
using T = TypeMapType<Backend::kDeviceType, ListGet<0>(list_tag)>;
68+
constexpr int kBlockSize = ListGet<1>(list_tag);
69+
70+
GeluInfinilmKernel<Backend::kDeviceType, T, kBlockSize>
71+
<<<grid, block, 0, cuda_stream>>>(
72+
reinterpret_cast<T*>(out.data()),
73+
reinterpret_cast<const T*>(input.data()), d_out_shape_,
74+
d_input_shape_, d_out_strides_, d_input_strides_,
75+
output_size_, ndim_, is_out_contiguous_,
76+
is_input_contiguous_);
77+
},
78+
"CudaGeluInfinilm::operator()");
79+
}
80+
81+
private:
82+
std::byte* d_metadata_{nullptr};
83+
84+
Tensor::Size* d_input_shape_{nullptr};
85+
86+
Tensor::Size* d_out_shape_{nullptr};
87+
88+
Tensor::Stride* d_input_strides_{nullptr};
89+
90+
Tensor::Stride* d_out_strides_{nullptr};
91+
};
92+
93+
} // namespace infini::ops
94+
95+
#endif

‎tests/test_gelu_infinilm.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
import infini.ops
2+
import pytest
3+
import torch
4+
5+
from tests.utils import Payload, empty_strided, get_stream, randn_strided
6+
7+
8+
@pytest.mark.auto_act_and_assert
9+
@pytest.mark.parametrize(
10+
"shape, input_strides, out_strides, inplace",
11+
(
12+
((13, 4), None, None, False),
13+
((13, 4), None, None, True),
14+
((13, 4), (10, 1), (10, 1), False),
15+
((13, 4), (10, 1), (10, 1), True),
16+
((13, 4, 4), None, None, False),
17+
((13, 4, 4), None, None, True),
18+
((13, 4, 4), (20, 4, 1), (20, 4, 1), False),
19+
((13, 4, 4), (20, 4, 1), (20, 4, 1), True),
20+
((16, 5632), None, None, False),
21+
((16, 5632), None, None, True),
22+
((16, 5632), (13312, 1), (13312, 1), False),
23+
((16, 5632), (13312, 1), (13312, 1), True),
24+
((4, 4, 5632), None, None, False),
25+
((4, 4, 5632), None, None, True),
26+
((4, 4, 5632), (45056, 5632, 1), (45056, 5632, 1), False),
27+
((4, 4, 5632), (45056, 5632, 1), (45056, 5632, 1), True),
28+
),
29+
)
30+
@pytest.mark.parametrize(
31+
("dtype", "rtol", "atol"),
32+
(
33+
(torch.float64, 1e-6, 1e-6),
34+
(torch.float32, 1e-5, 1e-5),
35+
(torch.float16, 1e-3, 1e-3),
36+
(torch.bfloat16, 1e-2, 1e-2),
37+
),
38+
)
39+
def test_gelu_infinilm(
40+
shape, input_strides, out_strides, inplace, dtype, device, rtol, atol
41+
):
42+
if device == "musa" and dtype == torch.float64:
43+
pytest.skip("MUSA does not support float64 GELU_INFINILM")
44+
45+
input = randn_strided(shape, input_strides, dtype=dtype, device=device)
46+
out = (
47+
input
48+
if inplace
49+
else empty_strided(shape, out_strides, dtype=dtype, device=device)
50+
)
51+
52+
return Payload(
53+
_gelu_infinilm,
54+
_torch_gelu_infinilm,
55+
(input, out),
56+
{},
57+
rtol=rtol,
58+
atol=atol,
59+
)
60+
61+
62+
def _gelu_infinilm(input, out):
63+
infini.ops.gelu_infinilm(input, "none", out, stream=get_stream(input.device))
64+
65+
return out
66+
67+
68+
def _torch_gelu_infinilm(input, out):
69+
result = torch.nn.functional.gelu(input, approximate="none")
70+
out.copy_(result)
71+
72+
return out

0 commit comments

Comments
 (0)