#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
/*
 * Copyright 2020-2026 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 */

#pragma once

#include <ATen/core/Array.h>
#include <ATen/detail/FunctionTraits.h>
#include <c10/core/DynamicCast.h>
#include <c10/macros/Macros.h>
#include <c10/util/Exception.h>
#include <c10/util/TypeCast.h>
#include <cstdint>
#include <type_traits>

#include <ATen/native/xpu/sycl/MemoryAccessUtils.h>
#include <ATen/native/xpu/sycl/OffsetCalculator.h>
#include <comm/SYCLContext.h>

namespace at {
namespace native {
namespace memory {
namespace policies {

// Assumption:
// all tensors are contiguous, that is: stride == sizeof(type) for all tensors
template <
    int item_work_size,
    typename data_t,
    typename inp_calc_t,
    typename out_calc_t,
    typename loader_t,
    typename storer_t,
    int num_outputs = 1>
struct unroll {
  data_t data;
  int remaining;
  inp_calc_t input_offset_calculator;
  out_calc_t output_offset_calculator;
  loader_t loader;
  storer_t storer;
  int item_idx;
  int group_idx;
  int num_items_per_group;
  int group_work_size;

  unroll(
      data_t data,
      int remaining,
      inp_calc_t ic,
      out_calc_t oc,
      loader_t l,
      storer_t s,
      int item_idx,
      int group_idx,
      int num_items_per_group)
      : data(data),
        remaining(remaining),
        input_offset_calculator(ic),
        output_offset_calculator(oc),
        loader(l),
        storer(s),
        item_idx(item_idx),
        group_idx(group_idx),
        num_items_per_group(num_items_per_group),
        group_work_size(item_work_size * num_items_per_group) {}

  inline bool check_inbounds(int item_work_elem) const {
    return (item_idx + item_work_elem * num_items_per_group < remaining);
  }

  template <typename args_t>
  inline void load(args_t* args) {
    constexpr int arity = std::tuple_size<args_t>::value;
    int item_idx_ = item_idx;
#pragma unroll
    for (int i = 0; i < item_work_size; i++) {
      if (item_idx_ >= remaining) {
        return;
      }
      int linear_idx = item_idx_ + group_work_size * group_idx;
      auto offset = input_offset_calculator.get(linear_idx);
      detail::static_unroll<detail::unroll_load_helper, arity>::with_args(
          *this, args, offset, loader, i, num_outputs);
      item_idx_ += num_items_per_group;
    }
  }

  template <typename scalar_t>
  inline void store(scalar_t* from) {
    int item_idx_ = item_idx;
#pragma unroll
    for (int i = 0; i < item_work_size; i++) {
      if (item_idx_ >= remaining) {
        return;
      }
      int linear_idx = item_idx_ + group_work_size * group_idx;
      int offset = output_offset_calculator.get(linear_idx)[0];
      storer.store(from[i], data[0], offset);
      item_idx_ += num_items_per_group;
    }
  }
};

// Assumption:
// 1. tensors could be contiguous, that is: stride == sizeof(type).
// 2. tensors could be broadcasted.
template <int vec_size, typename data_t, typename inp_calc_t>
struct vectorized {
  data_t data;
  inp_calc_t input_offset_calculator;
  int item_idx;
  int group_idx;
  int num_items_per_group;
  int group_work_size;

  vectorized(
      data_t data,
      inp_calc_t ic,
      int item_idx,
      int group_idx,
      int num_items_per_group)
      : data(data),
        input_offset_calculator(ic),
        item_idx(item_idx),
        group_idx(group_idx),
        num_items_per_group(num_items_per_group),
        group_work_size(vec_size * num_items_per_group) {}

  inline constexpr bool check_inbounds(int item_work_elem) const {
    return true;
  }

  template <typename accessor_t, typename scalar_t>
  inline void load_single_arg(accessor_t to, scalar_t* from) {
    auto v = load_vector<vec_size>(from, 0);
#pragma unroll
    for (int j = 0; j < vec_size; j++) {
      to(j) = v.val[j];
    }
  }

  template <typename args_t>
  inline void load(args_t* args) {
    constexpr int arity = std::tuple_size<args_t>::value;
    int group_offset = group_work_size * group_idx;
    // `Unroll` policy cannot feed memory bandwidth well on Intel GPU,
    // 1. Small loop size cannot provide enough payloads, specially for small
    //    size data type s8/u8/f16/bf16.
    // 2. Large loop size leads to register pressure and register spill.
    // Apply `Vectorized` policy for some 'non-contiguous' cases,
    // like broadcast. The broadcasted operands are dense and could be
    // optimized by the policy if satisfying vectorization conditions.
    // auto offset = input_offset_calculator.get(group_offset);
    auto linear_idx = group_offset + item_idx * vec_size;
    auto offset = input_offset_calculator.get(linear_idx);
    detail::static_unroll<detail::vectorized_load_helper, arity>::with_args(
        *this, args, offset);
  }

  template <typename scalar_t>
  inline void store(scalar_t* from) {
    using vec_t = aligned_vector<scalar_t, vec_size>;
    scalar_t* to =
        reinterpret_cast<scalar_t*>(data[0]) + group_work_size * group_idx;
    vec_t* to_ = reinterpret_cast<vec_t*>(to);
    vec_t v;
#pragma unroll
    for (int j = 0; j < vec_size; j++) {
      v.val[j] = from[j];
    }
    to_[item_idx] = v;
  }
};

template <
    int item_work_size,
    typename data_t,
    typename inp_calc_t,
    typename out_calc_t,
    int num_outputs>
struct multi_outputs_unroll {
  data_t data;
  int remaining;
  inp_calc_t input_offset_calculator;
  out_calc_t output_offset_calculator;
  LoadWithoutCast loader;
  StoreWithoutCast storer;
  int item_idx;
  int group_idx;
  int num_items_per_group;
  int group_work_size;

  multi_outputs_unroll(
      data_t data,
      int remaining,
      inp_calc_t ic,
      out_calc_t oc,
      int item_idx,
      int group_idx,
      int num_items_per_group)
      : data(data),
        remaining(remaining),
        input_offset_calculator(ic),
        output_offset_calculator(oc),
        item_idx(item_idx),
        group_idx(group_idx),
        num_items_per_group(num_items_per_group),
        group_work_size(item_work_size * num_items_per_group) {}

  inline bool check_inbounds(int item_work_elem) const {
    return (item_idx + item_work_elem * num_items_per_group < remaining);
  }

  template <typename args_t>
  inline void load(args_t* args) {
    constexpr int arity = std::tuple_size<args_t>::value;
    int item_idx_ = item_idx;
#pragma unroll
    for (int i = 0; i < item_work_size; i++) {
      if (item_idx_ >= remaining) {
        return;
      }
      int linear_idx = item_idx_ + group_work_size * group_idx;
      auto offset = input_offset_calculator.get(linear_idx);
      detail::static_unroll<detail::unroll_load_helper, arity>::with_args(
          *this, args, offset, loader, i, num_outputs);
      item_idx_ += num_items_per_group;
    }
  }

  template <typename return_t>
  inline void store(return_t* from) {
    int item_idx_ = item_idx;
#pragma unroll
    for (int i = 0; i < item_work_size; i++) {
      if (item_idx_ >= this->remaining) {
        return;
      }
      int linear_idx = item_idx_ + group_work_size * group_idx;
      auto offsets = this->output_offset_calculator.get(linear_idx);
      detail::static_unroll<detail::multi_outputs_store_helper, num_outputs>::
          with_args(this->data, offsets, from[i]);
      item_idx_ += num_items_per_group;
    }
  }
};

} // namespace policies

static inline int preferred_vector_width(at::DeviceIndex dev_id, int elem_sz) {
  size_t ret;
  switch (elem_sz) {
    case 1:
      static_assert(sizeof(char) == 1, "the char size is not 1 bytes");
      ret = syclPrefVectorWidth<char>(dev_id);
      break;
    case 2:
      static_assert(sizeof(short) == 2, "the short size is not 2 bytes");
      ret = syclPrefVectorWidth<short>(dev_id);
      break;
    case 4:
      ret = syclPrefVectorWidth<int>(dev_id);
      static_assert(sizeof(int) == 4, "the long size is not 4 bytes");
      break;
    case 8:
      static_assert(sizeof(int64_t) == 8, "the long size is not 8");
      ret = syclPrefVectorWidth<int64_t>(dev_id);
      break;
    default:
      // no vectorize
      ret = 1;
  }
  return ret;
}

// Query vector size for specific `scalar_t` on current device.
// Enlarge payloads on each work item with a definite data type
// specific vector size preference queried from SYCL runtime to feed
// memory bandwidth well on the specified hardward platform.
// In `Vectorized` policy, using definite vector size for some data type could
// get enough payloads without outer loop. Outer loop may bring additional
// instructions and potential registers usage.
template <typename scalar_t>
inline int can_vectorize_up_to(const char* pointer) {
  int elem_size = sizeof(scalar_t);
  at::DeviceIndex dev_id = c10::xpu::current_device();
  int preferred_width = preferred_vector_width(dev_id, elem_size);
  uint64_t address = reinterpret_cast<uint64_t>(pointer);
  constexpr int vec2_alignment =
      std::alignment_of<aligned_vector<scalar_t, 2>>::value;
  constexpr int vec4_alignment =
      std::alignment_of<aligned_vector<scalar_t, 4>>::value;
  constexpr int vec8_alignment =
      std::alignment_of<aligned_vector<scalar_t, 8>>::value;
  constexpr int vec16_alignment =
      std::alignment_of<aligned_vector<scalar_t, 16>>::value;
  if (address % vec16_alignment == 0) {
    return std::min<int>(preferred_width, 16);
  } else if (address % vec8_alignment == 0) {
    return std::min<int>(preferred_width, 8);
  } else if (address % vec4_alignment == 0) {
    return std::min<int>(preferred_width, 4);
  } else if (address % vec2_alignment == 0) {
    return std::min<int>(preferred_width, 2);
  }
  return 1;
}

template <typename scalar_t>
inline int can_vectorize_up_to(char* pointer) {
  return can_vectorize_up_to<scalar_t>(static_cast<const char*>(pointer));
}

template <int i>
struct can_vectorize_up_to_helper {
  template <typename array_t, typename traits>
  static void apply(int& result, array_t pointers, traits _) {
    using arg_t = typename traits::template arg<i>::type;
    result = std::min<int>(result, can_vectorize_up_to<arg_t>(pointers[i + 1]));
  }
};

template <typename func_t, typename array_t>
inline int can_vectorize_up_to(array_t pointers) {
  using traits = function_traits<func_t>;
  using return_t = typename traits::result_type;
  constexpr int arity = traits::arity;
  int result = can_vectorize_up_to<return_t>(pointers[0]);
  detail::static_unroll<can_vectorize_up_to_helper, arity>::with_args(
      result, pointers, traits());
  return result;
}

template <typename T>
struct alignas(sizeof(T) * 4) m_vec4 {
  union {
    T val[4];
    struct {
      T x, y, z, w;
    };
  };
  m_vec4() = default;
  m_vec4(T x_, T y_, T z_, T w_) : x(x_), y(y_), z(z_), w(w_) {}
};
typedef m_vec4<uint32_t> uint4;

template <typename T>
inline size_t get_alignment(T ptr_or_size) {
  auto val = reinterpret_cast<uintptr_t>(ptr_or_size);
  if (val % 16 == 0) {
    return 16;
  } else if (val % 8 == 0) {
    return 8;
  } else if (val % 4 == 0) {
    return 4;
  } else if (val % 2 == 0) {
    return 2;
  } else {
    return 1;
  }
}

template <>
inline size_t get_alignment<size_t>(size_t size) {
  return get_alignment(reinterpret_cast<void*>(size));
}

template <bool Value, class... Args>
inline constexpr bool dependent_bool_value = Value;

template <class... Args>
inline constexpr bool dependent_false = dependent_bool_value<false, Args...>;

template <int Size>
union Vec;

template <>
union Vec<4> {
  uint16_t u16[2];
  uint32_t u32, as_scalar;
  float f32;
};

template <>
union Vec<8> {
  uint16_t u16[4];
  uint32_t u32[2];
  uint64_t u64, as_scalar;
  float f32[2];
};

template <>
union alignas(16) Vec<16> {
  uint16_t u16[8];
  uint32_t u32[4];
  uint64_t u64[2];
  uint4 u128, as_scalar;
  float f32[4];
};

template <int Alignment, typename T>
inline Vec<Alignment> ld_vec(const T* addr) {
  Vec<Alignment> vec;
  if constexpr (Alignment == 16) {
    vec.u128 = *reinterpret_cast<const uint4*>(addr);
  } else if constexpr (Alignment == 8) {
    vec.u64 = *reinterpret_cast<const uint64_t*>(addr);
  } else if constexpr (Alignment == 4) {
    vec.u32 = *reinterpret_cast<const uint32_t*>(addr);
  } else {
    static_assert(dependent_false<T>);
  }
  return vec;
}

template <int Alignment, typename T>
inline void st_vec(T* addr, const Vec<Alignment>& vec) {
  if constexpr (Alignment == 16) {
    reinterpret_cast<uint64_t*>(addr)[0] = vec.u64[0];
    reinterpret_cast<uint64_t*>(addr)[1] = vec.u64[1];
  } else if constexpr (Alignment == 8) {
    *reinterpret_cast<uint64_t*>(addr) = vec.u64;
  } else if constexpr (Alignment == 4) {
    *reinterpret_cast<uint32_t*>(addr) = vec.u32;
  } else {
    static_assert(dependent_false<T>);
  }
}

} // namespace memory
} // namespace native
} // namespace at

#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif  // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
