FM / cuda

/usr/local/lib64/python3.6/site-packages/torch/include/ATen/native/cuda/

/usr/local/lib64/python3.6/site-packages/torch/include/ATen/native/cuda

Name	Size	Mode	Actions
BatchLinearAlgebraLib.h	3114	0644	edit dl rm
block_reduce.cuh	2549	0644	edit dl rm
CompositeRandomAccessor.h	929	0644	edit dl rm
CUDALoops.cuh	7598	0644	edit dl rm
CuFFTPlanCache.h	19282	0644	edit dl rm
CuFFTUtils.h	1892	0644	edit dl rm
DeviceSqrt.cuh	585	0644	edit dl rm
DistributionTemplates.h	27435	0644	edit dl rm
EmbeddingBackwardKernel.cuh	715	0644	edit dl rm
ForeachFunctors.cuh	16851	0644	edit dl rm
GridSampler.cuh	11316	0644	edit dl rm
im2col.cuh	6577	0644	edit dl rm
KernelUtils.cuh	2553	0644	edit dl rm
LaunchUtils.h	306	0644	edit dl rm
Loops.cuh	9997	0644	edit dl rm
Math.cuh	13840	0644	edit dl rm
MemoryAccess.cuh	12463	0644	edit dl rm
MiscUtils.h	3341	0644	edit dl rm
MultiTensorApply.cuh	7552	0644	edit dl rm
Normalization.cuh	74441	0644	edit dl rm
PersistentSoftmax.cuh	14635	0644	edit dl rm
Randperm.cuh	2114	0644	edit dl rm
Reduce.cuh	38784	0644	edit dl rm
Resize.cuh	1919	0644	edit dl rm
ROCmLoops.cuh	13526	0644	edit dl rm
SortingCommon.cuh	5688	0644	edit dl rm
SortingRadixSelect.cuh	11918	0644	edit dl rm
SortUtils.cuh	5549	0644	edit dl rm
TensorModeKernel.cuh	14391	0644	edit dl rm
UniqueCub.cuh	345	0644	edit dl rm
UpSample.cuh	7552	0644	edit dl rm
vol2col.cuh	8297	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/include/ATen/native/cuda/SortUtils.cuh (5549B)

#pragma once
#include <c10/macros/Macros.h>
#include <c10/util/Optional.h>

#include <ATen/ATen.h>
#include <ATen/cuda/detail/TensorInfo.cuh>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/native/cuda/SortingCommon.cuh>
#include <ATen/native/Resize.h>
#include <THC/THCNumerics.cuh> // for ScalarConvert

namespace at { namespace native {

template <typename T>
__device__ inline void swapVars(T& t1, T& t2) {
  T tmp = t1;
  t1 = t2;
  t2 = tmp;
}

template <typename Comparator, typename K, typename V>
__device__ inline void bitonicSwap(K& kA, V& vA, bool& validA,
                                   K& kB, V& vB, bool& validB,
                                   bool dir,
                                   const Comparator& comp) {
  // Invalid entries always sort to the end
  bool swap = (comp(kA, kB) && validA) || !validB;
  if (swap == dir) {
    swapVars(kA, kB);
    swapVars(vA, vB);
    swapVars(validA, validB);
  }
};

template <typename Comparator, typename K, typename V,
          typename IndexType, int Power2SortSize>
__device__ inline void bitonicSort(K keys[Power2SortSize],
                                   V values[Power2SortSize],
                                   bool valid[Power2SortSize],
                                   const Comparator& comp) {
#ifndef __HIP_PLATFORM_HCC__
#pragma unroll
#endif
  for (unsigned int size = 2; size < Power2SortSize; size *= 2) {
    bool flag = ((threadIdx.x & (size / 2)) != 0);

#ifndef __HIP_PLATFORM_HCC__
#pragma unroll
#endif
    for (unsigned int stride = size / 2; stride > 0; stride /= 2) {

__syncthreads();

unsigned int pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
      bitonicSwap<Comparator, K, V>(
        keys[pos], values[pos], valid[pos],
        keys[pos + stride], values[pos + stride], valid[pos + stride],
        flag, comp);
    }
  }

#ifndef __HIP_PLATFORM_HCC__
#pragma unroll
#endif
  for (unsigned int stride = Power2SortSize / 2; stride > 0; stride /= 2) {

__syncthreads();

unsigned int pos = 2 * threadIdx.x - (threadIdx.x & (stride - 1));
    bitonicSwap<Comparator, K, V>(
      keys[pos], values[pos], valid[pos],
      keys[pos + stride], values[pos + stride], valid[pos + stride],
      false, comp);
  }

__syncthreads();

}

// at::cuda::detail::TensorInfo version
// Sorts (key, value) pairs (in different tensors) in-place; i.e.,
// modifies the input `keys` and `values`
template <typename K, typename V,
          int KeyDims, int ValueDims,
          typename Comparator, typename IndexType, int Power2SortSize>
C10_LAUNCH_BOUNDS_1(1024)
__global__ void
bitonicSortKVInPlace(at::cuda::detail::TensorInfo<K, IndexType> keys,
                     IndexType keySlices,
                     IndexType keySliceSize,
                     IndexType keySliceStride,
                     at::cuda::detail::TensorInfo<V, IndexType> values,
                     IndexType valueSliceStride,
                     Comparator comp) {
  // Find the slice of the tensor that we are sorting
  const IndexType linearIndex = getLinearBlockId<IndexType>();
  // Tiling the slices could have us be out of bounds, if there are a
  // lot of slices to sort
  if (linearIndex >= keySlices) {
    return;
  }

__shared__ K sharedKeys[Power2SortSize];
  __shared__ V sharedValues[Power2SortSize];
  __shared__ bool sharedValid[Power2SortSize];

const IndexType keyStartOffset =
    at::cuda::detail::IndexToOffset<K, IndexType, KeyDims>::get(linearIndex, keys);
  const IndexType valueStartOffset =
    at::cuda::detail::IndexToOffset<V, IndexType, ValueDims>::get(linearIndex, values);

// If the sort size is 1, the data is already sorted
  if (Power2SortSize == 1) {
    return;
  } else {
    // Otherwise, each thread is responsible for loading and storing 2
    // elements. The sort size is guaranteed to be >= 2
    const int elem1 = threadIdx.x;
    const int elem2 = threadIdx.x + (Power2SortSize / 2);

bool valid1 = (elem1 < keySliceSize);
    K k1 = valid1 ?
      keys.data[keyStartOffset + elem1 * keySliceStride] : ScalarConvert<int, K>::to(0);
    V v1 = valid1 ?
      values.data[valueStartOffset + elem1 * valueSliceStride] : ScalarConvert<int, V>::to(0);

sharedKeys[elem1] = k1;
    sharedValues[elem1] = v1;
    sharedValid[elem1] = valid1;

bool valid2 = (elem2 < keySliceSize);
    K k2 = valid2 ?
      keys.data[keyStartOffset + elem2 * keySliceStride] : ScalarConvert<int, K>::to(0);
    V v2 = valid2 ?
      values.data[valueStartOffset + elem2 * valueSliceStride] : ScalarConvert<int, V>::to(0);

sharedKeys[elem2] = k2;
    sharedValues[elem2] = v2;
    sharedValid[elem2] = valid2;

// Sort!
    bitonicSort<Comparator, K, V, IndexType, Power2SortSize>(
      sharedKeys, sharedValues, sharedValid, comp);

// elem1 and elem2 values might be out-of-range, if the data size we are
    // sorting is smaller than half the power2 size
    if (valid1) {
      keys.data[keyStartOffset + elem1 * keySliceStride] =
        sharedKeys[elem1];
      values.data[valueStartOffset + elem1 * valueSliceStride] =
        sharedValues[elem1];
    }

if (valid2) {
      keys.data[keyStartOffset + elem2 * keySliceStride] =
        sharedKeys[elem2];
      values.data[valueStartOffset + elem2 * valueSliceStride] =
        sharedValues[elem2];
    }
  }
}

bool should_use_small_sort(const Tensor &self, int64_t dim);
void sortKeyValueInplace(const Tensor& key,
                         const Tensor& value,
                         int dim, bool dir);
}} // at::native

cmd: