<commit_before>#include <primitiv/config.h>

#include <primitiv/devices/cuda/device.h>
#include <primitiv/devices/cuda/ops/common.h>
#include <primitiv/internal/cuda/utils.h>

namespace {

__constant__ std::uint32_t permute_dims_x_strides[primitiv::Shape::MAX_DEPTH];
__constant__ std::uint32_t permute_dims_y_strides[primitiv::Shape::MAX_DEPTH];

// TODO(vbkaisetsu):
// Implove implementation of permute_dims.
// This function uses for-loops in the kernel code. It becomes slower than
// no-loop implementation.
__global__ void permute_dims_fw_dev(
    const float *px, const std::uint32_t ndims, const std::uint32_t size,
    float *py) {
  const std::uint32_t i = IDX;
  const std::uint32_t bid_z = IDY;
  const std::uint32_t ofs = bid_z * size;
  if (i < size) {
    std::uint32_t tmp = i;
    std::uint32_t j = 0;
    for (std::uint32_t d = 0; d < ndims; ++d) {
      const std::uint32_t p = tmp / permute_dims_x_strides[d];
      tmp -= p * permute_dims_x_strides[d];
      j += p * permute_dims_y_strides[d];
    }
    py[ofs + j] = px[ofs + i];
  }
}

__global__ void permute_dims_bw_dev(
    const float *py, const std::uint32_t ndims, const std::uint32_t size,
    float *px) {
  const std::uint32_t i = IDX;
  const std::uint32_t bid_z = IDY;
  const std::uint32_t ofs = bid_z * size;
  if (i < size) {
    std::uint32_t tmp = i;
    std::uint32_t j = 0;
    for (std::uint32_t d = 0; d < ndims; ++d) {
      const std::uint32_t p = tmp / permute_dims_x_strides[d];
      tmp -= p * permute_dims_x_strides[d];
      j += p * permute_dims_y_strides[d];
    }
    px[ofs + i] += py[ofs + j];
  }
}

}  // namespace

namespace primitiv {
namespace devices {

void CUDA::permute_dims_fw_impl(
    const Tensor &x, const std::vector<std::uint32_t> &perm,
    Tensor &y) {
  const std::uint32_t ndims = perm.size();
  const std::uint32_t bs = x.shape().batch();
  const std::uint32_t size = x.shape().volume();
  const std::uint32_t g1 = GRID_SIZE(size, dim1_x_);
  std::vector<std::uint32_t> x_strides(ndims);
  std::vector<std::uint32_t> y_strides(ndims);
  for (std::uint32_t i = 0; i < ndims; ++i) {<commit_msg>Add comments<commit_after>#include <primitiv/config.h>

#include <primitiv/devices/cuda/device.h>
#include <primitiv/devices/cuda/ops/common.h>
#include <primitiv/internal/cuda/utils.h>

namespace {

__constant__ std::uint32_t permute_dims_x_strides[primitiv::Shape::MAX_DEPTH];
__constant__ std::uint32_t permute_dims_y_strides[primitiv::Shape::MAX_DEPTH];

// TODO(vbkaisetsu):
// Implove implementation of permute_dims.
// This function uses for-loops in the kernel code. It becomes slower than
// no-loop implementation.
__global__ void permute_dims_fw_dev(
    const float *px, const std::uint32_t ndims, const std::uint32_t size,
    float *py) {
  const std::uint32_t i = IDX;
  const std::uint32_t bid_z = IDY;
  const std::uint32_t ofs = bid_z * size;
  if (i < size) {
    std::uint32_t tmp = i;
    std::uint32_t j = 0;
    // TODO(vbkaisetsu):
    // Implove implementation
    for (std::uint32_t d = 0; d < ndims; ++d) {
      const std::uint32_t p = tmp / permute_dims_x_strides[d];
      tmp -= p * permute_dims_x_strides[d];
      j += p * permute_dims_y_strides[d];
    }
    py[ofs + j] = px[ofs + i];
  }
}

__global__ void permute_dims_bw_dev(
    const float *py, const std::uint32_t ndims, const std::uint32_t size,
    float *px) {
  const std::uint32_t i = IDX;
  const std::uint32_t bid_z = IDY;
  const std::uint32_t ofs = bid_z * size;
  if (i < size) {
    std::uint32_t tmp = i;
    std::uint32_t j = 0;
    // TODO(vbkaisetsu):
    // Implove implementation
    for (std::uint32_t d = 0; d < ndims; ++d) {
      const std::uint32_t p = tmp / permute_dims_x_strides[d];
      tmp -= p * permute_dims_x_strides[d];
      j += p * permute_dims_y_strides[d];
    }
    px[ofs + i] += py[ofs + j];
  }
}

}  // namespace

namespace primitiv {
namespace devices {

void CUDA::permute_dims_fw_impl(
    const Tensor &x, const std::vector<std::uint32_t> &perm,
    Tensor &y) {
  const std::uint32_t ndims = perm.size();
  const std::uint32_t bs = x.shape().batch();
  const std::uint32_t size = x.shape().volume();
  const std::uint32_t g1 = GRID_SIZE(size, dim1_x_);
  std::vector<std::uint32_t> x_strides(ndims);
  std::vector<std::uint32_t> y_strides(ndims);
  for (std::uint32_t i = 0; i < ndims; ++i) {
{'fraction_non_alphanumeric': 0.14387497127097218, 'fraction_numerical': 0.02895886003217651, 'mean_word_length': 3.761487964989059, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, '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': '5383554', 'n_tokens_mistral': 1830, 'n_tokens_neox': 1793, 'n_words': 393}