Skip to content

Commit 4f6ef35

Browse files
wooway777wooway777
authored
feat: add cuda conv infinilm (#660)
1 parent 65e4ed4 commit 4f6ef35

8 files changed

Lines changed: 662 additions & 0 deletions

File tree

‎src/base/conv_infinilm.h‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
#ifndef INFINI_OPS_BASE_CONV_INFINILM_H_
2+
#define INFINI_OPS_BASE_CONV_INFINILM_H_
3+
4+
#include <cassert>
5+
#include <cstdint>
6+
#include <optional>
7+
#include <vector>
8+
9+
#include "operator.h"
10+
11+
namespace infini::ops {
12+
13+
class ConvInfinilm : public Operator<ConvInfinilm> {
14+
public:
15+
ConvInfinilm(const Tensor input, const Tensor weight,
16+
std::optional<Tensor> bias, const std::vector<int64_t> padding,
17+
const std::vector<int64_t> stride,
18+
const std::vector<int64_t> dilation, const int64_t groups,
19+
Tensor out)
20+
: input_shape_{input.shape()},
21+
input_strides_{input.strides()},
22+
weight_shape_{weight.shape()},
23+
weight_strides_{weight.strides()},
24+
out_shape_{out.shape()},
25+
out_strides_{out.strides()},
26+
bias_shape_{bias.has_value() ? bias->shape() : Tensor::Shape{}},
27+
bias_strides_{bias.has_value() ? bias->strides() : Tensor::Strides{}},
28+
input_type_{input.dtype()},
29+
weight_type_{weight.dtype()},
30+
out_type_{out.dtype()},
31+
bias_type_{bias.has_value() ? bias->dtype() : out.dtype()},
32+
padding_{padding},
33+
stride_{stride},
34+
dilation_{dilation},
35+
groups_{groups},
36+
spatial_ndim_{input.ndim() - 2},
37+
output_size_{out.numel()},
38+
kernel_size_{1},
39+
device_index_{out.device().index()},
40+
has_bias_{bias.has_value()} {
41+
assert(input.ndim() >= 3 && input.ndim() <= 5 &&
42+
"`ConvInfinilm` supports 1D, 2D, and 3D conv_infinilmolution");
43+
assert(input.ndim() == weight.ndim() && input.ndim() == out.ndim() &&
44+
"`ConvInfinilm` input, weight, and output ranks must match");
45+
assert(padding.size() == spatial_ndim_ && stride.size() == spatial_ndim_ &&
46+
dilation.size() == spatial_ndim_ &&
47+
"`ConvInfinilm` padding, stride, and dilation rank mismatch");
48+
assert(groups > 0 && "`ConvInfinilm` groups must be positive");
49+
assert(input_type_ == weight_type_ && input_type_ == out_type_ &&
50+
"`ConvInfinilm` input, weight, and output dtypes must match");
51+
assert(input_shape_[1] % groups == 0 &&
52+
"`ConvInfinilm` input channels must be divisible by groups");
53+
assert(weight_shape_[0] % groups == 0 &&
54+
"`ConvInfinilm` output channels must be divisible by groups");
55+
assert(weight_shape_[1] == input_shape_[1] / groups &&
56+
"`ConvInfinilm` weight input channels mismatch");
57+
assert(out_shape_[0] == input_shape_[0] &&
58+
"`ConvInfinilm` output batch size mismatch");
59+
assert(out_shape_[1] == weight_shape_[0] &&
60+
"`ConvInfinilm` output channels mismatch");
61+
assert(!out.HasBroadcastDim() &&
62+
"`ConvInfinilm` output must not have broadcasted dimensions");
63+
64+
if (has_bias_) {
65+
assert(bias_type_ == out_type_ && "`ConvInfinilm` bias dtype mismatch");
66+
assert(bias_shape_.size() == 1 && bias_shape_[0] == out_shape_[1] &&
67+
"`ConvInfinilm` bias shape must be `(out_channels,)`");
68+
}
69+
70+
for (std::size_t i = 0; i < spatial_ndim_; ++i) {
71+
assert(stride_[i] > 0 && "`ConvInfinilm` stride values must be positive");
72+
assert(dilation_[i] > 0 &&
73+
"`ConvInfinilm` dilation values must be positive");
74+
assert(padding_[i] >= 0 &&
75+
"`ConvInfinilm` padding values must be non-negative");
76+
77+
const auto expected = (input_shape_[i + 2] + 2 * padding_[i] -
78+
dilation_[i] * (weight_shape_[i + 2] - 1) - 1) /
79+
stride_[i] +
80+
1;
81+
assert(out_shape_[i + 2] == expected &&
82+
"`ConvInfinilm` output spatial shape mismatch");
83+
kernel_size_ *= weight_shape_[i + 2];
84+
}
85+
}
86+
87+
virtual void operator()(const Tensor input, const Tensor weight,
88+
std::optional<Tensor> bias,
89+
const std::vector<int64_t> padding,
90+
const std::vector<int64_t> stride,
91+
const std::vector<int64_t> dilation,
92+
const int64_t groups, Tensor out) const = 0;
93+
94+
protected:
95+
Tensor::Shape input_shape_;
96+
97+
Tensor::Strides input_strides_;
98+
99+
Tensor::Shape weight_shape_;
100+
101+
Tensor::Strides weight_strides_;
102+
103+
Tensor::Shape out_shape_;
104+
105+
Tensor::Strides out_strides_;
106+
107+
Tensor::Shape bias_shape_;
108+
109+
Tensor::Strides bias_strides_;
110+
111+
DataType input_type_;
112+
113+
DataType weight_type_;
114+
115+
DataType out_type_;
116+
117+
DataType bias_type_;
118+
119+
std::vector<int64_t> padding_;
120+
121+
std::vector<int64_t> stride_;
122+
123+
std::vector<int64_t> dilation_;
124+
125+
int64_t groups_{1};
126+
127+
Tensor::Size spatial_ndim_{0};
128+
129+
Tensor::Size output_size_{0};
130+
131+
Tensor::Size kernel_size_{1};
132+
133+
int device_index_{0};
134+
135+
bool has_bias_{false};
136+
};
137+
138+
} // namespace infini::ops
139+
140+
#endif

‎src/native/cuda/iluvatar/ops/conv_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_CONV_INFINILM_KERNEL_H_
2+
#define INFINI_OPS_ILUVATAR_CONV_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/conv_infinilm/kernel.h"
9+
10+
namespace infini::ops {
11+
12+
template <>
13+
class Operator<ConvInfinilm, Device::Type::kIluvatar>
14+
: public CudaConvInfinilm<Runtime<Device::Type::kIluvatar>> {
15+
public:
16+
using CudaConvInfinilm<Runtime<Device::Type::kIluvatar>>::CudaConvInfinilm;
17+
};
18+
19+
} // namespace infini::ops
20+
21+
#endif

‎src/native/cuda/metax/ops/conv_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_CONV_INFINILM_KERNEL_H_
2+
#define INFINI_OPS_METAX_CONV_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/conv_infinilm/kernel.h"
9+
10+
namespace infini::ops {
11+
12+
template <>
13+
class Operator<ConvInfinilm, Device::Type::kMetax>
14+
: public CudaConvInfinilm<Runtime<Device::Type::kMetax>> {
15+
public:
16+
using CudaConvInfinilm<Runtime<Device::Type::kMetax>>::CudaConvInfinilm;
17+
};
18+
19+
} // namespace infini::ops
20+
21+
#endif

‎src/native/cuda/moore/ops/conv_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_CONV_INFINILM_KERNEL_H_
2+
#define INFINI_OPS_MOORE_CONV_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/conv_infinilm/kernel.h"
10+
11+
namespace infini::ops {
12+
13+
template <>
14+
class Operator<ConvInfinilm, Device::Type::kMoore>
15+
: public CudaConvInfinilm<Runtime<Device::Type::kMoore>> {
16+
public:
17+
using CudaConvInfinilm<Runtime<Device::Type::kMoore>>::CudaConvInfinilm;
18+
};
19+
20+
} // namespace infini::ops
21+
22+
#endif

‎src/native/cuda/nvidia/ops/conv_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_CONV_INFINILM_KERNEL_H_
2+
#define INFINI_OPS_NVIDIA_CONV_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/conv_infinilm/kernel.h"
9+
10+
namespace infini::ops {
11+
12+
template <>
13+
class Operator<ConvInfinilm, Device::Type::kNvidia>
14+
: public CudaConvInfinilm<Runtime<Device::Type::kNvidia>> {
15+
public:
16+
using CudaConvInfinilm<Runtime<Device::Type::kNvidia>>::CudaConvInfinilm;
17+
};
18+
19+
} // namespace infini::ops
20+
21+
#endif

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

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
#ifndef INFINI_OPS_CUDA_CONV_INFINILM_KERNEL_CUH_
2+
#define INFINI_OPS_CUDA_CONV_INFINILM_KERNEL_CUH_
3+
4+
#include <cstddef>
5+
6+
#include "native/cuda/caster.cuh"
7+
#include "native/cuda/kernel_commons.cuh"
8+
9+
namespace infini::ops {
10+
11+
namespace {
12+
13+
__device__ __forceinline__ size_t
14+
OffsetFromCoordinates(const size_t* __restrict__ coords, size_t ndim,
15+
const ptrdiff_t* __restrict__ strides) {
16+
size_t offset = 0;
17+
for (size_t i = 0; i < ndim; ++i) {
18+
offset += coords[i] * strides[i];
19+
}
20+
return offset;
21+
}
22+
23+
} // namespace
24+
25+
template <Device::Type kDev, typename T, unsigned int block_size>
26+
__global__ void ConvInfinilmKernel(
27+
T* __restrict__ out, const T* __restrict__ input,
28+
const T* __restrict__ weight, const T* __restrict__ bias,
29+
const size_t* __restrict__ input_shape,
30+
const size_t* __restrict__ weight_shape,
31+
const size_t* __restrict__ out_shape,
32+
const ptrdiff_t* __restrict__ input_strides,
33+
const ptrdiff_t* __restrict__ weight_strides,
34+
const ptrdiff_t* __restrict__ out_strides,
35+
const ptrdiff_t* __restrict__ bias_strides,
36+
const int64_t* __restrict__ padding, const int64_t* __restrict__ stride,
37+
const int64_t* __restrict__ dilation, size_t output_size,
38+
size_t spatial_ndim, size_t kernel_size, int64_t groups, bool has_bias) {
39+
size_t linear = blockIdx.x * blockDim.x + threadIdx.x;
40+
if (linear >= output_size) {
41+
return;
42+
}
43+
44+
size_t coords[5] = {0, 0, 0, 0, 0};
45+
size_t tmp = linear;
46+
size_t ndim = spatial_ndim + 2;
47+
for (size_t axis = ndim; axis > 0; --axis) {
48+
size_t i = axis - 1;
49+
coords[i] = tmp % out_shape[i];
50+
tmp /= out_shape[i];
51+
}
52+
53+
size_t batch = coords[0];
54+
size_t out_channel = coords[1];
55+
size_t out_offset = OffsetFromCoordinates(coords, ndim, out_strides);
56+
size_t out_channels_per_group = weight_shape[0] / groups;
57+
size_t input_channels_per_group = weight_shape[1];
58+
size_t group = out_channel / out_channels_per_group;
59+
60+
float acc = 0.0f;
61+
if (has_bias) {
62+
acc =
63+
Caster<kDev>::template Cast<float>(bias[out_channel * bias_strides[0]]);
64+
}
65+
66+
for (size_t in_group_channel = 0; in_group_channel < input_channels_per_group;
67+
++in_group_channel) {
68+
size_t input_channel = group * input_channels_per_group + in_group_channel;
69+
70+
for (size_t kernel_linear = 0; kernel_linear < kernel_size;
71+
++kernel_linear) {
72+
size_t kernel_coords[3] = {0, 0, 0};
73+
size_t rem = kernel_linear;
74+
bool inside = true;
75+
size_t input_spatial[3] = {0, 0, 0};
76+
77+
for (size_t rev = spatial_ndim; rev > 0; --rev) {
78+
size_t d = rev - 1;
79+
kernel_coords[d] = rem % weight_shape[d + 2];
80+
rem /= weight_shape[d + 2];
81+
}
82+
83+
for (size_t d = 0; d < spatial_ndim; ++d) {
84+
int64_t pos = static_cast<int64_t>(coords[d + 2]) * stride[d] -
85+
padding[d] +
86+
static_cast<int64_t>(kernel_coords[d]) * dilation[d];
87+
if (pos < 0 || pos >= static_cast<int64_t>(input_shape[d + 2])) {
88+
inside = false;
89+
break;
90+
}
91+
input_spatial[d] = static_cast<size_t>(pos);
92+
}
93+
94+
if (!inside) {
95+
continue;
96+
}
97+
98+
size_t input_coords[5] = {batch, input_channel, 0, 0, 0};
99+
size_t weight_coords[5] = {out_channel, in_group_channel, 0, 0, 0};
100+
for (size_t d = 0; d < spatial_ndim; ++d) {
101+
input_coords[d + 2] = input_spatial[d];
102+
weight_coords[d + 2] = kernel_coords[d];
103+
}
104+
105+
float x = Caster<kDev>::template Cast<float>(
106+
input[OffsetFromCoordinates(input_coords, ndim, input_strides)]);
107+
float w = Caster<kDev>::template Cast<float>(
108+
weight[OffsetFromCoordinates(weight_coords, ndim, weight_strides)]);
109+
acc += x * w;
110+
}
111+
}
112+
113+
out[out_offset] = Caster<kDev>::template Cast<T>(acc);
114+
}
115+
116+
} // namespace infini::ops
117+
118+
#endif

0 commit comments

Comments
 (0)