// Copyright 2023 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

// Polyfill Metal kernels for buffer fills without aligned offsets / 4-byte
// patterns.

struct FillSpec {
  uint64_t buffer_offset;  // Buffer offset to fill (in bytes)
  uint64_t buffer_length;  // Buffer length to fill (in bytes)
  uint32_t pattern;        // 32-bit fill pattern
};

// Fills target |buffer| with the given |spec|ification.
//
// The target |buffer| is assumed to have 16-byte aligned offset/length.
// Each thread fills one 4-component 32-bit element vector.
kernel void fill_buffer_16byte(device uint4* buffer [[buffer(0)]],
                               constant FillSpec& spec [[buffer(1)]],
                               uint id [[thread_position_in_grid]]) {
  uint64_t end = spec.buffer_length / 16;
  if (id >= end) return;
  uint64_t start = spec.buffer_offset / 16;
  buffer[start + id] =
      uint4(spec.pattern, spec.pattern, spec.pattern, spec.pattern);
}

// Fills target |buffer| with the given |spec|ification.
//
// The target |buffer| is assumed to have 4-byte aligned offset/length.
// Each thread fills one 32-bit scalar.
kernel void fill_buffer_4byte(device uint32_t* buffer [[buffer(0)]],
                              constant FillSpec& spec [[buffer(1)]],
                              uint id [[thread_position_in_grid]]) {
  uint64_t end = spec.buffer_length / 4;
  if (id >= end) return;
  uint64_t start = spec.buffer_offset / 4;
  buffer[start + id] = spec.pattern;
}

// Fills target |buffer| with the given |spec|ification.
//
// The target |buffer| is assumed to have 1-byte aligned offset/length.
// Each thread fills one 8-bit scalar.
kernel void fill_buffer_1byte(device uint32_t* buffer [[buffer(0)]],
                              constant FillSpec& spec [[buffer(1)]],
                              uint id [[thread_position_in_grid]]) {
  // We split the full buffer fill range into three parts:
  // 1. Left bytes: containing (0 to 3) bytes before the first 4-byte aligned
  // address
  // 2. Middle bytes: aligned 32-bit scalars in the middle
  // 3. Right bytes: containing (0 to 3) bytes since the last 4-byte aligned
  // address
  //
  // Threads are distributed from the perspective of handling middle 32-bit
  // scalars. We use the first thread to *additionally* handle left and right
  // bytes.
  uint8_t left_byte_count = spec.buffer_offset % 4;
  uint8_t right_byte_count = (spec.buffer_offset + spec.buffer_length) % 4;
  uint32_t middle_pattern =
      metal::rotate(spec.pattern, uint32_t(8 * left_byte_count));

  // Masks for left bytes and right bytes from rotated pattern. Note that for
  // *little endian*, left bytes will replace high bits of the leftmost touched
  // 32-bit scalar, while right bytes will replace low bits of the rightmost
  // touched 32-bit scalar.
  uint32_t left_mask = ~((uint64_t(1) << (8 * left_byte_count)) - 1);
  uint32_t right_mask = (uint64_t(1) << (8 * right_byte_count)) - 1;

  // Indexing start points in |buffer| for the three parts.
  uint64_t left_start = spec.buffer_offset / 4;
  uint64_t middle_start = (spec.buffer_offset + 3) / 4;
  uint64_t right_start = (spec.buffer_offset + spec.buffer_length) / 4;

  if (middle_start < right_start) {  // Middle bytes
    if (middle_start + id >= right_start) return;
    buffer[middle_start + id] = middle_pattern;
  }

  if (id == 0) {
    if (left_start == right_start) {
      // The entire fill range fits within a single 32-bit scalar; read-modify-
      // write only the bytes in [left_byte_count, right_byte_count).
      uint32_t mask = left_mask & right_mask;
      uint32_t old = buffer[left_start];
      buffer[left_start] = (old & (~mask)) | (middle_pattern & mask);
    } else {
      if (left_byte_count != 0) {  // Left bytes
        uint32_t old = buffer[left_start];
        buffer[left_start] =
            (old & (~left_mask)) | (middle_pattern & left_mask);
      }
      if (right_byte_count != 0) {  // Right bytes
        uint32_t old = buffer[right_start];
        buffer[right_start] =
            (old & (~right_mask)) | (middle_pattern & right_mask);
      }
    }
  }
}
