/*
 * SPDX-FileCopyrightText: Copyright (c) 2024-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 <cassert>

#include "cuda_hint.cuh"
#include "defines.h"
#if !USE_CUSTOM_BARRIER
#include <cuda/std/barrier>
using CtaBarrier = cuda::barrier<cuda::thread_scope_block>;
#else

#ifndef __CUDACC__
#include <cuda_runtime_api.h>
#endif

#if __CUDACC_VER_MAJOR__ < 12
#define STR_REL_CTA ""
#define STR_ACQ_CTA ""
#else
#define STR_REL_CTA ".release.cta"
#define STR_ACQ_CTA ".acquire.cta"
#endif

enum class Scope : uint32_t {
  CTA = 0,
  CGA = 1,
};

enum class ArriveOrder : uint32_t {
  RELEASE = 0,
  RELAXED = 1,
};

enum class ArrivalToken : uint64_t {};

template <Scope defaultScope_ = Scope::CTA>
class MBarrier  // rename this to MBarrier
{
 public:
  using ArrivalToken = ::ArrivalToken;
  static constexpr Scope defaultScope = defaultScope_;
  using arrival_token = ArrivalToken;

  __device__ inline MBarrier(uint32_t count) {
    assert(count > 0);
    asm volatile("mbarrier.init.b64 [%0], %1;\n" ::"l"(addr()), "r"(count) : "memory");
  }

  __device__ ~MBarrier() { asm volatile("mbarrier.inval.b64 [%0];\n" ::"l"(addr()) : "memory"); }

  template <Scope scope = defaultScope, ArriveOrder order = ArriveOrder::RELEASE>
  __device__ inline mha::conditional_t<scope == Scope::CTA, ArrivalToken, void> arrive(
      uint32_t update = 1) {
    ArrivalToken token{};
#if __CUDA_ARCH__ == 1210
    // GB10 (sm121) does not reliably honor the mbarrier release ordering under
    // tight scheduling: make prior writes visible before the arrival is observed.
    __threadfence_block();
#endif
#if __CUDA_ARCH__ >= 900
    if constexpr (scope == Scope::CTA) {
      switch (order) {
        case ArriveOrder::RELEASE:
          asm volatile("mbarrier.arrive.release.cta.b64 %0, [%1], %2;\n"
                       : "=l"(token)
                       : "l"(addr()), "r"(update)
                       : "memory");
          break;
        case ArriveOrder::RELAXED:
          asm volatile("mbarrier.arrive.relaxed.cta.b64 %0, [%1], %2;\n"
                       : "=l"(token)
                       : "l"(addr()), "r"(update)
                       : "memory");
          break;
      }
      return token;
    } else {
      static_assert(scope == Scope::CGA);
      switch (order) {
        case ArriveOrder::RELEASE:
          asm volatile("mbarrier.arrive.release.cluster.b64 _, [%0], %1;\n" ::"l"(addr()),
                       "r"(update)
                       : "memory");
          break;
        case ArriveOrder::RELAXED:
          asm volatile("mbarrier.arrive.relaxed.cluster.b64 _, [%0], %1;\n" ::"l"(addr()),
                       "r"(update)
                       : "memory");
          break;
      }
      return;
    }
#else
    static_assert(scope == Scope::CTA && order == ArriveOrder::RELEASE);
    if (update > 1) {
      asm volatile("mbarrier.arrive.noComplete" STR_REL_CTA ".b64 %0, [%1], %2;\n"
                   : "=l"(token)
                   : "l"(addr()), "r"(update - 1U)
                   : "memory");
      ArrivalToken refToken;
      asm volatile("mbarrier.arrive" STR_REL_CTA ".b64 %0, [%1];\n"
                   : "=l"(refToken)
                   : "l"(addr())
                   : "memory");
      assert(token == refToken);
      return token;
    } else {
      asm volatile("mbarrier.arrive" STR_REL_CTA ".b64 %0, [%1];\n"
                   : "=l"(token)
                   : "l"(addr())
                   : "memory");
      return token;
    }
#endif
  }

  __device__ inline bool isLocal() const {
    uint32_t addrCtaRank{};
    asm("getctarank.u64 %0, %1;\n" : "=r"(addrCtaRank) : "l"(addr()));
    uint32_t ctaRank{};
    asm("mov.u32 %0, %%cluster_ctarank;\n" : "=r"(ctaRank));
    return addrCtaRank == ctaRank;
  }

  __device__ inline void remoteArrive(uint32_t update = 1) {
#if __CUDA_ARCH__ >= 900
    assert(!isLocal());
    asm volatile("mbarrier.arrive.release.cluster.shared::cluster.b64 _, [%0], %1;\n"
                 :
                 : "l"(__cvta_generic_to_shared(&mBar)), "r"(update)
                 : "memory");
#else
    asm volatile("trap;\n");
#endif
  }

  template <Scope scope = defaultScope, ArriveOrder order = ArriveOrder::RELEASE>
  __device__ inline mha::conditional_t<scope == Scope::CTA, ArrivalToken, void> arrive_tx_relaxed(
      uint32_t txCount) {
#if __CUDA_ARCH__ >= 900
    if constexpr (scope == Scope::CTA) {
      ArrivalToken token{};
      asm volatile("mbarrier.arrive.expect_tx.relaxed.cta.b64 %0, [%1], %2;\n"
                   : "=l"(token)
                   : "l"(addr()), "r"(txCount)
                   : "memory");
      return token;
    } else {
      asm volatile("mbarrier.arrive.expect_tx.relaxed.cluster.b64 _, [%0], %1;\n" ::"l"(addr()),
                   "r"(txCount)
                   : "memory");
      return;
    }
#else
    asm volatile("trap;\n");
#endif
  }

  template <Scope scope = defaultScope, ArriveOrder order = ArriveOrder::RELEASE>
  __device__ inline mha::conditional_t<scope == Scope::CTA, ArrivalToken, void> arrive_tx(
      uint32_t txCount, uint32_t arriveCount = 1) {
#if __CUDA_ARCH__ >= 900
    if (arriveCount == 1) {
      if constexpr (scope == Scope::CTA) {
        ArrivalToken token{};
        switch (order) {
          case ArriveOrder::RELEASE:
            asm volatile("mbarrier.arrive.expect_tx.release.cta.b64 %0, [%1], %2;\n"
                         : "=l"(token)
                         : "l"(addr()), "r"(txCount)
                         : "memory");
            break;
          case ArriveOrder::RELAXED:
            asm volatile("mbarrier.arrive.expect_tx.relaxed.cta.b64 %0, [%1], %2;\n"
                         : "=l"(token)
                         : "l"(addr()), "r"(txCount)
                         : "memory");
            break;
        }
        return token;
      } else {
        static_assert(scope == Scope::CGA);
        switch (order) {
          case ArriveOrder::RELEASE:
            asm volatile(
                "mbarrier.arrive.expect_tx.release.cluster.b64 _, [%0], %1;\n" ::"l"(addr()),
                "r"(txCount)
                : "memory");
            break;
          case ArriveOrder::RELAXED:
            asm volatile(
                "mbarrier.arrive.expect_tx.relaxed.cluster.b64 _, [%0], %1;\n" ::"l"(addr()),
                "r"(txCount)
                : "memory");
            break;
        }
        return;
      }
    } else {
      if constexpr (scope == Scope::CTA) {
        asm volatile("mbarrier.expect_tx.relaxed.cta.b64 [%0], %1;\n" ::"l"(addr()), "r"(txCount)
                     : "memory");
      } else {
        asm volatile("mbarrier.expect_tx.relaxed.cluster.b64 [%0], %1;\n" ::"l"(addr()),
                     "r"(txCount)
                     : "memory");
      }
      return arrive<scope, order>(arriveCount);
    }
#else
    asm volatile("trap;\n");
#endif
  }

  template <Scope scope = defaultScope>
  __device__ inline bool test_wait(ArrivalToken&& token) {
    uint32_t ready{};
    if constexpr (scope == Scope::CGA) {
      asm volatile(
          "{\n"
          ".reg .pred ready;\n"
          "mbarrier.test_wait.acquire.cluster.b64 ready, [%1], %2;\n"
          "selp.b32 %0, 1, 0, ready;\n"
          "}\n"
          : "=r"(ready)
          : "l"(addr()), "l"(token)
          : "memory");
    } else {
      asm volatile(
          "{\n"
          ".reg .pred ready;\n"
          "mbarrier.test_wait" STR_ACQ_CTA
          ".b64 ready, [%1], %2;\n"
          "selp.b32 %0, 1, 0, ready;\n"
          "}\n"
          : "=r"(ready)
          : "l"(addr()), "l"(token)
          : "memory");
    }
    return ready != 0;
  }

  template <Scope scope = defaultScope>
  __device__ inline bool test_wait_parity(bool parity) {
    uint32_t ready{};
    if constexpr (scope == Scope::CGA) {
      asm volatile(
          "{\n"
          ".reg .pred ready;\n"
          "mbarrier.test_wait.parity.acquire.cluster.b64 ready, [%1], %2;\n"
          "selp.b32 %0, 1, 0, ready;\n"
          "}\n"
          : "=r"(ready)
          : "l"(addr()), "r"(uint32_t{parity})
          : "memory");
    } else {
      asm volatile(
          "{\n"
          ".reg .pred ready;\n"
          "mbarrier.test_wait.parity" STR_ACQ_CTA
          ".b64 ready, [%1], %2;\n"
          "selp.b32 %0, 1, 0, ready;\n"
          "}\n"
          : "=r"(ready)
          : "l"(addr()), "r"(uint32_t{parity})
          : "memory");
    }
    return ready != 0;
  }
#if __CUDA_ARCH__ >= 900
  template <Scope scope = defaultScope>
  __device__ inline bool try_wait(ArrivalToken&& token) {
    uint32_t ready{};
    if constexpr (scope == Scope::CGA) {
      asm volatile(
          "{\n"
          ".reg .pred ready;\n"
          "mbarrier.try_wait.acquire.cluster.b64 ready, [%1], %2, %3;\n"
          "selp.b32 %0, 1, 0, ready;\n"
          "}\n"
          : "=r"(ready)
          : "l"(addr()), "l"(token), "n"(kSUSPEND_TIME_HINT)
          : "memory");
    } else {
      asm volatile(
          "{\n"
          ".reg .pred ready;\n"
          "mbarrier.try_wait.acquire.cta.b64 ready, [%1], %2, %3;\n"
          "selp.b32 %0, 1, 0, ready;\n"
          "}\n"
          : "=r"(ready)
          : "l"(addr()), "l"(token), "n"(kSUSPEND_TIME_HINT)
          : "memory");
    }
    return ready != 0;
  }

  template <Scope scope = defaultScope>
  __device__ inline bool try_wait_parity(bool parity) {
    uint32_t ready{};
    if constexpr (scope == Scope::CGA) {
      asm volatile(
          "{\n"
          ".reg .pred ready;\n"
          "mbarrier.try_wait.parity.acquire.cluster.b64 ready, [%1], %2, %3;\n"
          "selp.b32 %0, 1, 0, ready;\n"
          "}\n"
          : "=r"(ready)
          : "l"(addr()), "r"(uint32_t{parity}), "n"(kSUSPEND_TIME_HINT)
          : "memory");
    } else {
      asm volatile(
          "{\n"
          ".reg .pred ready;\n"
          "mbarrier.try_wait.parity.acquire.cta.b64 ready, [%1], %2, %3;\n"
          "selp.b32 %0, 1, 0, ready;\n"
          "}\n"
          : "=r"(ready)
          : "l"(addr()), "r"(uint32_t{parity}), "n"(kSUSPEND_TIME_HINT)
          : "memory");
    }
    return ready != 0;
  }
#endif
  template <Scope scope = defaultScope>
  __device__ inline void wait(ArrivalToken&& token) {
#if __CUDA_ARCH__ >= 900
    poll<true>([&]() { return try_wait<scope>(ArrivalToken{token}); });
#else
    poll<false>([&]() { return test_wait<scope>(ArrivalToken{token}); });
#endif
  }

  // starting from `parity = false`.
  template <Scope scope = defaultScope>
  __device__ inline void wait_parity(bool parity) {
#if __CUDA_ARCH__ >= 900
    poll<true>([&]() { return try_wait_parity<scope>(parity); });
#else
    poll<false>([&]() { return test_wait_parity<scope>(parity); });
#endif
  }

  template <Scope scope = defaultScope, ArriveOrder order = ArriveOrder::RELEASE>
  __device__ inline mha::enable_if_t<scope == Scope::CTA, void> arrive_and_wait(
      uint32_t update = 1) {
    wait<scope>(arrive<scope, order>(update));
  }

 private:
  __device__ inline uint64_t addr() const { return reinterpret_cast<uint64_t>(&mBar); }

  template <bool funcSupportsBlocking, typename F>
  __device__ inline static void poll(F&& func) {
    if constexpr (funcSupportsBlocking) {
      while (!func()) {
      }
    } else {
      float sleepDuration = 0.125F;
      while (!func()) {
        // if (sleepDuration > 1) {
        __nanosleep(uint32_t(sleepDuration));
        // }
        sleepDuration = sleepDuration * 1.25F + 0.F;
      }
    }
  }

 public:
  static constexpr uint32_t kSUSPEND_TIME_HINT = 0xFFFFFFFFU;

 private:
  uint64_t mBar;
};

template <Scope defaultScope>
__device__ inline void init(MBarrier<defaultScope>* bar, uint32_t count) {
  new (bar) MBarrier<defaultScope>{count};
}

using CtaBarrier = MBarrier<Scope::CTA>;
using CgaBarrier = MBarrier<Scope::CGA>;

template <uint32_t nbBars>
__device__ inline constexpr bool toParity(uint32_t i) {
  return i % (nbBars * 2) / nbBars;
}

class NamedBarrier {
 public:
  __device__ inline NamedBarrier(uint32_t idxBar, uint32_t arriveCount)
      : mName{idxBar}, mArriveCount{arriveCount} {
    assert(idxBar < 16 && arriveCount % 32 == 0);
  }

  __device__ inline void arrive() const {
    asm volatile("barrier.cta.arrive %0, %1;\n" ::"r"(mName), "r"(mArriveCount) : "memory");
  }

  __device__ inline void arrive_and_wait() const {
    asm volatile("barrier.cta.sync %0, %1;\n" ::"r"(mName), "r"(mArriveCount) : "memory");
  }

 private:
  uint32_t const mName;
  uint32_t const mArriveCount;
};

__device__ inline void namedBarSync(uint32_t idxBar, uint32_t arriveCount) {
  NamedBarrier bar{idxBar, arriveCount};
  bar.arrive_and_wait();
}
#endif
