/*
Copyright 2020 The OneFlow Authors. All rights reserved.

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

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include "hip/hip_runtime.h"
#include "oneflow/core/framework/framework.h"
#include "oneflow/core/ep/rocm/cuda_stream.h"
#include "oneflow/core/ep/include/primitive/copy_nd.h"
#include "oneflow/core/ep/include/primitive/batch_matmul.h"
#include "oneflow/core/kernel/cuda_graph_support.h"

namespace oneflow {

namespace {

__global__ void GenerateGatherIndicesGpu(const int32_t elem_cnt, const int32_t stride,
                                         const int32_t in_cols, const int32_t offset,
                                         int32_t* gather_indices) {
  CUDA_1D_KERNEL_LOOP(i, elem_cnt) {
    const int32_t row = i / stride;
    const int32_t col = i - row * stride;
    if (col < row + offset) {
      int32_t in_index = row * in_cols + col;
      int32_t idx = row * (offset + row - 1 + offset) / 2 + col;
      gather_indices[idx] = in_index;
    }
  }
}

template<typename T>
__global__ void GatherConcatGpu(int32_t elem_cnt, int32_t out_cols, int32_t valid_out_cols,
                                int32_t in_cols, int32_t output_concat_end_dim,
                                const int32_t* gather_indices, const T* in,
                                const T* output_concat_ptr, T* out_ptr) {
  CUDA_1D_KERNEL_LOOP(i, elem_cnt) {
    const int32_t row = i / out_cols;
    const int32_t col = i - row * out_cols;
    T out_val;
    if (col < output_concat_end_dim) {
      const int32_t output_concat_idx = row * output_concat_end_dim + col;
      out_val = output_concat_ptr[output_concat_idx];
    } else if (col < valid_out_cols) {
      const int32_t gather_col_idx = gather_indices[col - output_concat_end_dim];
      const int32_t in_offset = row * in_cols + gather_col_idx;
      out_val = in[in_offset];
    } else {
      out_val = 0;
    }
    out_ptr[i] = out_val;
  }
}

template<typename T>
__global__ void ScatterSplitAddTransposeGpu(int32_t elem_cnt, int32_t stride_dim, int32_t out_dim,
                                            int32_t in_grad_stride, int32_t in_grad_matrix_dim,
                                            int32_t in_grad_matrix_valid_dim,
                                            int32_t output_concat_end_dim, const int32_t offset,
                                            const T* dy, T* output_concat_grad, T* in_grad) {
  CUDA_1D_KERNEL_LOOP(i, elem_cnt) {
    const int32_t row = i / stride_dim;
    const int32_t col = i - row * stride_dim;
    if (col < output_concat_end_dim) {
      output_concat_grad[row * output_concat_end_dim + col] = dy[row * out_dim + col];
    } else {
      int32_t in_col_id = col - output_concat_end_dim;
      const int32_t matrix_row = in_col_id / in_grad_matrix_dim;
      const int32_t matrix_col = in_col_id - matrix_row * in_grad_matrix_dim;
      T grad_val = 0;
      const T* row_dy = dy + row * out_dim + output_concat_end_dim;
      if (matrix_row < in_grad_matrix_valid_dim && matrix_col < in_grad_matrix_valid_dim) {
        if (matrix_col < matrix_row) {
          int32_t dy_col_idx = matrix_row * (offset + matrix_row - 1 + offset) / 2 + matrix_col;
          grad_val = row_dy[dy_col_idx];
        } else if (matrix_row < matrix_col) {
          // transpose add
          int32_t trans_row_id = matrix_col;
          int32_t trans_col_id = matrix_row;
          int32_t dy_col_idx =
              trans_row_id * (offset + trans_row_id - 1 + offset) / 2 + trans_col_id;
          grad_val = row_dy[dy_col_idx];
        } else if ((matrix_row == matrix_col) && (offset == 1)) {
          int32_t dy_col_idx = matrix_row * (offset + matrix_row - 1 + offset) / 2 + matrix_col;
          grad_val = row_dy[dy_col_idx] * static_cast<T>(2);
        }
      }
      int32_t in_grad_offset = row * in_grad_stride + in_col_id;
      in_grad[in_grad_offset] = grad_val;
    }
  }
}

template<typename T>
void ConcatFeatures(user_op::KernelComputeContext* ctx, int64_t dst_rows, int64_t dst_cols,
                    void* dst_ptr) {
  const int64_t feature_input_size = ctx->input_size("features");
  auto primitive = ep::primitive::NewPrimitive<ep::primitive::CopyNdFactory>(DeviceType::kCUDA, 2);
  DimVector dst_shape = {dst_rows, dst_cols};
  int64_t out_col_offset = 0;
  for (int64_t i = 0; i < feature_input_size; ++i) {
    const user_op::Tensor* feature = ctx->Tensor4ArgNameAndIndex("features", i);
    const int64_t feature_rows = feature->shape_view().At(0);
    const int64_t feature_cols = feature->shape_view().Count(1);
    DimVector dst_pos_vec = {0, out_col_offset};
    DimVector src_shape = {feature_rows, feature_cols};
    DimVector src_pos_vec = {0, 0};
    DimVector extent_vec = {feature_rows, feature_cols};
    primitive->Launch(ctx->stream(), feature->data_type(), 2, dst_ptr, dst_shape.data(),
                      dst_pos_vec.data(), feature->dptr<T>(), src_shape.data(), src_pos_vec.data(),
                      extent_vec.data());
    out_col_offset += feature_cols;
  }
  int64_t pad_dim = dst_cols - out_col_offset;
  if (pad_dim > 0) {
    char* out_ptr = reinterpret_cast<char*>(dst_ptr) + out_col_offset * sizeof(T);
    OF_CUDA_CHECK(hipMemset2DAsync(out_ptr, dst_cols * sizeof(T), 0, pad_dim * sizeof(T), dst_rows,
                                    ctx->stream()->As<ep::CudaStream>()->cuda_stream()));
  }
}

template<typename T>
void GatherConcatKernel(ep::Stream* stream, int32_t elem_cnt, int32_t out_dim,
                        int32_t valid_out_dim, int32_t features_concated_dim,
                        int32_t concated_padded_dim, int32_t output_concat_end_dim,
                        bool self_interaction, const T* matmul_out, const T* output_concat_ptr,
                        int32_t* gather_indices_ptr, T* out_ptr) {
  hipStream_t cuda_stream = stream->As<ep::CudaStream>()->cuda_stream();
  const int32_t gen_indices_elem_cnt = features_concated_dim * features_concated_dim;
  int32_t offset = self_interaction ? 1 : 0;
  hipLaunchKernelGGL(GenerateGatherIndicesGpu, BlocksNum4ThreadsNum(gen_indices_elem_cnt), kCudaThreadsNumPerBlock, 0, cuda_stream, gen_indices_elem_cnt, features_concated_dim,
                                            concated_padded_dim, offset, gather_indices_ptr);

  int32_t matmul_stride = concated_padded_dim * concated_padded_dim;
  hipLaunchKernelGGL(GatherConcatGpu, BlocksNum4ThreadsNum(elem_cnt), kCudaThreadsNumPerBlock, 0, cuda_stream, 
      elem_cnt, out_dim, valid_out_dim, matmul_stride, output_concat_end_dim, gather_indices_ptr,
      matmul_out, output_concat_ptr, out_ptr);
}

template<typename T>
void ScatterSplitAddTranspose(ep::Stream* stream, int32_t batch_size, int32_t out_dim,
                              int32_t concated_padded_dim, int32_t features_concated_dim,
                              int32_t output_concat_end_dim, const bool self_interaction,
                              const T* dy, T* output_concat_grad, T* matmul_out_grad_ptr) {
  int32_t stride_dim = output_concat_end_dim + concated_padded_dim * concated_padded_dim;
  int32_t matmul_stride = concated_padded_dim * concated_padded_dim;
  const int32_t elem_cnt = batch_size * stride_dim;
  int32_t offset = self_interaction ? 1 : 0;
  ScatterSplitAddTransposeGpu<<<BlocksNum4ThreadsNum(elem_cnt), kCudaThreadsNumPerBlock, 0,
                                stream->As<ep::CudaStream>()->cuda_stream()>>>(
      elem_cnt, stride_dim, out_dim, matmul_stride, concated_padded_dim, features_concated_dim,
      output_concat_end_dim, offset, dy, output_concat_grad, matmul_out_grad_ptr);
}

template<typename T>
void ConcatFeaturesGrad(user_op::KernelComputeContext* ctx, const int64_t batch_size,
                        const int64_t concated_padded_dim, const int64_t vector_size,
                        const T* concated_features_grad) {
  auto primitive = ep::primitive::NewPrimitive<ep::primitive::CopyNdFactory>(DeviceType::kCUDA, 2);
  DimVector src_shape = {batch_size, concated_padded_dim * vector_size};
  int64_t in_col_offset = 0;
  for (int64_t i = 0; i < ctx->output_size("features_grad"); ++i) {
    user_op::Tensor* feature_grad = ctx->Tensor4ArgNameAndIndex("features_grad", i);
    const int64_t feature_grad_rows = feature_grad->shape_view().At(0);
    const int64_t feature_grad_cols = feature_grad->shape_view().Count(1);
    DimVector dst_shape = {feature_grad_rows, feature_grad_cols};
    DimVector dst_pos_vec = {0, 0};
    DimVector src_pos_vec = {0, in_col_offset};
    DimVector extent_vec = {feature_grad_rows, feature_grad_cols};
    in_col_offset += feature_grad_cols;
    primitive->Launch(ctx->stream(), feature_grad->data_type(), 2, feature_grad->mut_dptr(),
                      dst_shape.data(), dst_pos_vec.data(), concated_features_grad,
                      src_shape.data(), src_pos_vec.data(), extent_vec.data());
  }
}

template<typename T>
struct DefaultComputeType {
  using type = T;
};

template<>
struct DefaultComputeType<half> {
  using type = float;
};

template<typename T, size_t pack_size>
struct alignas(sizeof(T) * pack_size) Pack {
  T elem[pack_size];
};

int64_t GetPaddedDim(int64_t dim) {
  const int64_t align_dim = 16;
  const int64_t padded_dim = (dim + align_dim - 1) / align_dim * align_dim;
  return padded_dim;
}

template<typename T, int32_t max_in>
struct DotFwdParam {
  const T* in[max_in];
  int32_t in_feature_dim[max_in];
  int32_t dim_start_offset[max_in];
  int32_t features_dim;
  const T* output_concat;
  int32_t output_concat_size;
  T* out;
  int32_t num_in;
};

constexpr int kUnrollDim = 2;
template<typename T, typename ComputeType, int32_t max_in, int32_t pack_size, int mn_tile_dim,
         int k_tile_dim>
__global__ void DotFeatureInteractionWmmaImpl(
    int m_num_tiles, int k_num_tiles, int64_t batch_size, int padded_num_rows, int vector_num_pack,
    int padded_vector_num_pack, int out_num_cols, int out_num_cols_num_pack, int in_shared_mem_cols,
    int in_shared_mem_cols_num_pack, int acc_shared_mem_cols, int acc_shared_mem_cols_num_pack,
    int offset, int output_padding, DotFwdParam<T, max_in> param) {
  asm volatile("s_trap 0;");
}

template<typename T>
struct KTileDim {
  static const int val = 16;
};

template<>
struct KTileDim<float> {
  static const int val = 8;
};

template<typename T, int max_in, int32_t pack_size>
struct DotFeatureInteractionKernel {
  static bool Launch(ep::Stream* stream, int64_t batch_size, int concated_padded_dim,
                     int vector_size, int out_num_cols, bool self_interaction, int output_padding,
                     const DotFwdParam<T, max_in>& param) {
    const int block_size = 128;
    const int block_dim_x = 32;
    const int block_dim_y = block_size / block_dim_x;
    const int num_blocks = batch_size;
    const int mn_tile_dim = 16;
    const int k_tile_dim = KTileDim<T>::val;
    const int64_t padded_vector_size = GetPaddedDim(vector_size);
    const int m_num_tiles = concated_padded_dim / mn_tile_dim;
    const int k_num_tiles = padded_vector_size / k_tile_dim;
    const int skew_in = 8;
    const int skew_acc = 8;
    const int in_shared_mem_num_cols = padded_vector_size + skew_in;
    const int acc_shared_mem_num_cols = concated_padded_dim + skew_acc;
    const size_t in_shared_mem_bytes = concated_padded_dim * in_shared_mem_num_cols * sizeof(T);
    using ComputeType = typename DefaultComputeType<T>::type;
    const size_t acc_shared_mem_bytes =
        concated_padded_dim * acc_shared_mem_num_cols * sizeof(ComputeType);
    const size_t total_shared_mem_bytes = in_shared_mem_bytes + acc_shared_mem_bytes;
    const int32_t offset = self_interaction ? 1 : 0;
    const int out_num_cols_num_pack = out_num_cols / pack_size;
    const int vector_num_pack = vector_size / pack_size;
    const int padded_vector_num_pack = padded_vector_size / pack_size;
    const int in_shared_mem_cols_num_pack = in_shared_mem_num_cols / pack_size;
    const int acc_shared_mem_cols_num_pack = acc_shared_mem_num_cols / pack_size;
    int max_active_blocks;
    OF_CUDA_CHECK(hipOccupancyMaxActiveBlocksPerMultiprocessor(
        &max_active_blocks,
        DotFeatureInteractionWmmaImpl<T, ComputeType, max_in, pack_size, mn_tile_dim, k_tile_dim>,
        block_size, total_shared_mem_bytes));
    if (max_active_blocks <= 0) { return false; }
    hipStream_t cuda_stream = stream->As<ep::CudaStream>()->cuda_stream();
    hipLaunchKernelGGL(HIP_KERNEL_NAME(DotFeatureInteractionWmmaImpl<T, ComputeType, max_in, pack_size, mn_tile_dim, k_tile_dim>), num_blocks, dim3(block_dim_x, block_dim_y), total_shared_mem_bytes, cuda_stream, 
            m_num_tiles, k_num_tiles, batch_size, concated_padded_dim, vector_num_pack,
            padded_vector_num_pack, out_num_cols, out_num_cols_num_pack, in_shared_mem_num_cols,
            in_shared_mem_cols_num_pack, acc_shared_mem_num_cols, acc_shared_mem_cols_num_pack,
            offset, output_padding, param);
    return true;
  }
};

template<typename T, int32_t max_in>
struct DotBwdParam {
  const T* out_grad;
  const T* in[max_in];
  T* in_grad[max_in];
  T* output_concat_grad;
  int32_t output_concat_size;
  int32_t in_feature_dim[max_in];
  int32_t dim_start_offset[max_in];
  int32_t features_dim;
  int32_t num_in;
};

template<typename T, typename ComputeType, int32_t max_in, int32_t pack_size, int mn_tile_dim,
         int k_tile_dim>
__global__ void DotFeatureInteractionBackwardWmmaImpl(
    int m_num_tiles, int n_num_tiles, int k_num_tiles, int64_t batch_size, int padded_num_rows,
    int vector_num_pack, int padded_vector_num_pack, int out_num_cols, int in_shared_mem_cols,
    int in_shared_mem_cols_num_pack, int matrix_out_grad_shared_mem_cols, int offset,
    DotBwdParam<T, max_in> param) {
  asm volatile("s_trap 0;");
}

template<typename T, int max_in, int32_t pack_size>
struct DotFeatureInteractionBackwardKernel {
  static bool Launch(ep::Stream* stream, int64_t batch_size, int concated_padded_dim,
                     int vector_size, int out_num_cols, bool self_interaction,
                     const DotBwdParam<T, max_in>& param) {
    const int block_size = 256;
    const int block_dim_x = 32;
    const int block_dim_y = block_size / block_dim_x;
    const int num_blocks = batch_size;
    const int mn_tile_dim = 16;
    const int k_tile_dim = KTileDim<T>::val;
    const int64_t padded_vector_size = GetPaddedDim(vector_size);
    const int m_num_tiles = concated_padded_dim / mn_tile_dim;
    const int k_num_tiles = concated_padded_dim / k_tile_dim;
    const int n_num_tiles = padded_vector_size / mn_tile_dim;
    const int skew_in = 8;
    const int in_shared_mem_num_cols = padded_vector_size + skew_in;
    const int matrix_out_grad_shared_mem_cols = concated_padded_dim + skew_in;
    const size_t in_shared_mem_bytes = concated_padded_dim * in_shared_mem_num_cols * sizeof(T);
    const size_t matrix_out_grad_shared_mem_bytes =
        concated_padded_dim * matrix_out_grad_shared_mem_cols * sizeof(T);
    using ComputeType = typename DefaultComputeType<T>::type;
    const size_t in_grad_shared_mem_bytes =
        concated_padded_dim * in_shared_mem_num_cols * sizeof(ComputeType);
    const size_t total_shared_mem_bytes =
        in_shared_mem_bytes + matrix_out_grad_shared_mem_bytes + in_grad_shared_mem_bytes;
    const int32_t offset = self_interaction ? 1 : 0;
    const int vector_num_pack = vector_size / pack_size;
    const int padded_vector_num_pack = padded_vector_size / pack_size;
    const int in_shared_mem_cols_num_pack = in_shared_mem_num_cols / pack_size;
    int max_active_blocks;
    OF_CUDA_CHECK(hipOccupancyMaxActiveBlocksPerMultiprocessor(
        &max_active_blocks,
        DotFeatureInteractionBackwardWmmaImpl<T, ComputeType, max_in, pack_size, mn_tile_dim,
                                              k_tile_dim>,
        block_size, total_shared_mem_bytes));
    if (max_active_blocks <= 0) { return false; }
    hipStream_t cuda_stream = stream->As<ep::CudaStream>()->cuda_stream();
    DotFeatureInteractionBackwardWmmaImpl<T, ComputeType, max_in, pack_size, mn_tile_dim,
                                          k_tile_dim>
        <<<num_blocks, dim3(block_dim_x, block_dim_y), total_shared_mem_bytes, cuda_stream>>>(
            m_num_tiles, n_num_tiles, k_num_tiles, batch_size, concated_padded_dim, vector_num_pack,
            padded_vector_num_pack, out_num_cols, in_shared_mem_num_cols,
            in_shared_mem_cols_num_pack, matrix_out_grad_shared_mem_cols, offset, param);

    return true;
  }
};

template<typename T, int max_in>
bool DispatchFeatureInteractionDotPackSize(user_op::KernelComputeContext* ctx,
                                           const int32_t input_size) {
  CHECK_LE(input_size, max_in) << input_size;
  user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);
  const int64_t batch_size = out->shape_view().At(0);
  const int64_t out_num_cols = out->shape_view().At(1);
  const int64_t vector_size = ctx->TensorDesc4ArgNameAndIndex("features", 0)->shape().At(2);
  DotFwdParam<T, max_in> param;
  param.num_in = input_size;
  param.out = out->mut_dptr<T>();
  int64_t features_concated_dim = 0;
  for (int i = 0; i < input_size; ++i) {
    param.in[i] = ctx->Tensor4ArgNameAndIndex("features", i)->dptr<T>();
    param.in_feature_dim[i] = ctx->TensorDesc4ArgNameAndIndex("features", i)->shape().At(1);
    param.dim_start_offset[i] = features_concated_dim;
    features_concated_dim += param.in_feature_dim[i];
  }
  const int64_t concated_padded_dim = GetPaddedDim(features_concated_dim);
  param.features_dim = features_concated_dim;
  if (ctx->has_input("output_concat", 0)) {
    const user_op::Tensor* output_concat = ctx->Tensor4ArgNameAndIndex("output_concat", 0);
    param.output_concat = output_concat->dptr<T>();
    param.output_concat_size = output_concat->shape_view().At(1);
  } else {
    param.output_concat = nullptr;
    param.output_concat_size = 0;
  }
  const bool self_interaction = ctx->Attr<bool>("self_interaction");
  const int32_t output_padding = ctx->Attr<int32_t>("output_padding");
  if (vector_size % 4 == 0 && out_num_cols % 4 == 0) {
    return DotFeatureInteractionKernel<T, max_in, 4>::Launch(
        ctx->stream(), batch_size, concated_padded_dim, vector_size, out_num_cols, self_interaction,
        output_padding, param);
  } else if (vector_size % 2 == 0 && out_num_cols % 2 == 0) {
    return DotFeatureInteractionKernel<T, max_in, 2>::Launch(
        ctx->stream(), batch_size, concated_padded_dim, vector_size, out_num_cols, self_interaction,
        output_padding, param);
  } else {
    return DotFeatureInteractionKernel<T, max_in, 1>::Launch(
        ctx->stream(), batch_size, concated_padded_dim, vector_size, out_num_cols, self_interaction,
        output_padding, param);
  }
}

template<typename T, int max_in>
bool DispatchFeatureInteractionDotBackwardPackSize(user_op::KernelComputeContext* ctx,
                                                   const int32_t input_size) {
  CHECK_LE(input_size, max_in) << input_size;
  user_op::Tensor* dy = ctx->Tensor4ArgNameAndIndex("dy", 0);
  const int64_t batch_size = dy->shape_view().At(0);
  const int64_t out_num_cols = dy->shape_view().At(1);
  const int64_t vector_size = ctx->TensorDesc4ArgNameAndIndex("features", 0)->shape().At(2);
  DotBwdParam<T, max_in> param;
  param.num_in = input_size;
  param.out_grad = dy->dptr<T>();
  int64_t features_concated_dim = 0;
  for (int i = 0; i < input_size; ++i) {
    param.in[i] = ctx->Tensor4ArgNameAndIndex("features", i)->dptr<T>();
    param.in_grad[i] = ctx->Tensor4ArgNameAndIndex("features_grad", i)->mut_dptr<T>();
    param.in_feature_dim[i] = ctx->TensorDesc4ArgNameAndIndex("features", i)->shape().At(1);
    param.dim_start_offset[i] = features_concated_dim;
    features_concated_dim += param.in_feature_dim[i];
  }
  const int64_t concated_padded_dim = GetPaddedDim(features_concated_dim);
  param.features_dim = features_concated_dim;
  if (ctx->has_output("output_concat_grad", 0)) {
    user_op::Tensor* output_concat_grad = ctx->Tensor4ArgNameAndIndex("output_concat_grad", 0);
    param.output_concat_grad = output_concat_grad->mut_dptr<T>();
    param.output_concat_size = output_concat_grad->shape_view().At(1);
  } else {
    param.output_concat_grad = nullptr;
    param.output_concat_size = 0;
  }
  const bool self_interaction = ctx->Attr<bool>("self_interaction");
  if (vector_size % 4 == 0) {
    return DotFeatureInteractionBackwardKernel<T, max_in, 4>::Launch(
        ctx->stream(), batch_size, concated_padded_dim, vector_size, out_num_cols, self_interaction,
        param);
  } else if (vector_size % 2 == 0) {
    return DotFeatureInteractionBackwardKernel<T, max_in, 2>::Launch(
        ctx->stream(), batch_size, concated_padded_dim, vector_size, out_num_cols, self_interaction,
        param);
  } else {
    return DotFeatureInteractionBackwardKernel<T, max_in, 1>::Launch(
        ctx->stream(), batch_size, concated_padded_dim, vector_size, out_num_cols, self_interaction,
        param);
  }
}

template<typename T, int32_t max_in>
struct Param {
  const T* in[max_in];
  int32_t in_feature_dim[max_in];
  T* out;
  int32_t num_in;
};

template<typename T, int32_t max_in, int32_t pack_size>
__global__ void FeatureInteractionSum(int64_t batch_size, int64_t vector_num_pack,
                                      Param<T, max_in> param) {
  using ComputeType = typename DefaultComputeType<T>::type;
  Pack<T, pack_size>* dst_pack = reinterpret_cast<Pack<T, pack_size>*>(param.out);
  for (int batch_idx = blockIdx.x * blockDim.y + threadIdx.y; batch_idx < batch_size;
       batch_idx += gridDim.x * blockDim.y) {
    Pack<T, pack_size>* batch_out = dst_pack + batch_idx * vector_num_pack;
    for (int col_id = threadIdx.x; col_id < vector_num_pack; col_id += blockDim.x) {
      Pack<ComputeType, pack_size> sum;
      Pack<ComputeType, pack_size> square_sum;
#pragma unroll
      for (int k = 0; k < pack_size; ++k) {
        sum.elem[k] = static_cast<ComputeType>(0);
        square_sum.elem[k] = static_cast<ComputeType>(0);
      }
      for (int i = 0; i < max_in; ++i) {
        if (i >= param.num_in) { break; }
        const Pack<T, pack_size>* batch_in =
            reinterpret_cast<const Pack<T, pack_size>*>(param.in[i])
            + batch_idx * param.in_feature_dim[i] * vector_num_pack;
#pragma unroll
        for (int j = 0; j < param.in_feature_dim[i]; ++j) {
          Pack<T, pack_size> val = batch_in[j * vector_num_pack + col_id];
#pragma unroll
          for (int k = 0; k < pack_size; ++k) {
            const ComputeType compute_val = static_cast<ComputeType>(val.elem[k]);
            sum.elem[k] += compute_val;
            square_sum.elem[k] += compute_val * compute_val;
          }
        }
      }
      Pack<T, pack_size> out;
#pragma unroll
      for (int k = 0; k < pack_size; ++k) {
        out.elem[k] = static_cast<T>((sum.elem[k] * sum.elem[k] - square_sum.elem[k])
                                     * static_cast<ComputeType>(0.5));
      }
      batch_out[col_id] = out;
    }
  }
}

template<typename T, int32_t max_in>
struct GradParam {
  const T* out_grad;
  const T* in[max_in];
  int32_t in_feature_dim[max_in];
  T* in_grad[max_in];
  int32_t num_in;
};

template<typename T, int32_t max_in>
__global__ void FeatureInteractionSumGrad(int64_t batch_size, int64_t vector_size,
                                          GradParam<T, max_in> param) {
  using ComputeType = typename DefaultComputeType<T>::type;
  for (int batch_idx = blockIdx.x * blockDim.y + threadIdx.y; batch_idx < batch_size;
       batch_idx += gridDim.x * blockDim.y) {
    const T* batch_out_grad = param.out_grad + batch_idx * vector_size;
    for (int col_id = threadIdx.x; col_id < vector_size; col_id += blockDim.x) {
      ComputeType sum = 0;
      for (int i = 0; i < max_in; ++i) {
        if (i >= param.num_in) { break; }
        const T* batch_in = param.in[i] + batch_idx * param.in_feature_dim[i] * vector_size;
        for (int j = 0; j < param.in_feature_dim[i]; ++j) {
          sum += static_cast<ComputeType>(batch_in[j * vector_size + col_id]);
        }
      }
      for (int i = 0; i < max_in; ++i) {
        if (i >= param.num_in) { break; }
        const int64_t in_batch_offset = batch_idx * param.in_feature_dim[i] * vector_size;
        const T* batch_in = param.in[i] + in_batch_offset;
        T* batch_in_grad = param.in_grad[i] + in_batch_offset;
        for (int j = 0; j < param.in_feature_dim[i]; ++j) {
          const int64_t offset = j * vector_size + col_id;
          batch_in_grad[offset] =
              static_cast<T>(static_cast<ComputeType>(batch_out_grad[col_id])
                             * (sum - static_cast<ComputeType>(batch_in[offset])));
        }
      }
    }
  }
}

void GetBlockDims(const int64_t vector_size, int* block_dim_x, int* block_dim_y) {
  const int block_size = 256;
  if (vector_size < block_size) {
    *block_dim_x = std::ceil(static_cast<float>(vector_size) / 8) * 8;
    *block_dim_y = (block_size + *block_dim_x - 1) / *block_dim_x;
  } else {
    *block_dim_x = block_size;
    *block_dim_y = 1;
  }
}

int GetNumBlocks(const int64_t num_instances, const int64_t instance_per_block) {
  int max_blocks = (num_instances + instance_per_block - 1) / instance_per_block;
  return std::min(max_blocks, kCudaMaxBlocksNum);
}

template<typename T, int32_t max_in>
void DispatchFeatureInteractionSumPackSize(ep::Stream* stream, const int64_t batch_size,
                                           const int64_t vector_size,
                                           const Param<T, max_in>& param) {
  int block_dim_x;
  int block_dim_y;
  const int pack_size = (vector_size % 2 == 0) ? 2 : 1;
  const int64_t vector_num_pack = vector_size / pack_size;
  GetBlockDims(vector_num_pack, &block_dim_x, &block_dim_y);
  const int num_blocks = GetNumBlocks(batch_size, block_dim_y);
  dim3 block_dims = dim3(block_dim_x, block_dim_y);
  hipStream_t cuda_stream = stream->As<ep::CudaStream>()->cuda_stream();
  if (pack_size == 2) {
    hipLaunchKernelGGL(HIP_KERNEL_NAME(FeatureInteractionSum<T, max_in, 2>), num_blocks, block_dims, 0, cuda_stream, batch_size, vector_num_pack, param);
  } else {
    hipLaunchKernelGGL(HIP_KERNEL_NAME(FeatureInteractionSum<T, max_in, 1>), num_blocks, block_dims, 0, cuda_stream, batch_size, vector_num_pack, param);
  }
}

template<typename T, int max_in>
void DispatchFeatureInteractionSumInputSize(user_op::KernelComputeContext* ctx,
                                            const int32_t input_size) {
  CHECK_LE(input_size, max_in) << input_size;
  user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);
  const int64_t batch_size = out->shape_view().At(0);
  const int64_t vector_size = out->shape_view().At(1);
  Param<T, max_in> param;
  param.num_in = input_size;
  param.out = out->mut_dptr<T>();
  for (int i = 0; i < input_size; ++i) {
    param.in[i] = ctx->Tensor4ArgNameAndIndex("features", i)->dptr<T>();
    param.in_feature_dim[i] = ctx->TensorDesc4ArgNameAndIndex("features", i)->shape().At(1);
  }
  DispatchFeatureInteractionSumPackSize<T, max_in>(ctx->stream(), batch_size, vector_size, param);
}

template<typename T, int max_in>
void DispatchFeatureInteractionSumGradInputSize(user_op::KernelComputeContext* ctx,
                                                const int32_t input_size) {
  CHECK_LE(input_size, max_in) << input_size;
  const user_op::Tensor* dy = ctx->Tensor4ArgNameAndIndex("dy", 0);
  const int64_t batch_size = dy->shape_view().At(0);
  const int64_t vector_size = dy->shape_view().At(1);
  int block_dim_x;
  int block_dim_y;
  GetBlockDims(vector_size, &block_dim_x, &block_dim_y);
  const int num_blocks = GetNumBlocks(batch_size, block_dim_y);
  dim3 block_dims = dim3(block_dim_x, block_dim_y);
  GradParam<T, max_in> param;
  param.num_in = input_size;
  param.out_grad = dy->dptr<T>();
  for (int i = 0; i < input_size; ++i) {
    param.in[i] = ctx->Tensor4ArgNameAndIndex("features", i)->dptr<T>();
    param.in_grad[i] = ctx->Tensor4ArgNameAndIndex("features_grad", i)->mut_dptr<T>();
    param.in_feature_dim[i] = ctx->TensorDesc4ArgNameAndIndex("features_grad", i)->shape().At(1);
  }
  FeatureInteractionSumGrad<T, max_in>
      <<<num_blocks, block_dims, 0, ctx->stream()->As<ep::CudaStream>()->cuda_stream()>>>(
          batch_size, vector_size, param);
}

}  // namespace

template<typename T>
class FusedDotFeatureInteractionPoolingSumKernel final : public user_op::OpKernel,
                                                         public user_op::CudaGraphSupport {
 public:
  FusedDotFeatureInteractionPoolingSumKernel() = default;
  ~FusedDotFeatureInteractionPoolingSumKernel() override = default;

 private:
  using user_op::OpKernel::Compute;
  void Compute(user_op::KernelComputeContext* ctx) const override {
    const int input_size = ctx->input_size("features");
    if (input_size == 1) {
      DispatchFeatureInteractionSumInputSize<T, 1>(ctx, input_size);
    } else if (input_size == 2) {
      DispatchFeatureInteractionSumInputSize<T, 2>(ctx, input_size);
    } else if (input_size <= 8) {
      DispatchFeatureInteractionSumInputSize<T, 8>(ctx, input_size);
    } else {
      CHECK_LE(input_size, 128) << "input_size must not greater than 128. ";
      DispatchFeatureInteractionSumInputSize<T, 128>(ctx, input_size);
    }
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};

#define REGISTER_FUSED_DOT_FEATURE_INTERACTION_POOLING_SUM_KERNEL(dtype)                \
  REGISTER_USER_KERNEL("fused_dot_feature_interaction")                                 \
      .SetCreateFn<FusedDotFeatureInteractionPoolingSumKernel<dtype>>()                 \
      .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA)                  \
                       && (user_op::HobDataType("out", 0) == GetDataType<dtype>::value) \
                       && (user_op::HobAttr<std::string>("pooling") == "sum"));

REGISTER_FUSED_DOT_FEATURE_INTERACTION_POOLING_SUM_KERNEL(float)
REGISTER_FUSED_DOT_FEATURE_INTERACTION_POOLING_SUM_KERNEL(half)

template<typename T>
bool TryLaunchTensorCoreDotKernel(user_op::KernelComputeContext* ctx) {
  const int input_size = ctx->input_size("features");
  if (input_size == 1) {
    return DispatchFeatureInteractionDotPackSize<T, 1>(ctx, input_size);
  } else if (input_size == 2) {
    return DispatchFeatureInteractionDotPackSize<T, 2>(ctx, input_size);
  } else if (input_size <= 8) {
    return DispatchFeatureInteractionDotPackSize<T, 8>(ctx, input_size);
  } else {
    CHECK_LE(input_size, 128) << "input_size must not greater than 128. ";
    return DispatchFeatureInteractionDotPackSize<T, 128>(ctx, input_size);
  }
}

template<typename T>
bool TryLaunchTensorCoreDotBackwardKernel(user_op::KernelComputeContext* ctx) {
  const int input_size = ctx->input_size("features");
  if (input_size == 1) {
    return DispatchFeatureInteractionDotBackwardPackSize<T, 1>(ctx, input_size);
  } else if (input_size == 2) {
    return DispatchFeatureInteractionDotBackwardPackSize<T, 2>(ctx, input_size);
  } else if (input_size <= 8) {
    return DispatchFeatureInteractionDotBackwardPackSize<T, 8>(ctx, input_size);
  } else {
    CHECK_LE(input_size, 128) << "input_size must not greater than 128. ";
    return DispatchFeatureInteractionDotBackwardPackSize<T, 128>(ctx, input_size);
  }
}
template<typename T>
class FusedDotFeatureInteractionKernel final : public user_op::OpKernel,
                                               public user_op::CudaGraphSupport {
 public:
  FusedDotFeatureInteractionKernel() = default;
  ~FusedDotFeatureInteractionKernel() override = default;

 private:
  using user_op::OpKernel::Compute;
  void Compute(user_op::KernelComputeContext* ctx) const override {
    user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);
    const DataType data_type = out->data_type();
    CHECK_LT(out->shape_view().elem_cnt(), GetMaxVal<int32_t>());
    auto* cuda_stream = ctx->stream()->As<ep::CudaStream>();
    // if ((cuda_stream->device_properties().major >= 7 && data_type == DataType::kFloat16)
    //     || (cuda_stream->device_properties().major >= 8 && data_type == DataType::kFloat)) {
    //   bool success = TryLaunchTensorCoreDotKernel<T>(ctx);
    //   if (success == true) { return; }
    // }
    user_op::Tensor* tmp_buffer = ctx->Tensor4ArgNameAndIndex("tmp_buffer", 0);
    const int64_t batch_size = out->shape_view().At(0);
    int64_t features_concated_dim = 0;
    for (int64_t i = 0; i < ctx->input_size("features"); ++i) {
      features_concated_dim += ctx->TensorDesc4ArgNameAndIndex("features", i)->shape().At(1);
    }
    const int64_t concated_padded_dim = GetPaddedDim(features_concated_dim);
    const int64_t vector_size = ctx->TensorDesc4ArgNameAndIndex("features", 0)->shape().At(2);
    const int64_t out_dim = out->shape_view().At(1);
    const int32_t output_padding = ctx->Attr<int32_t>("output_padding");
    const int64_t valid_out_dim = out_dim - output_padding;
    const bool self_interaction = ctx->Attr<bool>("self_interaction");

    T* matmul_out = reinterpret_cast<T*>(tmp_buffer->mut_dptr<char>());
    size_t matmul_out_size =
        GetCudaAlignedSize(batch_size * concated_padded_dim * concated_padded_dim * sizeof(T));
    const int64_t interaction_dim = self_interaction
                                        ? features_concated_dim * (features_concated_dim + 1) / 2
                                        : features_concated_dim * (features_concated_dim - 1) / 2;
    int32_t* gather_indices_ptr =
        reinterpret_cast<int32_t*>(tmp_buffer->mut_dptr<char>() + matmul_out_size);
    size_t gather_indices_size = GetCudaAlignedSize(interaction_dim * sizeof(int32_t));
    T* padded_concated_features_ptr =
        reinterpret_cast<T*>(tmp_buffer->mut_dptr<char>() + matmul_out_size + gather_indices_size);
    size_t padded_concated_features_size =
        GetCudaAlignedSize(batch_size * concated_padded_dim * vector_size * sizeof(T));
    CHECK_GE(tmp_buffer->shape_view().elem_cnt(),
             matmul_out_size + gather_indices_size + padded_concated_features_size);
    ConcatFeatures<T>(ctx, batch_size, concated_padded_dim * vector_size,
                      padded_concated_features_ptr);
    auto batch_matmul = ep::primitive::NewPrimitive<ep::primitive::BatchMatmulFactory>(
        ctx->device_type(), data_type, ep::primitive::BlasTransposeType::N,
        ep::primitive::BlasTransposeType::T);
    batch_matmul->Launch(ctx->stream(), batch_size, concated_padded_dim, concated_padded_dim,
                         vector_size, 1.0, padded_concated_features_ptr,
                         padded_concated_features_ptr, 0.0, matmul_out);

    int64_t output_concat_end_dim = 0;
    const T* output_concat_ptr = nullptr;
    if (ctx->has_input("output_concat", 0)) {
      user_op::Tensor* output_concat = ctx->Tensor4ArgNameAndIndex("output_concat", 0);
      output_concat_end_dim = output_concat->shape_view().At(1);
      output_concat_ptr = output_concat->dptr<T>();
    }
    CHECK_EQ(valid_out_dim, output_concat_end_dim + interaction_dim);
    GatherConcatKernel<T>(ctx->stream(), out->shape_view().elem_cnt(), out_dim, valid_out_dim,
                          features_concated_dim, concated_padded_dim, output_concat_end_dim,
                          self_interaction, matmul_out, output_concat_ptr, gather_indices_ptr,
                          out->mut_dptr<T>());
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};

template<typename T>
user_op::InferTmpSizeFn GenFusedDotFeatureInteractionInferTmpSizeFn() {
  return [](user_op::InferContext* ctx) {
    const Shape& first_feature_shape = ctx->InputShape("features", 0);
    const int64_t batch_size = first_feature_shape.At(0);
    const int64_t vector_size = first_feature_shape.At(2);
    int64_t features_concated_dim = 0;
    for (int32_t i = 0; i < ctx->input_size("features"); ++i) {
      features_concated_dim += ctx->InputShape("features", i).At(1);
    }
    const int64_t concated_padded_dim = GetPaddedDim(features_concated_dim);
    size_t matmul_out_size =
        GetCudaAlignedSize(batch_size * concated_padded_dim * concated_padded_dim * sizeof(T));
    const bool self_interaction = ctx->Attr<bool>("self_interaction");
    const int64_t interaction_dim = self_interaction
                                        ? features_concated_dim * (features_concated_dim + 1) / 2
                                        : features_concated_dim * (features_concated_dim - 1) / 2;
    size_t gather_indices_size = GetCudaAlignedSize(interaction_dim * sizeof(int32_t));
    size_t padded_concated_features_size =
        GetCudaAlignedSize(batch_size * concated_padded_dim * vector_size * sizeof(T));
    return matmul_out_size + gather_indices_size + padded_concated_features_size;
  };
}

#define REGISTER_FUSED_DOT_FEATURE_INTERACTION_KERNEL(dtype)                            \
  REGISTER_USER_KERNEL("fused_dot_feature_interaction")                                 \
      .SetCreateFn<FusedDotFeatureInteractionKernel<dtype>>()                           \
      .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA)                  \
                       && (user_op::HobDataType("out", 0) == GetDataType<dtype>::value) \
                       && (user_op::HobAttr<std::string>("pooling") == "none"))         \
      .SetInferTmpSizeFn(GenFusedDotFeatureInteractionInferTmpSizeFn<dtype>());

REGISTER_FUSED_DOT_FEATURE_INTERACTION_KERNEL(float)
REGISTER_FUSED_DOT_FEATURE_INTERACTION_KERNEL(half)

template<typename T>
class FusedDotFeatureInteractionGradKernel final : public user_op::OpKernel,
                                                   public user_op::CudaGraphSupport {
 public:
  FusedDotFeatureInteractionGradKernel() = default;
  ~FusedDotFeatureInteractionGradKernel() override = default;

 private:
  using user_op::OpKernel::Compute;
  void Compute(user_op::KernelComputeContext* ctx) const override {
    const user_op::Tensor* dy = ctx->Tensor4ArgNameAndIndex("dy", 0);
    user_op::Tensor* tmp_buffer = ctx->Tensor4ArgNameAndIndex("tmp_buffer", 0);
    const DataType data_type = dy->data_type();
    auto* cuda_stream = ctx->stream()->As<ep::CudaStream>();
    // if ((cuda_stream->device_properties().major >= 7 && data_type == DataType::kFloat16)
    //     || (cuda_stream->device_properties().major >= 8 && data_type == DataType::kFloat)) {
    //   bool success = TryLaunchTensorCoreDotBackwardKernel<T>(ctx);
    //   if (success == true) { return; }
    // }
    const int64_t batch_size = dy->shape_view().At(0);
    int64_t features_concated_dim = 0;
    for (int32_t i = 0; i < ctx->output_size("features_grad"); ++i) {
      features_concated_dim += ctx->TensorDesc4ArgNameAndIndex("features_grad", i)->shape().At(1);
    }
    const int64_t concated_padded_dim = GetPaddedDim(features_concated_dim);
    const int64_t vector_size = ctx->TensorDesc4ArgNameAndIndex("features_grad", 0)->shape().At(2);
    const int64_t out_dim = dy->shape_view().At(1);
    const bool self_interaction = ctx->Attr<bool>("self_interaction");
    T* matmul_out_grad_ptr = reinterpret_cast<T*>(tmp_buffer->mut_dptr<char>());
    size_t matmul_out_grad_size =
        GetCudaAlignedSize(batch_size * concated_padded_dim * concated_padded_dim * sizeof(T));
    T* padded_concated_features_grad_ptr =
        reinterpret_cast<T*>(tmp_buffer->mut_dptr<char>() + matmul_out_grad_size);
    size_t padded_concated_features_grad_size =
        GetCudaAlignedSize(batch_size * concated_padded_dim * vector_size * sizeof(T));
    T* padded_concated_features_ptr = reinterpret_cast<T*>(
        tmp_buffer->mut_dptr<char>() + matmul_out_grad_size + padded_concated_features_grad_size);
    size_t padded_concated_features_size = padded_concated_features_grad_size;
    CHECK_LE(
        matmul_out_grad_size + padded_concated_features_grad_size + padded_concated_features_size,
        tmp_buffer->shape_view().elem_cnt());
    ConcatFeatures<T>(ctx, batch_size, concated_padded_dim * vector_size,
                      padded_concated_features_ptr);

    T* output_concat_grad_ptr = nullptr;
    int64_t output_concat_end_dim = 0;
    if (ctx->has_output("output_concat_grad", 0)) {
      user_op::Tensor* output_concat_grad = ctx->Tensor4ArgNameAndIndex("output_concat_grad", 0);
      output_concat_grad_ptr = output_concat_grad->mut_dptr<T>();
      output_concat_end_dim = output_concat_grad->shape_view().At(1);
    }
    ScatterSplitAddTranspose(ctx->stream(), batch_size, out_dim, concated_padded_dim,
                             features_concated_dim, output_concat_end_dim, self_interaction,
                             dy->dptr<T>(), output_concat_grad_ptr, matmul_out_grad_ptr);

    auto batch_matmul = ep::primitive::NewPrimitive<ep::primitive::BatchMatmulFactory>(
        ctx->device_type(), data_type, ep::primitive::BlasTransposeType::N,
        ep::primitive::BlasTransposeType::N);
    batch_matmul->Launch(ctx->stream(), batch_size, concated_padded_dim, vector_size,
                         concated_padded_dim, 1.0, matmul_out_grad_ptr,
                         padded_concated_features_ptr, 0.0, padded_concated_features_grad_ptr);

    ConcatFeaturesGrad(ctx, batch_size, concated_padded_dim, vector_size,
                       padded_concated_features_grad_ptr);
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};

template<typename T>
user_op::InferTmpSizeFn GenFusedDotFeatureInteractionGradInferTmpSizeFn() {
  return [](user_op::InferContext* ctx) {
    int64_t features_concated_dim = 0;
    for (int32_t i = 0; i < ctx->output_size("features_grad"); ++i) {
      features_concated_dim += ctx->InputShape("features_grad", i).At(1);
    }
    const int64_t concated_padded_dim = GetPaddedDim(features_concated_dim);
    const int64_t batch_size = ctx->InputShape("features_grad", 0).At(0);
    const int64_t vector_size = ctx->InputShape("features_grad", 0).At(2);
    size_t matmul_out_grad_size =
        GetCudaAlignedSize(batch_size * concated_padded_dim * concated_padded_dim * sizeof(T));
    size_t padded_concated_features_grad_size =
        GetCudaAlignedSize(batch_size * concated_padded_dim * vector_size * sizeof(T));
    size_t padded_concated_features_size = padded_concated_features_grad_size;
    return matmul_out_grad_size + padded_concated_features_grad_size
           + padded_concated_features_size;
  };
}

#define REGISTER_FUSED_DOT_FEATURE_INTERACTION_GRAD_KERNEL(dtype)                      \
  REGISTER_USER_KERNEL("fused_dot_feature_interaction_grad")                           \
      .SetCreateFn<FusedDotFeatureInteractionGradKernel<dtype>>()                      \
      .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA)                 \
                       && (user_op::HobDataType("dy", 0) == GetDataType<dtype>::value) \
                       && (user_op::HobAttr<std::string>("pooling") == "none"))        \
      .SetInferTmpSizeFn(GenFusedDotFeatureInteractionGradInferTmpSizeFn<dtype>());

REGISTER_FUSED_DOT_FEATURE_INTERACTION_GRAD_KERNEL(float)
REGISTER_FUSED_DOT_FEATURE_INTERACTION_GRAD_KERNEL(half)

template<typename T>
class FusedDotFeatureInteractionPoolingSumGradKernel final : public user_op::OpKernel,
                                                             public user_op::CudaGraphSupport {
 public:
  FusedDotFeatureInteractionPoolingSumGradKernel() = default;
  ~FusedDotFeatureInteractionPoolingSumGradKernel() override = default;

 private:
  using user_op::OpKernel::Compute;
  void Compute(user_op::KernelComputeContext* ctx) const override {
    const int input_size = ctx->input_size("features");
    if (input_size == 1) {
      DispatchFeatureInteractionSumGradInputSize<T, 1>(ctx, input_size);
    } else if (input_size == 2) {
      DispatchFeatureInteractionSumGradInputSize<T, 2>(ctx, input_size);
    } else if (input_size <= 8) {
      DispatchFeatureInteractionSumGradInputSize<T, 8>(ctx, input_size);
    } else {
      CHECK_LE(input_size, 128) << "input_size must not greater than 128. ";
      DispatchFeatureInteractionSumGradInputSize<T, 128>(ctx, input_size);
    }
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};

#define REGISTER_FUSED_DOT_FEATURE_INTERACTION_POOLING_SUM_GRAD_KERNEL(dtype)          \
  REGISTER_USER_KERNEL("fused_dot_feature_interaction_grad")                           \
      .SetCreateFn<FusedDotFeatureInteractionPoolingSumGradKernel<dtype>>()            \
      .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA)                 \
                       && (user_op::HobDataType("dy", 0) == GetDataType<dtype>::value) \
                       && (user_op::HobAttr<std::string>("pooling") == "sum"));

REGISTER_FUSED_DOT_FEATURE_INTERACTION_POOLING_SUM_GRAD_KERNEL(float)
REGISTER_FUSED_DOT_FEATURE_INTERACTION_POOLING_SUM_GRAD_KERNEL(half)

}  // namespace oneflow