segment_kernel.cu 16.9 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
rusty1s's avatar
rusty1s committed
3
4
#include <ATen/cuda/detail/IndexUtils.cuh>
#include <ATen/cuda/detail/TensorInfo.cuh>
rusty1s's avatar
rusty1s committed
5

rusty1s's avatar
rusty1s committed
6
#include "atomics.cuh"
rusty1s's avatar
rusty1s committed
7
#include "compat.cuh"
rusty1s's avatar
rusty1s committed
8
#include "indptr.cuh"
rusty1s's avatar
rusty1s committed
9

rusty1s's avatar
rusty1s committed
10
#define THREADS 256
rusty1s's avatar
rusty1s committed
11
#define BLOCKS(TB, N) (TB * N + THREADS - 1) / THREADS
rusty1s's avatar
rusty1s committed
12
13
#define FULL_MASK 0xffffffff

rusty1s's avatar
rusty1s committed
14
enum ReductionType { ADD, MEAN, MIN, MAX };
rusty1s's avatar
rusty1s committed
15

rusty1s's avatar
rusty1s committed
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
#define AT_DISPATCH_REDUCTION_TYPES(reduce, ...)                               \
  [&] {                                                                        \
    if (reduce == "add") {                                                     \
      const ReductionType REDUCE = ADD;                                        \
      return __VA_ARGS__();                                                    \
    } else if (reduce == "mean") {                                             \
      const ReductionType REDUCE = MEAN;                                       \
      return __VA_ARGS__();                                                    \
    } else if (reduce == "min") {                                              \
      const ReductionType REDUCE = MIN;                                        \
      return __VA_ARGS__();                                                    \
    } else if (reduce == "max") {                                              \
      const ReductionType REDUCE = MAX;                                        \
      return __VA_ARGS__();                                                    \
    }                                                                          \
  }()

template <typename scalar_t, ReductionType REDUCE> struct Reducer {
rusty1s's avatar
rusty1s committed
34
  static inline __host__ __device__ scalar_t init() {
rusty1s's avatar
rusty1s committed
35
36
37
    if (REDUCE == MIN) {
      return std::numeric_limits<scalar_t>::max();
    } else if (REDUCE == MAX) {
rusty1s's avatar
rusty1s committed
38
      return std::numeric_limits<scalar_t>::lowest();
rusty1s's avatar
rusty1s committed
39
40
41
42
43
    } else {
      return (scalar_t)0;
    }
  }

rusty1s's avatar
rusty1s committed
44
45
46
47
48
49
50
51
52
53
  static inline __host__ __device__ void update(scalar_t *val,
                                                scalar_t new_val) {
    if (REDUCE == ADD || REDUCE == MEAN) {
      *val = *val + new_val;
    } else if ((REDUCE == MIN && new_val < *val) ||
               (REDUCE == MAX && new_val > *val)) {
      *val = new_val;
    }
  }

rusty1s's avatar
rusty1s committed
54
55
  static inline __host__ __device__ void update(scalar_t *val, scalar_t new_val,
                                                int64_t *arg, int64_t new_arg) {
rusty1s's avatar
rusty1s committed
56
57
58
59
    if (REDUCE == ADD || REDUCE == MEAN) {
      *val = *val + new_val;
    } else if ((REDUCE == MIN && new_val < *val) ||
               (REDUCE == MAX && new_val > *val)) {
rusty1s's avatar
rusty1s committed
60
61
62
63
64
      *val = new_val;
      *arg = new_arg;
    }
  }

rusty1s's avatar
rusty1s committed
65
66
67
  static inline __host__ __device__ void write(scalar_t *address, scalar_t val,
                                               int64_t *arg_address,
                                               int64_t arg, int count) {
rusty1s's avatar
rusty1s committed
68
69
70
71
72
73
74
75
76
77
78
79
80
    if (REDUCE == ADD) {
      *address = val;
    } else if (REDUCE == MEAN) {
      *address = val / (scalar_t)max(count, 1);
    } else if (REDUCE == MIN || REDUCE == MAX) {
      if (count > 0) {
        *address = val;
        *arg_address = arg;
      } else {
        *address = (scalar_t)0;
      }
    }
  }
rusty1s's avatar
atomics  
rusty1s committed
81

rusty1s's avatar
rusty1s committed
82
  static inline __device__ void atomic_write(scalar_t *address, scalar_t val) {
rusty1s's avatar
atomics  
rusty1s committed
83
84
85
86
87
88
89
90
91
92
    if (REDUCE == ADD) {
      atomAdd(address, val);
    } else if (REDUCE == MEAN) {
      atomAdd(address, val);
    } else if (REDUCE == MIN && val < *address) {
      atomMin(address, val);
    } else if (REDUCE == MAX && val > *address) {
      atomMax(address, val);
    }
  }
rusty1s's avatar
rusty1s committed
93
};
rusty1s's avatar
rusty1s committed
94

rusty1s's avatar
rusty1s committed
95
96
97
98
99
100
template <typename scalar_t, ReductionType REDUCE, int TB>
__global__ void
segment_csr_kernel(const scalar_t *src_data,
                   const at::cuda::detail::TensorInfo<int64_t, int> indptr_info,
                   scalar_t *out_data, int64_t *arg_out_data, size_t N,
                   size_t E) {
rusty1s's avatar
rusty1s committed
101

rusty1s's avatar
atomics  
rusty1s committed
102
103
  // Each warp processes exactly `32/TB` rows and aggregates all row values
  // via a parallel reduction.
rusty1s's avatar
rusty1s committed
104

rusty1s's avatar
rusty1s committed
105
  int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
rusty1s's avatar
rusty1s committed
106
  int row_idx = thread_idx / TB;
rusty1s's avatar
rusty1s committed
107
108
  int lane_idx = thread_idx & (TB - 1);

rusty1s's avatar
rusty1s committed
109
  if (row_idx < N) {
rusty1s's avatar
rusty1s committed
110
    int offset = IndexPtrToOffset<int64_t>::get(row_idx, indptr_info);
rusty1s's avatar
rusty1s committed
111
    int row_start = __ldg(indptr_info.data + offset);
rusty1s's avatar
rusty1s committed
112
113
    int row_end = __ldg(indptr_info.data + offset +
                        indptr_info.strides[indptr_info.dims - 1]);
rusty1s's avatar
rusty1s committed
114

rusty1s's avatar
rusty1s committed
115
    scalar_t val = Reducer<scalar_t, REDUCE>::init();
rusty1s's avatar
atomics  
rusty1s committed
116
    int64_t arg, arg_tmp;
rusty1s's avatar
rusty1s committed
117

rusty1s's avatar
rusty1s committed
118
    offset = (row_idx / (indptr_info.sizes[indptr_info.dims - 1] - 1)) * E;
rusty1s's avatar
rusty1s committed
119
    for (int src_idx = row_start + lane_idx; src_idx < row_end; src_idx += TB) {
rusty1s's avatar
rusty1s committed
120
121
      Reducer<scalar_t, REDUCE>::update(&val, src_data[offset + src_idx], &arg,
                                        src_idx);
rusty1s's avatar
rusty1s committed
122
123
124
    }

#pragma unroll
rusty1s's avatar
rusty1s committed
125
126
    for (int i = TB / 2; i > 0; i /= 2) {
      // Parallel reduction inside a single warp.
rusty1s's avatar
rusty1s committed
127
      if (REDUCE == MIN || REDUCE == MAX)
rusty1s's avatar
atomics  
rusty1s committed
128
        arg_tmp = __shfl_down_sync(FULL_MASK, arg, i);
rusty1s's avatar
rusty1s committed
129
130
      Reducer<scalar_t, REDUCE>::update(
          &val, __shfl_down_sync(FULL_MASK, val, i), &arg, arg_tmp);
rusty1s's avatar
rusty1s committed
131
    }
rusty1s's avatar
rusty1s committed
132
133

    if (lane_idx == 0) {
rusty1s's avatar
rusty1s committed
134
135
136
      Reducer<scalar_t, REDUCE>::write(out_data + row_idx, val,
                                       arg_out_data + row_idx, arg,
                                       row_end - row_start);
rusty1s's avatar
rusty1s committed
137
138
139
140
    }
  }
}

rusty1s's avatar
rusty1s committed
141
142
template <typename scalar_t, ReductionType REDUCE>
__global__ void segment_csr_broadcast_kernel(
rusty1s's avatar
rusty1s committed
143
144
    const scalar_t *src_data,
    const at::cuda::detail::TensorInfo<int64_t, int> indptr_info,
rusty1s's avatar
rusty1s committed
145
    scalar_t *out_data, int64_t *arg_out_data, size_t N, size_t K, size_t E) {
rusty1s's avatar
rusty1s committed
146

rusty1s's avatar
rusty1s committed
147
148
149
  // Each thread processes exactly one row. It turned out that is more
  // efficient than using shared memory due to avoiding synchronization
  // barriers.
rusty1s's avatar
rusty1s committed
150

rusty1s's avatar
rusty1s committed
151
152
153
154
155
  int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
  int row_idx = thread_idx / K;
  int lane_idx = thread_idx % K;

  if (thread_idx < N * K) {
rusty1s's avatar
rusty1s committed
156
    int offset = IndexPtrToOffset<int64_t>::get(row_idx, indptr_info);
rusty1s's avatar
rusty1s committed
157
158
159
    int row_start = __ldg(indptr_info.data + offset);
    int row_end = __ldg(indptr_info.data + offset +
                        indptr_info.strides[indptr_info.dims - 1]);
rusty1s's avatar
rusty1s committed
160

rusty1s's avatar
rusty1s committed
161
162
    scalar_t val = Reducer<scalar_t, REDUCE>::init();
    int64_t arg;
rusty1s's avatar
rusty1s committed
163
164
165

    offset = (row_idx / (indptr_info.sizes[indptr_info.dims - 1] - 1)) * E * K;
    for (int src_idx = row_start; src_idx < row_end; src_idx++) {
rusty1s's avatar
rusty1s committed
166
167
      Reducer<scalar_t, REDUCE>::update(
          &val, src_data[offset + K * src_idx + lane_idx], &arg, src_idx);
rusty1s's avatar
rusty1s committed
168
169
    }

rusty1s's avatar
rusty1s committed
170
171
172
    Reducer<scalar_t, REDUCE>::write(out_data + thread_idx, val,
                                     arg_out_data + thread_idx, arg,
                                     row_end - row_start);
rusty1s's avatar
rusty1s committed
173
174
175
  }
}

rusty1s's avatar
rusty1s committed
176
177
178
std::tuple<at::Tensor, at::optional<at::Tensor>>
segment_csr_cuda(at::Tensor src, at::Tensor indptr,
                 at::optional<at::Tensor> out_opt, std::string reduce) {
179

rusty1s's avatar
rusty1s committed
180
  AT_ASSERTM(src.dim() >= indptr.dim(), "Input mismatch");
rusty1s's avatar
rusty1s committed
181
  for (int i = 0; i < indptr.dim() - 1; i++)
rusty1s's avatar
rusty1s committed
182
    AT_ASSERTM(src.size(i) == indptr.size(i), "Input mismatch");
rusty1s's avatar
rusty1s committed
183

rusty1s's avatar
rusty1s committed
184
  src = src.contiguous();
rusty1s's avatar
rusty1s committed
185
  auto reduce_dim = indptr.dim() - 1;
186
187
188

  at::Tensor out;
  if (out_opt.has_value()) {
rusty1s's avatar
rusty1s committed
189
    out = out_opt.value().contiguous();
190
191
    for (int i = 0; i < out.dim(); i++)
      if (i != reduce_dim)
rusty1s's avatar
rusty1s committed
192
        AT_ASSERTM(src.size(i) == out.size(i), "Input mismatch");
rusty1s's avatar
rusty1s committed
193
194
    AT_ASSERTM(out.size(reduce_dim) == indptr.size(reduce_dim) - 1,
               "Input mismatch");
195
196
197
198
199
  } else {
    auto sizes = src.sizes().vec();
    sizes[reduce_dim] = indptr.size(reduce_dim) - 1;
    out = at::empty(sizes, src.options());
  }
rusty1s's avatar
rusty1s committed
200

rusty1s's avatar
rusty1s committed
201
  at::optional<at::Tensor> arg_out = at::nullopt;
rusty1s's avatar
rusty1s committed
202
  int64_t *arg_out_data = nullptr;
rusty1s's avatar
rusty1s committed
203
204
  if (reduce == "min" || reduce == "max") {
    arg_out = at::full_like(out, src.size(reduce_dim), indptr.options());
rusty1s's avatar
rusty1s committed
205
    arg_out_data = arg_out.value().DATA_PTR<int64_t>();
rusty1s's avatar
rusty1s committed
206
207
  }

rusty1s's avatar
rusty1s committed
208
209
  auto N = out.size(reduce_dim) * (indptr.numel() / indptr.size(-1));
  auto K = out.numel() / N;
rusty1s's avatar
rusty1s committed
210
  auto E = src.size(reduce_dim);
rusty1s's avatar
rusty1s committed
211

rusty1s's avatar
rusty1s committed
212
  auto indptr_info = at::cuda::detail::getTensorInfo<int64_t, int>(indptr);
rusty1s's avatar
rusty1s committed
213
  auto stream = at::cuda::getCurrentCUDAStream();
rusty1s's avatar
rusty1s committed
214
  AT_DISPATCH_ALL_TYPES(src.scalar_type(), "segment_csr_kernel", [&] {
rusty1s's avatar
rusty1s committed
215
216
217
    auto src_data = src.DATA_PTR<scalar_t>();
    auto out_data = out.DATA_PTR<scalar_t>();

rusty1s's avatar
rusty1s committed
218
219
220
221
222
223
224
225
226
227
228
    AT_DISPATCH_REDUCTION_TYPES(reduce, [&] {
      if (K == 1) {
        segment_csr_kernel<scalar_t, REDUCE, 1>
            <<<BLOCKS(32, N), THREADS, 0, stream>>>(
                src_data, indptr_info, out_data, arg_out_data, N, E);
      } else {
        segment_csr_broadcast_kernel<scalar_t, REDUCE>
            <<<BLOCKS(1, N * K), THREADS, 0, stream>>>(
                src_data, indptr_info, out_data, arg_out_data, N, K, E);
      }
    });
rusty1s's avatar
rusty1s committed
229
230
  });

rusty1s's avatar
rusty1s committed
231
  return std::make_tuple(out, arg_out);
rusty1s's avatar
rusty1s committed
232
233
}

rusty1s's avatar
rusty1s committed
234
template <typename scalar_t, ReductionType REDUCE, bool HAS_VAL>
rusty1s's avatar
rusty1s committed
235
236
237
__global__ void
segment_coo_kernel(const scalar_t *src_data,
                   const at::cuda::detail::TensorInfo<int64_t, int> index_info,
rusty1s's avatar
rusty1s committed
238
                   scalar_t *out_data, size_t E, size_t N) {
rusty1s's avatar
rusty1s committed
239

rusty1s's avatar
rusty1s committed
240
241
242
243
244
245
  // Each thread processes exactly one entry. Within a warp, we perform a
  // parallel reduction across equal indices, and write the intermediate
  // result via atomics.

  int row_idx = blockIdx.x * blockDim.x + threadIdx.x;
  int lane_idx = row_idx & (32 - 1);
rusty1s's avatar
rusty1s committed
246
  int D = index_info.sizes[index_info.dims - 1];
rusty1s's avatar
rusty1s committed
247

rusty1s's avatar
rusty1s committed
248
249
250
251
  if (row_idx < E) {
    int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(
        row_idx, index_info);
    int idx = index_info.data[offset], next_idx;
rusty1s's avatar
rusty1s committed
252
    int out_idx = (row_idx / D) * N + idx;
rusty1s's avatar
atomics  
rusty1s committed
253

rusty1s's avatar
rusty1s committed
254
    scalar_t val = HAS_VAL ? src_data[row_idx] : (scalar_t)1, tmp;
rusty1s's avatar
rusty1s committed
255
256

#pragma unroll
rusty1s's avatar
rusty1s committed
257
    for (int i = 1; i < 32; i *= 2) {
rusty1s's avatar
atomics  
rusty1s committed
258
      // Parallel reduction inside a single warp.
rusty1s's avatar
rusty1s committed
259
260
      tmp = __shfl_up_sync(FULL_MASK, val, i);
      next_idx = __shfl_up_sync(FULL_MASK, idx, i);
rusty1s's avatar
rusty1s committed
261
262
263
264
265
      if (lane_idx >= i && row_idx / D == (row_idx - i) / D) {
        assert(idx >= next_idx);
        if (idx == next_idx)
          Reducer<scalar_t, REDUCE>::update(&val, tmp);
      }
rusty1s's avatar
rusty1s committed
266
267
    }

rusty1s's avatar
rusty1s committed
268
    next_idx = __shfl_down_sync(FULL_MASK, idx, 1);
rusty1s's avatar
rusty1s committed
269
270
    if (lane_idx == 32 - 1 || row_idx / D != (row_idx + 1) / D ||
        idx != next_idx)
rusty1s's avatar
rusty1s committed
271
      Reducer<scalar_t, REDUCE>::atomic_write(out_data + out_idx, val);
rusty1s's avatar
rusty1s committed
272
273
274
  }
}

rusty1s's avatar
rusty1s committed
275
276
277
278
279
280
281
template <typename scalar_t>
__global__ void segment_coo_arg_kernel(
    const scalar_t *src_data,
    const at::cuda::detail::TensorInfo<int64_t, int> index_info,
    scalar_t *out_data, int64_t *arg_out_data, size_t E, size_t N) {

  int row_idx = blockIdx.x * blockDim.x + threadIdx.x;
rusty1s's avatar
rusty1s committed
282
  int D = index_info.sizes[index_info.dims - 1];
rusty1s's avatar
rusty1s committed
283
284
285
286
287

  if (row_idx < E) {
    int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(
        row_idx, index_info);
    int idx = index_info.data[offset];
rusty1s's avatar
rusty1s committed
288
    int out_idx = (row_idx / D) * N + idx;
rusty1s's avatar
rusty1s committed
289
290
291

    scalar_t val = __ldg(out_data + out_idx);
    if (src_data[row_idx] == val)
rusty1s's avatar
rusty1s committed
292
      arg_out_data[out_idx] = row_idx % D;
rusty1s's avatar
rusty1s committed
293
294
295
  }
}

rusty1s's avatar
rusty1s committed
296
297
template <typename scalar_t, ReductionType REDUCE, int TB>
__global__ void segment_coo_broadcast_kernel(
rusty1s's avatar
rusty1s committed
298
299
    const scalar_t *src_data,
    const at::cuda::detail::TensorInfo<int64_t, int> index_info,
rusty1s's avatar
rusty1s committed
300
    scalar_t *out_data, size_t E, size_t K, size_t N) {
rusty1s's avatar
rusty1s committed
301

rusty1s's avatar
rusty1s committed
302
303
304
  // Each thread processes a single column and `TB` index entries. Coalesced
  // read and write is performed in column-major order. The intermediate
  // results are written via atomics.
rusty1s's avatar
rusty1s committed
305

rusty1s's avatar
rusty1s committed
306
307
308
309
310
  int D = index_info.sizes[index_info.dims - 1];
  int E_1 = E / D;
  int E_2 = D + D % TB;

  int row_idx = blockIdx.x * blockDim.y + threadIdx.y;
rusty1s's avatar
rusty1s committed
311
312
  int col_idx = blockIdx.y * blockDim.x + threadIdx.x;

rusty1s's avatar
rusty1s committed
313
314
315
316
  int dim_start = (row_idx * TB) / E_2;
  int row_start = (row_idx * TB) % E_2;

  if (dim_start < E_1 && col_idx < K) {
rusty1s's avatar
rusty1s committed
317
    int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(
rusty1s's avatar
rusty1s committed
318
319
        dim_start * D + row_start, index_info);
    int idx1 = __ldg(index_info.data + offset), idx2;
rusty1s's avatar
rusty1s committed
320

rusty1s's avatar
rusty1s committed
321
    scalar_t val = src_data[K * (dim_start * D + row_start) + col_idx];
rusty1s's avatar
rusty1s committed
322
323
324

#pragma unroll
    for (int i = 1; i < TB; i++) {
rusty1s's avatar
rusty1s committed
325
      if (row_start + i >= D)
rusty1s's avatar
rusty1s committed
326
327
        break;

rusty1s's avatar
rusty1s committed
328
329
      idx2 = __ldg(index_info.data + offset +
                   i * index_info.strides[index_info.dims - 1]);
330
      assert(idx1 <= idx2);
rusty1s's avatar
rusty1s committed
331
      if (idx1 == idx2) {
rusty1s's avatar
rusty1s committed
332
        Reducer<scalar_t, REDUCE>::update(
rusty1s's avatar
rusty1s committed
333
            &val, src_data[K * (dim_start * D + row_start + i) + col_idx]);
rusty1s's avatar
rusty1s committed
334
      } else {
rusty1s's avatar
rusty1s committed
335
        Reducer<scalar_t, REDUCE>::atomic_write(
rusty1s's avatar
rusty1s committed
336
337
            out_data + (dim_start * N + idx1) * K + col_idx, val);
        val = src_data[K * (dim_start * D + row_start + i) + col_idx];
rusty1s's avatar
rusty1s committed
338
339
340
341
      }
      idx1 = idx2;
    }

rusty1s's avatar
rusty1s committed
342
    Reducer<scalar_t, REDUCE>::atomic_write(
rusty1s's avatar
rusty1s committed
343
        out_data + (dim_start * N + idx1) * K + col_idx, val);
rusty1s's avatar
rusty1s committed
344
345
346
347
348
349
350
351
352
353
354
355
  }
}

template <typename scalar_t>
__global__ void segment_coo_arg_broadcast_kernel(
    const scalar_t *src_data,
    const at::cuda::detail::TensorInfo<int64_t, int> index_info,
    scalar_t *out_data, int64_t *arg_out_data, size_t E, size_t K, size_t N) {

  int thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
  int row_idx = thread_idx / K;
  int col_idx = thread_idx % K;
rusty1s's avatar
rusty1s committed
356
  int D = index_info.sizes[index_info.dims - 1];
rusty1s's avatar
rusty1s committed
357
358
359
360
361

  if (row_idx < E && col_idx < K) {
    int offset = at::cuda::detail::IndexToOffset<int64_t, int, -1>::get(
        row_idx, index_info);
    int idx = __ldg(index_info.data + offset);
rusty1s's avatar
rusty1s committed
362
    int out_idx = ((row_idx / D) * N + idx) * K + col_idx;
rusty1s's avatar
rusty1s committed
363
364
365

    scalar_t val = __ldg(out_data + out_idx);
    if (src_data[thread_idx] == val)
rusty1s's avatar
rusty1s committed
366
      arg_out_data[out_idx] = row_idx % D;
rusty1s's avatar
rusty1s committed
367
  }
rusty1s's avatar
rusty1s committed
368
369
}

rusty1s's avatar
rusty1s committed
370
371
372
std::tuple<at::Tensor, at::optional<at::Tensor>>
segment_coo_cuda(at::Tensor src, at::Tensor index, at::Tensor out,
                 std::string reduce) {
rusty1s's avatar
rusty1s committed
373
  AT_ASSERTM(src.dim() >= index.dim(), "Input mismatch");
rusty1s's avatar
rusty1s committed
374
  for (int i = 0; i < index.dim(); i++)
rusty1s's avatar
rusty1s committed
375
    AT_ASSERTM(src.size(i) == index.size(i), "Input mismatch");
rusty1s's avatar
rusty1s committed
376
377

  src = src.contiguous();
rusty1s's avatar
rusty1s committed
378
  out = out.contiguous();
rusty1s's avatar
rusty1s committed
379
  auto reduce_dim = index.dim() - 1;
rusty1s's avatar
rusty1s committed
380

rusty1s's avatar
rusty1s committed
381
382
  for (int i = 0; i < out.dim(); i++)
    if (i != reduce_dim)
rusty1s's avatar
rusty1s committed
383
      AT_ASSERTM(src.size(i) == out.size(i), "Input mismatch");
rusty1s's avatar
rusty1s committed
384

rusty1s's avatar
rusty1s committed
385
  at::optional<at::Tensor> arg_out = at::nullopt;
rusty1s's avatar
rusty1s committed
386
  int64_t *arg_out_data = nullptr;
rusty1s's avatar
rusty1s committed
387
388
  if (reduce == "min" || reduce == "max") {
    arg_out = at::full_like(out, src.size(reduce_dim), index.options());
rusty1s's avatar
rusty1s committed
389
    arg_out_data = arg_out.value().DATA_PTR<int64_t>();
rusty1s's avatar
rusty1s committed
390
391
  }

rusty1s's avatar
rusty1s committed
392
  auto E = index.numel();
rusty1s's avatar
rusty1s committed
393
394
  auto E_2 = index.size(reduce_dim);
  auto E_1 = index.numel() / E_2;
rusty1s's avatar
rusty1s committed
395
  auto K = src.numel() / E;
rusty1s's avatar
rusty1s committed
396
  auto N = out.size(reduce_dim);
rusty1s's avatar
rusty1s committed
397
  auto avg_len = (float)src.size(reduce_dim) / (float)out.size(reduce_dim);
rusty1s's avatar
rusty1s committed
398

rusty1s's avatar
rusty1s committed
399
400
  auto index_info = at::cuda::detail::getTensorInfo<int64_t, int>(index);
  auto stream = at::cuda::getCurrentCUDAStream();
rusty1s's avatar
rusty1s committed
401
  AT_DISPATCH_ALL_TYPES(src.scalar_type(), "segment_coo_kernel", [&] {
rusty1s's avatar
rusty1s committed
402
403
    auto src_data = src.DATA_PTR<scalar_t>();
    auto out_data = out.DATA_PTR<scalar_t>();
rusty1s's avatar
rusty1s committed
404

rusty1s's avatar
rusty1s committed
405
406
    AT_DISPATCH_REDUCTION_TYPES(reduce, [&] {
      if (K == 1) {
rusty1s's avatar
rusty1s committed
407
        segment_coo_kernel<scalar_t, REDUCE, true>
rusty1s's avatar
rusty1s committed
408
            <<<BLOCKS(1, E), THREADS, 0, stream>>>(src_data, index_info,
rusty1s's avatar
rusty1s committed
409
                                                   out_data, E, N);
rusty1s's avatar
rusty1s committed
410
411
      } else if (avg_len <= 8) {
        segment_coo_broadcast_kernel<scalar_t, REDUCE, 4>
rusty1s's avatar
rusty1s committed
412
413
414
            <<<dim3((E_1 * ((E_2 + 3) / 4) + 7) / 8, (K + 31) / 32),
               dim3(32, 8), 0, stream>>>(src_data, index_info, out_data, E, K,
                                         N);
rusty1s's avatar
rusty1s committed
415
416
      } else if (avg_len <= 16) {
        segment_coo_broadcast_kernel<scalar_t, REDUCE, 8>
rusty1s's avatar
rusty1s committed
417
418
419
            <<<dim3((E_1 * ((E_2 + 7) / 8) + 7) / 8, (K + 31) / 32),
               dim3(32, 8), 0, stream>>>(src_data, index_info, out_data, E, K,
                                         N);
rusty1s's avatar
rusty1s committed
420
421
      } else if (avg_len <= 32) {
        segment_coo_broadcast_kernel<scalar_t, REDUCE, 16>
rusty1s's avatar
rusty1s committed
422
            <<<dim3((E_1 * ((E_2 + 15) / 16) + 7) / 8, (K + 31) / 32),
rusty1s's avatar
rusty1s committed
423
424
               dim3(32, 8), 0, stream>>>(src_data, index_info, out_data, E, K,
                                         N);
rusty1s's avatar
rusty1s committed
425
426
      } else {
        segment_coo_broadcast_kernel<scalar_t, REDUCE, 32>
rusty1s's avatar
rusty1s committed
427
            <<<dim3((E_1 * ((E_2 + 31) / 32) + 7) / 8, (K + 31) / 32),
rusty1s's avatar
rusty1s committed
428
429
430
431
432
433
434
435
436
437
438
439
440
441
               dim3(32, 8), 0, stream>>>(src_data, index_info, out_data, E, K,
                                         N);
      }

      if (REDUCE == MIN || REDUCE == MAX) {
        if (K == 1) {
          segment_coo_arg_kernel<scalar_t>
              <<<BLOCKS(1, E), THREADS, 0, stream>>>(
                  src_data, index_info, out_data, arg_out_data, E, N);
        } else {
          segment_coo_arg_broadcast_kernel<scalar_t>
              <<<BLOCKS(1, E * K), THREADS, 0, stream>>>(
                  src_data, index_info, out_data, arg_out_data, E, K, N);
        }
rusty1s's avatar
rusty1s committed
442
443
      }
    });
rusty1s's avatar
rusty1s committed
444
  });
445

rusty1s's avatar
atomics  
rusty1s committed
446
  if (reduce == "mean") {
rusty1s's avatar
rusty1s committed
447
448
449
450
    auto sizes = index.sizes().vec();
    sizes[reduce_dim] = out.size(reduce_dim);
    auto count = at::zeros(sizes, out.options());

rusty1s's avatar
rusty1s committed
451
452
    AT_DISPATCH_ALL_TYPES(out.scalar_type(), "count_kernel", [&] {
      auto count_data = count.DATA_PTR<scalar_t>();
rusty1s's avatar
rusty1s committed
453
454
      segment_coo_kernel<scalar_t, ADD, false>
          <<<BLOCKS(1, E), THREADS, 0, stream>>>(nullptr, index_info,
rusty1s's avatar
rusty1s committed
455
                                                 count_data, E, N);
rusty1s's avatar
rusty1s committed
456
    });
rusty1s's avatar
rusty1s committed
457
458

    count.clamp_(1);
rusty1s's avatar
rusty1s committed
459
    arg_out = count;
rusty1s's avatar
rusty1s committed
460
461
462
463
464
465

    for (int i = reduce_dim + 1; i < out.dim(); i++) {
      count = count.unsqueeze(-1);
    }

    out.div_(count);
rusty1s's avatar
atomics  
rusty1s committed
466
467
  }

rusty1s's avatar
rusty1s committed
468
  return std::make_tuple(out, arg_out);
rusty1s's avatar
rusty1s committed
469
}