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/ROCmLoops.cuh (13526B)

#pragma once

// This file provides two functions to help write GPU elementwise kernels:
//
//   gpu_kernel(TensorIterator iter, <lambda>)
//   gpu_kernel_with_scalars(TensorIterator iter, <lambda>)
//
// The gpu_kernel_with_scalars generates specializations that support a
// single scalar CPU argument, such as from `cuda_tensor + 5`. The CPU scalar
// is lifted to a kernel parameter instead of copying to device memory.
// This should be  used in conjunction with TensorIterator::allow_cpu_scalars_,
// which is the default for TensorIterator::binary_op. Otherwise, all inputs
// and the output must be on the GPU.
//
// For example, to write a reciprocal kernel for GPU float Tensors:
//
//   gpu_kernel(iter, []GPU_LAMBDA(float a) {
//    return 1.0f / a;
//   });
//
// To write a multiplication kernel for GPU float Tensors where one argument
// may be a CPU scalar:
//
//   gpu_kernel_with_scalars(iter, []GPU_LAMBDA(float a, float b) {
//     return a * b;
//   });
//
// See BinaryOpsKernel.cu for the complete implementation
//

#include <type_traits>

#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/core/Array.h>
#include <ATen/cuda/detail/OffsetCalculator.cuh>
#include <ATen/detail/FunctionTraits.h>
#include <ATen/native/TensorIterator.h>
#include <c10/macros/Macros.h>
#include <c10/core/ScalarType.h>
#include <c10/util/TypeCast.h>

// Marks a lambda as executable on both the host and device. The __host__
// attribute is important so that we can access static type information from
// the host, even if the function is typically only executed on the device.
#ifndef GPU_LAMBDA
#define GPU_LAMBDA __host__ __device__
#endif

#ifdef __NVCC__
#define ASSERT_HOST_DEVICE_LAMBDA(type) \
  static_assert(__nv_is_extended_host_device_lambda_closure_type(type), \
                #type " must be a __host__ __device__ lambda")
#else
#define ASSERT_HOST_DEVICE_LAMBDA(type)
#endif

static constexpr int launch_size_1d = 512;
static constexpr int launch_size_nd = 128;
static constexpr int launch_bound2 = 4;

namespace at { namespace native {

// See [NOTE: Complex Operator Unification]
// std::complex and thrust::complex don't work with some !needs_dynamic_casting optimizations.
// They always currently map to !needs_dynamic_casting even though we sometimes rely on the ability
// to reinterpret_cast between these representations.
// In order to separate these concerns, we have a check for non-c10 complex separately.
template<typename func_t, int nargs=function_traits<func_t>::arity>
struct uses_non_c10_complex {
  constexpr static bool check() {
    using traits = function_traits<func_t>;
    using type = typename traits::template arg<nargs - 1>::type;
    constexpr bool non_c10_complex =
        std::is_same<std::complex<float>, type>::value
        || std::is_same<std::complex<double>, type>::value
        || std::is_same<thrust::complex<float>, type>::value
        || std::is_same<thrust::complex<double>, type>::value;

return c10::guts::if_constexpr<non_c10_complex>([]() {
      return true;
    }, /* else */ []() {
      return uses_non_c10_complex<func_t, nargs - 1>::check();
    });
  }
};

template<typename func_t>
struct uses_non_c10_complex<func_t, 0> {
  constexpr static bool check() {
    using traits = function_traits<func_t>;
    using type = typename traits::result_type;
    constexpr bool non_c10_complex =
        std::is_same<std::complex<float>, type>::value
        || std::is_same<std::complex<double>, type>::value
        || std::is_same<thrust::complex<float>, type>::value
        || std::is_same<thrust::complex<double>, type>::value;

return non_c10_complex;
  }
};

// NOTE: @zasdfgbnm is currently working on rewriting the gpu loops.
// Some of the old codes has been moved to namespace legacy, and
// new codes will be put into namespace modern. These two namespaces
// will coexists for a while until the rewrite is done. Once the rewrite
// is done, we will remove the legacy and modern namespace and everything
// will be in at::native directly.
namespace legacy {

template<int nt, int vt, typename func_t>
C10_LAUNCH_BOUNDS_2(nt, launch_bound2)
__global__ void elementwise_kernel(int N, func_t f) {
  int tid = threadIdx.x;
  int nv = nt * vt;
  int idx = nv * blockIdx.x + tid;
  #pragma unroll
  for (int i = 0; i < vt; i++) {
    if (idx < N) {
      f(idx);
      idx += nt;
    }
  }
}

template<int nt, int vt, typename func_t>
static void launch_kernel(int64_t N, const func_t& f) {
  TORCH_INTERNAL_ASSERT(N >= 0 && N <= std::numeric_limits<int32_t>::max());
  if (N == 0) {
    return;
  }
  dim3 block(nt);
  dim3 grid((N + block.x * vt - 1) / (block.x * vt));
  auto stream = at::cuda::getCurrentCUDAStream();
  elementwise_kernel<nt, vt, func_t><<<grid, block, 0, stream>>>(N, f);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

template <typename traits, typename func_t, typename index_t, size_t... INDEX>
C10_HOST_DEVICE typename traits::result_type
invoke_impl(const func_t &f, char *const C10_RESTRICT data[], const index_t strides[], int i,
            std::index_sequence<INDEX...>) {
  return f(*(typename traits::template arg<INDEX>::type*)(data[INDEX] + i * strides[INDEX])...);
}

template <typename func_t, typename index_t, typename traits = function_traits<func_t>>
C10_HOST_DEVICE typename traits::result_type
invoke(const func_t &f, char *const C10_RESTRICT data[], const index_t strides[], int i) {
  using Indices = std::make_index_sequence<traits::arity>;
  return invoke_impl<traits>(f, data, strides, i, Indices{});
}

template <typename traits, typename func_t, typename index_t, size_t... I>
C10_HOST_DEVICE typename traits::result_type
invoke_impl(const func_t &f, char *const C10_RESTRICT data[], const index_t strides[], const ScalarType dtypes[], int i,
            std::index_sequence<I...>) {
  return f(c10::fetch_and_cast<typename traits::template arg<I>::type>(dtypes[I], data[I] + i * strides[I])...);
}

template <typename func_t, typename index_t, typename traits = function_traits<func_t>>
C10_HOST_DEVICE typename traits::result_type
invoke(const func_t &f, char *const C10_RESTRICT data[], const index_t strides[], const ScalarType dtypes[], int i) {
  using Indices = std::make_index_sequence<traits::arity>;
  return invoke_impl<traits>(f, data, strides, dtypes, i, Indices{});
}

} // namespace legacy

// See the note for namespace legacy above.
namespace modern {

namespace detail {

template <typename func_t, typename array_t, std::size_t... I>
__device__ inline constexpr decltype(auto) invoke_with_array_impl(func_t f, array_t t, std::index_sequence<I...>)
{
    return f(t[I]...);
}
template <typename func_t, typename array_t>
__device__ inline constexpr decltype(auto) invoke_with_array(func_t f, array_t a) {
  constexpr auto arity = function_traits<func_t>::arity;
  return invoke_with_array_impl(f, a, std::make_index_sequence<arity>{});
}

namespace arg_type {

// We need a way to compute the argument type of a function. But
// for nullary function, it does not really have an argument type
// in this case, we still need to return a valid type, but we don't
// really care what type this is.

struct dont_care {};

template <typename func_t, std::size_t arity>
struct arg_type_helper {
  using type = typename function_traits<func_t>::template arg<0>::type;
};

template <typename func_t>
struct arg_type_helper<func_t, 0> {
  using type = dont_care;
};

template <typename func_t>
using type = typename arg_type_helper<func_t, function_traits<func_t>::arity>::type;

}  // namespace arg_type

template<typename func_t, int remaining=function_traits<func_t>::arity-1>
struct has_same_arg_types {
  using traits = function_traits<func_t>;
  static constexpr bool value = std::is_same<
      typename traits::template arg<remaining>::type,
      typename traits::template arg<remaining-1>::type
    >::value && has_same_arg_types<func_t, remaining-1>::value;
};

template<typename func_t>
struct has_same_arg_types<func_t, 0> {
  static constexpr bool value = true;
};

template<typename func_t>
struct has_same_arg_types<func_t, -1> {
  static constexpr bool value = true;
};

}  // namespace detail

template<typename func_t, typename array_t>
C10_LAUNCH_BOUNDS_1(num_threads)
__global__ void elementwise_kernel(int N, func_t f, array_t data) {
  // Assumption:
  // 1. all arguments of `f` have the same type, which could be different from the return type of `f`
  // 2. all tensors are contiguous, that is: stride == sizeof(type) for all tensors

using traits = function_traits<func_t>;
  using return_t = typename traits::result_type;
  using arg_t = detail::arg_type::type<func_t>;
  constexpr int arity = traits::arity;

// We need to create array to hold all the arguments, for nullary `f`, this means array of size 0.
  // Unfortunately the compiler don't allow us to create array of 0 size, so for this case, we create
  // an array of size 1 and just don't use it.
  constexpr int nargs = traits::arity == 0 ? 1 : traits::arity;

int tid = threadIdx.x;
  int idx = block_work_size * blockIdx.x + tid;

// compute base pointers
  return_t *result_base = reinterpret_cast<return_t *>(data[0]) + idx;
  arg_t *args_base[nargs];
  #pragma unroll
  for (int i = 0; i < arity; i++) {
    args_base[i] = reinterpret_cast<arg_t *>(data[i + 1]) + idx;
  }

// fetch data
  return_t results[thread_work_size];
  arg_t args[thread_work_size][nargs];
  #pragma unroll
  for (int i = 0; i < thread_work_size; i++) {
    if (idx + num_threads * i < N) {
      #pragma unroll
      for (int j = 0; j < arity; j++) {
        args[i][j] = *(args_base[j] + i * num_threads);
      }
    }
  }

// compute
  #pragma unroll
  for (int i = 0; i < thread_work_size; i++) {
    if (idx + num_threads * i < N) {
      results[i] = detail::invoke_with_array<func_t, arg_t[nargs]>(f, args[i]);
    }
  }

// store data
  #pragma unroll
  for (int i = 0; i < thread_work_size; i++) {
    if (idx + num_threads * i < N) {
      *(result_base + i * num_threads) = results[i];
    }
  }
}

// TODO (@zasdfgbnm): this function assume trivial 1d and no dynamic casting
template<typename func_t, typename array_t, std::enable_if_t<detail::has_same_arg_types<func_t>::value, int> = 0>
static void launch_kernel(int64_t N, const func_t& f, array_t data) {
  TORCH_INTERNAL_ASSERT(N >= 0 && N <= std::numeric_limits<int32_t>::max());
  if (N == 0) {
    return;
  }
  int64_t grid = (N + block_work_size - 1) / block_work_size;
  auto stream = at::cuda::getCurrentCUDAStream();
  elementwise_kernel<func_t, array_t><<<grid, num_threads, 0, stream>>>(N, f, data);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

template<typename func_t, typename array_t, std::enable_if_t<!detail::has_same_arg_types<func_t>::value, int> = 0>
static void launch_kernel(int64_t N, const func_t& f, array_t data) {}

} // namespace modern

template <typename func_t>
void gpu_kernel_impl(TensorIteratorBase& iter, const func_t& f) {
  using traits = function_traits<func_t>;
  using arg0_t = typename traits::result_type;
  constexpr int ntensors = traits::arity + 1;

TORCH_INTERNAL_ASSERT(iter.can_use_32bit_indexing());
  TORCH_INTERNAL_ASSERT(iter.ntensors() == traits::arity + 1);
  bool non_c10_complex = uses_non_c10_complex<func_t>::check();

at::detail::Array<char*, ntensors> data;
  for (int i = 0; i < ntensors; i++) {
    data[i] = (char*)iter.data_ptr(i);
  }

at::detail::Array<ScalarType, ntensors> dtypes;
  for (int i = 0; i < ntensors; i++) {
    dtypes[i] = iter.tensor(i).scalar_type();
  }

int64_t numel = iter.numel();
  if (iter.is_trivial_1d()) {
    auto inner_strides = iter.get_inner_strides();
    at::detail::Array<int, ntensors> strides;
    for (int i = 0; i < ntensors; i++) {
      strides[i] = inner_strides[i];
    }

// TODO: can non_c10_complex go through the other path?  Need to verify.
    if (needs_dynamic_casting<func_t>::check(iter) || non_c10_complex) {
      legacy::launch_kernel<launch_size_1d, 1>(numel, [=]GPU_LAMBDA(int idx) {
        void* out = data[0] + strides[0] * idx;
        arg0_t result = legacy::invoke(f, &data.data[1], &strides.data[1], &dtypes.data[1], idx);
        c10::cast_and_store<arg0_t>(dtypes[0], out, result);
      });
    } else if (iter.has_contiguous_first_dim() && modern::detail::has_same_arg_types<func_t>::value) {
      modern::launch_kernel(numel, f, data);
    } else {
      legacy::launch_kernel<launch_size_1d, 1>(numel, [=]GPU_LAMBDA(int idx) {
        arg0_t* out = (arg0_t*)(data[0] + strides[0] * idx);
        *out = legacy::invoke(f, &data.data[1], &strides.data[1], idx);
      });
    }
  } else {
    auto offset_calc = ::make_offset_calculator<traits::arity + 1>(iter);
    // TODO: can non_c10_complex go through the other path?  Need to verify.
    if (needs_dynamic_casting<func_t>::check(iter) || non_c10_complex) {
      legacy::launch_kernel<launch_size_nd, launch_bound2>(numel, [=]GPU_LAMBDA(int idx) {
        auto offsets = offset_calc.get(idx);
        void* out = data[0] + offsets[0];
        arg0_t result = legacy::invoke(f, &data.data[1], &offsets.data[1], &dtypes.data[1], 1);
        c10::cast_and_store<arg0_t>(dtypes[0], out, result);
      });
    } else {
      legacy::launch_kernel<launch_size_nd, launch_bound2>(numel, [=]GPU_LAMBDA(int idx) {
        auto offsets = offset_calc.get(idx);
        arg0_t* out = (arg0_t*)(data[0] + offsets[0]);
        *out = legacy::invoke(f, &data.data[1], &offsets.data[1], 1);
      });
    }
  }
}

}} // namespace at::native

cmd: