<commit_before>#pragma once

#include <algorithm>
#include <cstdint>
#include <type_traits>

#include "xchainer/cuda/cuda_runtime.h"
#include "xchainer/macro.h"
#include "xchainer/reduction_kernel_arg.h"

namespace xchainer {
namespace cuda {
namespace reduce_detail {

static constexpr int kMaxReductionBlockSize{512};
static constexpr int64_t kMaxGridSize{0x7fffffff};

int64_t RoundUpToPowerOf2(int64_t x) {
    --x;
    x |= x >> 1;
    x |= x >> 2;
    x |= x >> 4;
    x |= x >> 8;
    x |= x >> 16;
    x |= x >> 32;
    return x + 1;
}

template <typename In, typename Out, typename ReductionImpl, int8_t InNdim = kDynamicNdim, int8_t OutNdim = kDynamicNdim>
__global__ void ReductionKernel(
        ReductionKernelArg<In, Out, InNdim, OutNdim> arg, int out_block_size, int reduce_block_size, ReductionImpl impl) {
    using T = decltype(impl.Identity());

    extern __shared__ __align__(8) uint8_t work_bytes[];
    T* work = reinterpret_cast<T*>(work_bytes);
    int tid = threadIdx.x;

    int64_t reduce_block_offset = tid / out_block_size;
    int64_t reduce_offset = reduce_block_offset * arg.out_indexer.total_size();
    int64_t reduce_stride = reduce_block_size * arg.out_indexer.total_size();

    int64_t out_offset = tid % out_block_size;
    int64_t out_base = blockIdx.x * out_block_size;
    int64_t out_stride = gridDim.x * out_block_size;

    for (auto it_out = arg.out_indexer.It(out_base + out_offset, out_stride); it_out; ++it_out) {
        T accum = impl.Identity();

        for (auto it_in = arg.in_indexer.It(it_out.raw_index() + reduce_offset, reduce_stride); it_in; ++it_in) {
            int64_t i_reduce = it_in.raw_index() / arg.out_indexer.total_size();
            impl.Reduce(impl.MapIn(arg.in[it_in], i_reduce), accum);
        }

        if (out_block_size <= kMaxReductionBlockSize / 2) {<commit_msg>reduction: avoid div<commit_after>#pragma once

#include <algorithm>
#include <cstdint>
#include <type_traits>

#include "xchainer/cuda/cuda_runtime.h"
#include "xchainer/macro.h"
#include "xchainer/reduction_kernel_arg.h"

namespace xchainer {
namespace cuda {
namespace reduce_detail {

static constexpr int kMaxReductionBlockSize{512};
static constexpr int64_t kMaxGridSize{0x7fffffff};

int64_t RoundUpToPowerOf2(int64_t x) {
    --x;
    x |= x >> 1;
    x |= x >> 2;
    x |= x >> 4;
    x |= x >> 8;
    x |= x >> 16;
    x |= x >> 32;
    return x + 1;
}

template <typename In, typename Out, typename ReductionImpl, int8_t InNdim = kDynamicNdim, int8_t OutNdim = kDynamicNdim>
__global__ void ReductionKernel(
        ReductionKernelArg<In, Out, InNdim, OutNdim> arg, int out_block_size, int reduce_block_size, ReductionImpl impl) {
    using T = decltype(impl.Identity());

    extern __shared__ __align__(8) uint8_t work_bytes[];
    T* work = reinterpret_cast<T*>(work_bytes);
    int tid = threadIdx.x;

    int64_t reduce_block_offset = tid / out_block_size;
    int64_t reduce_offset = reduce_block_offset * arg.out_indexer.total_size();
    int64_t reduce_stride = reduce_block_size * arg.out_indexer.total_size();

    int64_t out_offset = tid % out_block_size;
    int64_t out_base = blockIdx.x * out_block_size;
    int64_t out_stride = gridDim.x * out_block_size;

    for (auto it_out = arg.out_indexer.It(out_base + out_offset, out_stride); it_out; ++it_out) {
        T accum = impl.Identity();

        int64_t i_reduce = reduce_block_offset;
        for (auto it_in = arg.in_indexer.It(it_out.raw_index() + reduce_offset, reduce_stride); it_in;
             ++it_in, i_reduce += reduce_block_size) {
            impl.Reduce(impl.MapIn(arg.in[it_in], i_reduce), accum);
        }

        if (out_block_size <= kMaxReductionBlockSize / 2) {
{'fraction_non_alphanumeric': 0.11320754716981132, 'fraction_numerical': 0.0215633423180593, 'mean_word_length': 3.8257477243172953, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5737484', 'n_tokens_mistral': 1411, 'n_tokens_neox': 1335, 'n_words': 299}