Loops.hpp

#pragma once

#include "OffsetCalculator.hpp"
#include <array>
#include <tuple>



constexpr int group_size() { return 16 /* Use query interface for sub-group */ * 4; }
constexpr int item_work_size() { return 4; }
constexpr int group_work_size() { return item_work_size() * group_size(); }

#include "MemoryAccess.hpp"

namespace porting {

template<int N>
static OffsetCalculator<N> make_input_offset_calculator(const at::TensorIteratorBase& iter) {
  // array size can not be 0, this happens when N == 0
  constexpr int array_size = std::max<int>(N, 1);
  assert(N == iter.ntensors() - iter.noutputs());
  std::array<const int64_t*, array_size> strides;
  int64_t element_sizes[array_size];
  for (int i = 0; i < N; i++) {
    strides[i] = iter.strides(i + iter.noutputs()).data();
    element_sizes[i] = iter.element_size(i + iter.noutputs());
  }
  return OffsetCalculator<N>(iter.ndim(), iter.shape().data(), strides.data(), element_sizes);
}

template <int num_outputs = 1>
static OffsetCalculator<num_outputs> make_output_offset_calculator(const at::TensorIteratorBase& iter) {
  assert(num_outputs == iter.noutputs());
  std::array<const int64_t*, num_outputs> strides;
  int64_t element_sizes[num_outputs];
  for (int i = 0; i < num_outputs; i++) {
    strides[i] = iter.strides(i).data();
    element_sizes[i] = iter.element_size(i);
  }
  return OffsetCalculator<num_outputs>(iter.ndim(), iter.shape().data(), strides.data(), element_sizes);
}

template<typename func_t, typename policy_t>
inline void elementwise_kernel_helper(sycl::nd_item<1> pos, func_t f, policy_t policy) {
  using traits = function_traits<func_t>;
  using return_t = typename traits::result_type;
  using args_t = typename traits::ArgsTuple;

  int idx = pos.get_group(0);

  return_t results[item_work_size()];
  args_t args[item_work_size()];

  // load
  policy.load(pos, args, idx);

  // compute
  #pragma unroll (item_work_size())
  for (int i = 0; i < item_work_size(); i++) {
    if (policy.check_inbounds(pos.get_local_id(), i)) {
      results[i] = std::apply(f, args[i]);
    }
  }

  // store
  policy.store(pos, results, idx);
}

}

// Note:
// CUDA and ROCm get diverged in this PR:
//   https://github.com/pytorch/pytorch/pull/32383
// Because for some reason trying to enable vectorized
// memory access introduce regression on ROCm.

#include "SYCLLoops.hpp"

namespace porting {

/* Note [Jiterator]
The "jiterator" simply just-in-time compiles the same kernels that
Loops.cuh (and CUDALoops.cuh) usually build. This reduces build time,
build size, and CUDA context size.

The jiterator currently has some limitations, however. It cannot:
  - handle float16, bfloat16, or complex datatypes
  - handle scalar inputs

These improvements will likely come soon.

For examples of how to use the jiterator see the i1 and gcd kernel
implementations, which pass jittable strings implementing their
operations instead of the typical CUDA functors.
*/

// Entrypoint for jitted GPU kernels.
// Only handles elementwise unary and binary kernels with a
//   common dtype and a single output.
// NOTE: this assumes the op's iterator has a common_dtype.
// template <char const *name, typename return_type, typename compute_type, int arity>
// void jitted_gpu_kernel(TensorIteratorBase& iter, const std::string& f) {
//   // TODO: much of preamble is common to both jitted_gpu_kernel and gpu_kernel
//   //   Maybe it could be refactored?
//   for (int arg = 0; arg < iter.ntensors(); arg++) {
//     TORCH_INTERNAL_ASSERT(
//       iter.device(arg).is_cuda(),
//       "argument ", arg, ": expected a CUDA device but found ", iter.device(arg));
//   }
// 
//   if (iter.numel() == 0) {
//     return;
//   }
// 
//   if (!iter.can_use_32bit_indexing()) {
//     for (auto& sub_iter : iter.with_32bit_indexing()) {
//       jitted_gpu_kernel<name, return_type, compute_type, arity>(sub_iter, f);
//     }
// 
//     return;
//   }
// 
//   // Computes if dynamic casting is needed
//   // Dynamic casting is needed if an input's dtype differs from the common dtype
//   //   or if the result dtype differs from the output's dtype
//   // Note: this is intentionally divergent from calling needs_dynamic_casting,
//   //   which is more general and inspects a lambda to determine if dynamic
//   //   casting is needed.
//   bool needs_dynamic_casting = false;
// 
//   // Checks output
//   const ScalarType return_scalar_type = c10::CppTypeToScalarType<return_type>::value;
//   const auto dtype0 = iter.dtype(0);
//   if (dtype0 != return_scalar_type) {
//     needs_dynamic_casting = true;
//   }
//   // TODO: FIXME: support these datatypes!
//   TORCH_CHECK(dtype0 != kComplexDouble && dtype0 != kComplexFloat &&
//               dtype0 != kBFloat16 && dtype0 != at::kHalf,
//                 "Encountered an unsupported dtype ", dtype0, "!");
// 
//   // Checks input(s)
//   const ScalarType compute_scalar_type = c10::CppTypeToScalarType<compute_type>::value;
//   for (auto i = decltype(arity){1}; i < (arity + 1); ++i) {
//     const auto dtypei = iter.dtype(i);
//     if (dtypei != compute_scalar_type) {
//       needs_dynamic_casting = true;
//       // NOTE: can't short-circuit here yet because the dtype check below needs to run on every arg
//     }
//     TORCH_CHECK(dtypei != kComplexDouble && dtypei != kComplexFloat &&
//                 dtypei != kBFloat16 && dtypei != at::kHalf,
//                 "Encountered an unsupported dtype ", dtypei, "!");
//   }
// 
//   jitted_gpu_kernel_impl</*name*/ name,
//                   /*return_type=*/ return_type,
//                   /*compute_type=*/ compute_type,
//                   arity>(iter, f, needs_dynamic_casting);
// }

template <typename func_t>
void gpu_kernel(at::TensorIteratorBase& iter, const func_t& f) {
  if (iter.numel() == 0) {
    return;
  }

  if (!iter.can_use_32bit_indexing()) {
    // for (auto& sub_iter : iter.with_32bit_indexing()) {
    //   gpu_kernel(sub_iter, f);
    // }
    return; // Never be here
  }

  gpu_kernel_impl(iter, f);
}

template<typename arg1_t, typename arg2_t, typename return_t, typename func_t>
struct AUnaryFunctor {
  using traits = function_traits<func_t>;
  using opmath_arg1_t = typename traits::template arg<0>::type;
  return_t operator()(arg2_t b) const {
    return f(a, b);
  }
  // NB: scalar is stored in higher precision!
  AUnaryFunctor(func_t f_, opmath_arg1_t a_): f(f_), a(a_) {}
  private:
    func_t f;
    opmath_arg1_t a;
};

template<typename arg1_t, typename arg2_t, typename return_t, typename func_t>
struct BUnaryFunctor {
  using traits = function_traits<func_t>;
  using opmath_arg2_t = typename traits::template arg<1>::type;
  return_t operator()(arg1_t a) const {
    return f(a, b);
  }
  // NB: scalar is stored in higher precision!
  BUnaryFunctor(func_t f_, opmath_arg2_t b_): f(f_), b(b_) {}
  private:
    func_t f;
    opmath_arg2_t b;
};

// Though seemingly noop, this inserts casts from arg1_t to func_t's type
// (which may be higher precision), as well as casts to return_t
template <typename arg1_t, typename arg2_t, typename return_t, typename func_t>
struct BinaryFunctor {
  return_t operator()(arg1_t a, arg2_t b) const {
    return f(a, b);
  }
  BinaryFunctor(func_t f_): f(f_) {}
  private:
    func_t f;
};

// Unlike gpu_kernel_with_scalars, this allows you to pass a func_t which
// accepts inputs at higher precision (typically opmath_t), but then
// ensure that we load from memory at the correct precision (scalar_t)
// to avoid expensive loads.  For the whole sordid story see
// https://dev-discuss.pytorch.org/t/cuda-loops-case-study-code-generation-vs-templates/302
template <typename arg1_t, typename arg2_t = arg1_t, typename return_t = arg1_t, typename func_t>
void opmath_gpu_kernel_with_scalars(at::TensorIteratorBase& iter, const func_t& f) {
  assert(iter.ntensors() == 3);

  using traits = function_traits<func_t>;
  using opmath_arg1_t = typename traits::template arg<0>::type;
  using opmath_arg2_t = typename traits::template arg<1>::type;
  static_assert(
      traits::arity == 2,
      "gpu_kernel_with_scalars only supports two input arguments");

  if (iter.is_cpu_scalar(1)) {
    // AUnaryFunctor<arg1_t, arg2_t, return_t, func_t> af(f, iter.scalar_value<opmath_arg1_t>(1));
    // iter.remove_operand(1);
    // // TODO: When all kernels that use gpu_kernel_with_scalars are
    // // ported to structured, this device guard can be deleted.  This
    // // works around incorrect device guard generation for pre-structured
    // // kernels device guards, but structured kernels do it right and
    // // we can assume the device is already set correctly
    // const at::OptionalDeviceGuard device_guard(iter.device(1));
    // gpu_kernel(iter, af);
  } else if (iter.is_cpu_scalar(2)) {
    // BUnaryFunctor<arg1_t, arg2_t, return_t, func_t> bf(f, iter.scalar_value<opmath_arg2_t>(2));
    // iter.remove_operand(2);
    // gpu_kernel(iter, bf);
  } else { // Only here is support now
    gpu_kernel(iter, BinaryFunctor<arg1_t, arg2_t, return_t, func_t>(f));
  }
}

// Legacy variant that assumes that func_t has the correct types
// that we expect to load from memory
template <typename func_t>
void gpu_kernel_with_scalars(at::TensorIteratorBase& iter, const func_t& f) {
  using traits = function_traits<func_t>;
  static_assert(
      traits::arity == 2,
      "gpu_kernel_with_scalars only supports two input arguments");
  using arg1_t = typename traits::template arg<0>::type;
  using arg2_t = typename traits::template arg<1>::type;
  using return_t = typename traits::result_type;
  opmath_gpu_kernel_with_scalars<arg1_t, arg2_t, return_t, func_t>(iter, f);
}

namespace { // functions for `gpu_kernel_multiple_outputs`.

template <typename T> struct is_tuple: std::false_type {};

template <typename ...T> struct is_tuple<std::tuple<T...>>: std::true_type {};

template <int num_outputs, typename func_t, typename array_t, typename inp_calc_t, typename out_calc_t>
// C10_LAUNCH_BOUNDS_1(num_threads())
struct unrolled_elementwise_kernel_for_multi_outputs {
  unrolled_elementwise_kernel_for_multi_outputs(int N, func_t f, array_t data, inp_calc_t ic, out_calc_t oc)
  : N(N), f(f), data(data), ic(ic), oc(oc) {}

  void operator () (sycl::nd_item<1> pos) const {
    int remaining = N - group_work_size() * pos.get_local_id();
    elementwise_kernel_helper(
        pos, f,
        memory::policies::multi_outputs_unroll<array_t, inp_calc_t, out_calc_t, num_outputs>(
          data, remaining, ic, oc));
  }
private:
  int N;
  func_t f;
  array_t data;
  inp_calc_t ic;
  out_calc_t oc;
};

template <int num_outputs, typename func_t, typename array_t, typename inp_calc_t, typename out_calc_t>
static inline void launch_unrolled_kernel_for_multi_outputs(int64_t N, const func_t& f, array_t data, inp_calc_t ic, out_calc_t oc) {
  assert(N > 0 && N <= std::numeric_limits<int32_t>::max());
  int64_t grid = (N + group_work_size() - 1) / group_work_size();
  auto queue = currentQueue();
  queue.submit([&](sycl::handler &cgh) {
    cgh.parallel_for(sycl::nd_range<1> ({grid * group_size(), group_size()}),
      unrolled_elementwise_kernel_for_multi_outputs<num_outputs, func_t, array_t, inp_calc_t, out_calc_t>(N, f, data, ic, oc));
  });
  // C10_CUDA_KERNEL_LAUNCH_CHECK();
}

template <typename func_t>
void gpu_kernel_multiple_outputs_impl(at::TensorIteratorBase& iter, const func_t& f) {
  using traits = function_traits<func_t>;
  using output_t = typename traits::result_type;
  static_assert(is_tuple<output_t>::value, "f's return type must be `std::tuple`");
  constexpr int num_outputs = std::tuple_size<output_t>::value;
  constexpr int num_inputs = traits::arity;
  constexpr int ntensors = num_outputs + num_inputs;

  assert(iter.can_use_32bit_indexing());
  assert(iter.ntensors() == ntensors);

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

  int64_t numel = iter.numel();

  if (iter.is_contiguous()) {
    auto input_calc = TrivialOffsetCalculator<num_inputs>();
    auto output_calc = TrivialOffsetCalculator<num_outputs>();
    launch_unrolled_kernel_for_multi_outputs<num_outputs>(numel, f, data, input_calc, output_calc);
  } else {
    auto input_calc = make_input_offset_calculator<num_inputs>(iter);
    auto output_calc = make_output_offset_calculator<num_outputs>(iter);
    launch_unrolled_kernel_for_multi_outputs<num_outputs>(numel, f, data, input_calc, output_calc);
  }
}
} // namespace

template <typename func_t>
void gpu_kernel_multiple_outputs(at::TensorIteratorBase& iter, const func_t& f) {
  // ASSERT_HOST_DEVICE_LAMBDA(func_t);

  if (iter.numel() == 0) {
    return;
  }

  if (!iter.can_use_32bit_indexing()) {
    // for (auto& sub_iter : iter.with_32bit_indexing()) {
    //   gpu_kernel_multiple_outputs(sub_iter, f);
    // }
    return; // Never be here
  }

  gpu_kernel_multiple_outputs_impl(iter, f);
}

} //namespace at::native