// Copyright (c) 2025 PaddlePaddle 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.
/*
 * copyright (C) 2025 KUNLUNXIN, Inc
 */

#include "xpu/kernel/cluster.h"
#include "xpu/kernel/cluster_partition.h"
#include "xpu/kernel/cluster_primitive.h"
#include "xpu/kernel/cluster_primitive_template.h"

namespace xpu3 {
namespace plugin {

static __device__ void do_calc(const _shared_ptr_ int* lm_x,
                               int* lm_y1,
                               int* lm_y2,
                               int64_t size,
                               int& text_count,
                               int& images_count) {
  for (int j = 0; j < size; j++) {
    if (lm_x[j] == 0) {
      lm_y1[j] = text_count;
      text_count += 1;
    } else {
      lm_y2[j] = images_count;
      images_count += 1;
    }
  }
  mfence_lm_sm();
}

__global__ void text_image_index_out_kernel(const int* token_type_ids,  // x
                                            int* text_index,            // y1
                                            int* image_index,           // y2
                                            const int64_t token_num) {
  const int cid = core_id();
  const int tid = core_id() * cluster_num() + cluster_id();
  const int nthreads = core_num() * cluster_num();
  if (tid >= 1) return;
  constexpr int BUFSIZE = 1024;
  constexpr int READ_MAX_SIZE = BUFSIZE / sizeof(int);
  const int64_t len = token_num;

  __simd__ char buffer0[BUFSIZE * 3];
  __simd__ char buffer1[BUFSIZE * 3];
  __simd__ __shared__ char buffer2[64][BUFSIZE * 2];

  DoublePtr<READ_MAX_SIZE, SmPtr<int>> buffer_ptr_x(
      (SmPtr<int>((_shared_ptr_ int*)buffer2[cid])));
  TriplePtr<READ_MAX_SIZE, LmPtr<int>> buffer_ptr_y1(
      (LmPtr<int>((int*)buffer0)));
  TriplePtr<READ_MAX_SIZE, LmPtr<int>> buffer_ptr_y2(
      (LmPtr<int>((int*)buffer1)));
  int64_t buflen = get_1d_buflen(len, nthreads, READ_MAX_SIZE, 64);
  int64_t i = tid * buflen;
  int read_size = 0;
  int offset = nthreads * buflen;

  int text_count = 0;
  int images_count = 0;

  if (i < len) {
    read_size = min<int64_t>(buflen, len - i);
    buffer_ptr_y1.gm_load_async(text_index + tid * buflen, read_size);
    buffer_ptr_y2.gm_load_async(image_index + tid * buflen, read_size);
    buffer_ptr_x.gm_load_async(token_type_ids + tid * buflen, read_size);
    mfence();
  }
  while (i < len && i + offset < len) {
    i = i + offset;
    int read_size_next = min<int64_t>(buflen, len - i);
    buffer_ptr_x.next().gm_load_async(token_type_ids + i, read_size_next);
    buffer_ptr_y1.next().gm_load_async(text_index + i, read_size_next);
    buffer_ptr_y2.next().gm_load_async(image_index + i, read_size_next);

    do_calc(buffer_ptr_x.ptr,
            buffer_ptr_y1.ptr,
            buffer_ptr_y2.ptr,
            read_size,
            text_count,
            images_count);

    buffer_ptr_y1.gm_store_async(text_index + i - offset, read_size);
    buffer_ptr_y2.gm_store_async(image_index + i - offset, read_size);
    buffer_ptr_x.toggle();
    buffer_ptr_y1.toggle();
    buffer_ptr_y2.toggle();
    read_size = read_size_next;
  }
  if (i < len) {
    do_calc(buffer_ptr_x.ptr,
            buffer_ptr_y1.ptr,
            buffer_ptr_y2.ptr,
            read_size,
            text_count,
            images_count);
    buffer_ptr_y1.gm_store_async(text_index + i, read_size);
    buffer_ptr_y2.gm_store(image_index + i, read_size);
  }
}
}  // namespace plugin
}  // namespace xpu3
