/* Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of NVIDIA CORPORATION nor the names of its
 *    contributors may be used to endorse or promote products derived
 *    from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/*
 * This sample demonstrates a persistent LayerNorm forward pass using
 * CUDA Tile C++:  y = (x - mean) * rsqrt(var + eps) * weight + bias.
 * The grid launches NUM_SMS persistent blocks; each block walks the
 * row dimension with a grid-stride loop, processing BLOCK_N rows x
 * BLOCK_D cols per iteration and striding by NUM_SMS * BLOCK_N rows.
 * Per-row mean and rstd are reduced over the column dimension and
 * (when COMPUTE_MEAN_AND_RSTD and TRAINING are enabled) saved to
 * float32 side buffers. N, D, NUM_SMS, and EPS are template NTTPs so
 * the tile compiler can fold the loop step, the (1/D) reciprocal,
 * partition_view extents, and the (var + eps) broadcast.  A SIMT
 * kernel is used to initialize X, W, and B on device.
 */

#include "helper_cuda.h"
#include "cuda_tile.h"
#include "cuda_fp16.h"
#include <cstdio>
#include <cstdlib>
#include <cmath>

/* SIMT initializer for X (N x D), W (D,), B (D,) with deterministic data. */
__global__ void initializeInputs(__half* X, __half* W, __half* B,
                                 int N, int D) {
  auto idx   = blockIdx.x * blockDim.x + threadIdx.x;
  auto total = N * D;
  if (idx < total) {
    int m = idx / D;
    int n = idx - m * D;
    X[idx] = __half{float((m + n) % 7) - 3.5f};
  }
  if (idx < D) {
    W[idx] = __half{1.0f + 0.1f * float(idx % 5)};
    B[idx] = __half{0.1f * float(idx % 3)};
  }
}

template<typename T,
         int BLOCK_N,
         int BLOCK_D,
         bool TRAINING,
         bool COMPUTE_MEAN_AND_RSTD,
         int N,
         int D,        // compile-time so `(... / D)` and partition_view extents fold.
         int NUM_SMS,  // compile-time so the persistent for-loop step is constant.
         float EPS>    // compile-time so the (var + eps) reshape/broadcast hoists out of the for-loop
[[ using cutile : hint(0, num_cta_in_cga=1) ]]
__tile_global__ void persistent_layer_norm_fwd_kernel(
    const T* __restrict__ X,       // (N, D)
    T* __restrict__ Y,             // (N, D)
    const T* __restrict__ W,       // (D,)
    const T* __restrict__ B,       // (D,)
    float* __restrict__ Mean,      // (N,)
    float* __restrict__ Rstd       // (N,)
) {
    namespace ct = cuda::tiles;

    using f32_N     = ct::tile<float, ct::shape<BLOCK_N>>;

    X    = ct::assume_aligned<16>(X);
    Y    = ct::assume_aligned<16>(Y);
    W    = ct::assume_aligned<16>(W);
    B    = ct::assume_aligned<16>(B);

    int pid = ct::bid().x;
    constexpr int upper_bound = (N + BLOCK_N - 1) / BLOCK_N;

    // Partitioned views with compile-time extents (N, D are template NTTPs).
    using ExtND = ct::extents<uint32_t, static_cast<uint32_t>(N), static_cast<uint32_t>(D)>;
    using ExtD  = ct::extents<uint32_t, static_cast<uint32_t>(D)>;
    using ExtN  = ct::extents<uint32_t, static_cast<uint32_t>(N)>;
    auto pX = ct::partition_view(
        ct::tensor_span{X, ExtND{}},
        ct::shape<BLOCK_N, BLOCK_D>{});
    auto pY = ct::partition_view(
        ct::tensor_span{Y, ExtND{}},
        ct::shape<BLOCK_N, BLOCK_D>{});
    auto pW = ct::partition_view(
        ct::tensor_span{W, ExtD{}},
        ct::shape<BLOCK_D>{});
    auto pB = ct::partition_view(
        ct::tensor_span{B, ExtD{}},
        ct::shape<BLOCK_D>{});
    auto pMean = ct::partition_view(
        ct::tensor_span{Mean, ExtN{}},
        ct::shape<BLOCK_N>{});
    auto pRstd = ct::partition_view(
        ct::tensor_span{Rstd, ExtN{}},
        ct::shape<BLOCK_N>{});

    // Load weights once (hoisted out of the grid-stride loop).
    auto w = ct::element_cast<float>(pW.load(0));  // (BLOCK_D,)
    auto b = ct::element_cast<float>(pB.load(0));  // (BLOCK_D,)

    // Broadcast weights into (BLOCK_N, BLOCK_D) by reshape to (1, BLOCK_D).
    auto w_bcast = ct::reshape<ct::shape<1, BLOCK_D>>(w);
    auto b_bcast = ct::reshape<ct::shape<1, BLOCK_D>>(b);

    constexpr float inv_D_scalar = 1.0f / static_cast<float>(D);
    auto inv_D_tile = ct::full<f32_N>(inv_D_scalar);
    auto eps_tile   = ct::full<f32_N>(EPS);

    using TileXNxD = ct::tile<T, ct::shape<BLOCK_N, BLOCK_D>>;

    for (auto current_pid : ct::irange(pid, upper_bound, NUM_SMS)) {
        TileXNxD x_tile;
        [[ using cutile : hint(0, latency=4) ]]
        x_tile = pX.load_masked(current_pid, 0);
        auto x = ct::element_cast<float>(x_tile);

        f32_N mean;
        f32_N rstd;

        if constexpr (COMPUTE_MEAN_AND_RSTD) {
            // Step 1: Compute x^2 then sum/mean.  Use the loop-invariant
            // `inv_D_tile` and `eps_tile` (built outside the loop) so the
            // reshape + broadcast of those scalars stays hoisted.
            auto x_squared = x * x;
            auto avg_square_2d = ct::sum<1>(x_squared);
            auto avg_square = ct::reshape<ct::shape<BLOCK_N>>(avg_square_2d) * inv_D_tile;
            auto mean_2d = ct::sum<1>(x);
            mean = ct::reshape<ct::shape<BLOCK_N>>(mean_2d) * inv_D_tile;
            auto var = avg_square - mean * mean;

            rstd = ct::rsqrt(var + eps_tile);

            if constexpr (TRAINING) {
                [[ using cutile : hint(0, allow_tma=false) ]]
                pMean.store_masked(mean, current_pid);
                [[ using cutile : hint(0, allow_tma=false) ]]
                pRstd.store_masked(rstd, current_pid);
            }
        } else {
            mean = pMean.load_masked(current_pid);
            rstd = pRstd.load_masked(current_pid);
        }

        // Broadcast mean/rstd to (BLOCK_N, 1) then rely on implicit broadcast
        // against (BLOCK_N, BLOCK_D).
        auto mean_col = ct::reshape<ct::shape<BLOCK_N, 1>>(mean);
        auto rstd_col = ct::reshape<ct::shape<BLOCK_N, 1>>(rstd);

        auto x_hat = (x - mean_col) * rstd_col;
        auto y_f32 = x_hat * w_bcast + b_bcast;

        auto y_T = ct::element_cast<T>(y_f32);
        [[ using cutile : hint(0, allow_tma=false) ]]
        pY.store_masked(y_T, current_pid, 0);
    }
}

int main() {
  /* BLOCK_D == D so each persistent-loop iteration covers a full row's
   * columns in one tile per row. NUM_SMS is a template NTTP, hence
   * compile-time; 132 matches B200 / Hopper-class GPUs - adjust to
   * match the target device's `multiProcessorCount` for best perf. */
  constexpr int   N = 1024, D = 256;
  constexpr int   BLOCK_N = 4, BLOCK_D = 256;
  constexpr int   NUM_SMS = 132;
  constexpr float EPS     = 1e-5f;

  __half *d_X = nullptr, *d_Y = nullptr, *d_W = nullptr, *d_B = nullptr;
  float  *d_Mean = nullptr, *d_Rstd = nullptr;
  checkCudaErrors(cudaMalloc(&d_X,    N * D * sizeof(__half)));
  checkCudaErrors(cudaMalloc(&d_Y,    N * D * sizeof(__half)));
  checkCudaErrors(cudaMalloc(&d_W,    D     * sizeof(__half)));
  checkCudaErrors(cudaMalloc(&d_B,    D     * sizeof(__half)));
  checkCudaErrors(cudaMalloc(&d_Mean, N     * sizeof(float)));
  checkCudaErrors(cudaMalloc(&d_Rstd, N     * sizeof(float)));

  int init_threads = 256, init_blocks = 1 + ((N * D - 1) / init_threads);
  initializeInputs<<<init_blocks, init_threads>>>(d_X, d_W, d_B, N, D);
  checkCudaErrors(cudaGetLastError());

  /* NUM_SMS is a compile-time NTTP that doubles as the persistent-loop
   * stride; the launch grid x must equal NUM_SMS for correctness.
   * Adjust the constant (and recompile) for one block per SM on a
   * device with a different SM count. */
  persistent_layer_norm_fwd_kernel<__half, BLOCK_N, BLOCK_D,
                                   /*TRAINING=*/true, /*COMPUTE_MEAN_AND_RSTD=*/true,
                                   N, D, NUM_SMS, EPS>
      <<<dim3(NUM_SMS, 1, 1)>>>(d_X, d_Y, d_W, d_B, d_Mean, d_Rstd);
  checkCudaErrors(cudaGetLastError());
  checkCudaErrors(cudaDeviceSynchronize());

  __half* h_Y        = new __half[N * D];
  __half* h_Y_ref    = new __half[N * D];
  float*  h_Mean     = new float[N];
  float*  h_Rstd     = new float[N];
  float*  h_Mean_ref = new float[N];
  float*  h_Rstd_ref = new float[N];
  checkCudaErrors(cudaMemcpy(h_Y,    d_Y,    N * D * sizeof(__half), cudaMemcpyDeviceToHost));
  checkCudaErrors(cudaMemcpy(h_Mean, d_Mean, N     * sizeof(float),  cudaMemcpyDeviceToHost));
  checkCudaErrors(cudaMemcpy(h_Rstd, d_Rstd, N     * sizeof(float),  cudaMemcpyDeviceToHost));

  /* CPU reference in double precision; compare with 1e-1 fp16 tolerance
   * for Y and 1e-3 for the float32 Mean/Rstd outputs. */
  for (int m = 0; m < N; ++m) {
    double sum = 0.0, sumsq = 0.0;
    for (int n = 0; n < D; ++n) {
      double x = double(float((m + n) % 7) - 3.5f);
      sum += x; sumsq += x * x;
    }
    double mu      = sum / double(D);
    double var     = sumsq / double(D) - mu * mu;
    double inv_std = 1.0 / std::sqrt(var + double(EPS));
    h_Mean_ref[m]  = float(mu);
    h_Rstd_ref[m]  = float(inv_std);
    for (int n = 0; n < D; ++n) {
      double x = double(float((m + n) % 7) - 3.5f);
      double w = double(1.0f + 0.1f * float(n % 5));
      double b = double(0.1f * float(n % 3));
      h_Y_ref[m * D + n] = __half(float((x - mu) * inv_std * w + b));
    }
  }

  for (int idx = 0; idx < N * D; ++idx) {
    float expected = float(h_Y_ref[idx]), actual = float(h_Y[idx]);
    if (std::fabs(expected - actual) > 1e-1f) {
      printf("Mismatch h_Y[%d]: expected %f, actual %f\n", idx, expected, actual);
      return 1;
    }
  }
  for (int m = 0; m < N; ++m) {
    if (std::fabs(h_Mean_ref[m] - h_Mean[m]) > 1e-3f) {
      printf("Mismatch h_Mean[%d]: expected %f, actual %f\n", m, h_Mean_ref[m], h_Mean[m]);
      return 1;
    }
    if (std::fabs(h_Rstd_ref[m] - h_Rstd[m]) > 1e-3f) {
      printf("Mismatch h_Rstd[%d]: expected %f, actual %f\n", m, h_Rstd_ref[m], h_Rstd[m]);
      return 1;
    }
  }

  printf("Success! Persistent LayerNorm matches expected results.\n");

  checkCudaErrors(cudaFree(d_X));   checkCudaErrors(cudaFree(d_Y));
  checkCudaErrors(cudaFree(d_W));   checkCudaErrors(cudaFree(d_B));
  checkCudaErrors(cudaFree(d_Mean)); checkCudaErrors(cudaFree(d_Rstd));
  delete[] h_Y;        delete[] h_Y_ref;
  delete[] h_Mean;     delete[] h_Rstd;
  delete[] h_Mean_ref; delete[] h_Rstd_ref;
}