/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION &
* AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
*
* 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.
*/
#pragma once

#include <cstdint>

namespace batchedGemm {

namespace gemm {

enum class AllReduceAlgo : uint32_t {
  // Does not apply all-reduce.
  None = 0,
  // Reduction occurs at L2 cache; pulls N-1 partial outputs from peer devices. Result is
  // non-deterministic. Potentially lower latency at cost of higher memory traffic.
  OneShot,
  // Reduction occurs at switch; pulls 1/Nth of the output from switch (reduce-scatter phase) and
  // store to multicast mem (all-gather phase). Result is deterministic. Lower memory traffic at
  // cost of potentially higher latency.
  TwoShot,
};

////////////////////////////////////////////////////////////////////////////////////////////////////

enum class MatrixLayout {
  // K-major layout (default). [Mn, K]
  MajorK = 0,
  // M-major for A and N-major for B. [K, Mn]
  MajorMn,
  // Layout is blocked along the K dimension as seen in the diagram below. [K / blockK, Mn, blockK]
  // where blockK is fixed at 128B
  //
  //         ├──────────────  K  ──────────────┤
  //  ┬  ┬   ├──── K block ───┤
  //  │  │   │ 0   1   2   3  ║ 32  33  34  35 │
  //  │ CTA0 │ 4   5   6   7  ║ 36  37  38  39 │
  //  │  │   │ 8   9   10  11 ║ 40  41  42  43 │
  //  │  ┴   │ 12  13  14  15 ║ 44  45  46  47 │
  //  M  ┬   ├────────────────║────────────────┤
  //  │  │   │ 16  17  18  19 ║ 48  49  50  51 │
  //  │ CTA1 │ 20  21  22  23 ║ 52  53  54  55 │
  //  │  │   │ 24  25  26  27 ║ 56  57  58  59 │
  //  ┴  ┴   │ 28  29  30  31 ║ 60  61  62  63 │
  BlockMajorK
};

////////////////////////////////////////////////////////////////////////////////////////////////////

enum class SplitK : uint32_t {
  // No split-k is needed. I.e. mNumSlicesForSplitK == 1.
  None = 0,
  // CTAs computing one MN tile save partial results to global memory.
  // Then wait on the barrier and the last CTA in the group loads partial results from gmem,
  // sums them up and writes back to gmem.
  Gmem,
  // All CTAs in one CGA calculate partial sums. Then send the results to the smem of
  // the last CTA in the CGA, which sums them up and writes to gmem.
  Dsmem,
};

////////////////////////////////////////////////////////////////////////////////////////////////////

enum class BiasType : uint32_t {
  // No bias.
  None = 0,
  // One bias value per N of the output tensor.
  M = 1,
  // One bias value per row M of the output tensor.
  N = 2,
  // One bias value for each element of the output tensor.
  Mn = 3,
};

////////////////////////////////////////////////////////////////////////////////////////////////////

enum class TileScheduler {
  // Static scheduler (Non-persistent).
  Static = 0,
  // Dynamic persistent scheduler. This is either based on an atomically incremented global work id
  // prior to SM100 archs, or the HW supported work id scheduler based on UGETNEXTWORKID for SM100+.
  Persistent,
};

////////////////////////////////////////////////////////////////////////////////////////////////////

// Helper functions to check the SplitK type.

#define SPLIT_K_FUNCTION(Mode)                                                                     \
  inline bool doesSplitKUse##Mode(SplitK mode) {                                                   \
    return (mode == SplitK::Mode);                                                                 \
  }

SPLIT_K_FUNCTION(Gmem)
SPLIT_K_FUNCTION(Dsmem)

#undef SPLIT_K_FUNCTION

////////////////////////////////////////////////////////////////////////////////////////////////////

// Helper functions to check the Bias type.

#define BIAS_TYPE_FUNCTION(Mode)                                                                   \
  inline bool isBiasType##Mode(BiasType type) {                                                    \
    return (type == BiasType::Mode);                                                               \
  }

BIAS_TYPE_FUNCTION(None)
BIAS_TYPE_FUNCTION(N)
BIAS_TYPE_FUNCTION(M)
BIAS_TYPE_FUNCTION(Mn)

#undef BIAS_TYPE_FUNCTION

////////////////////////////////////////////////////////////////////////////////////////////////////

} // namespace gemm

} // namespace batchedGemm
