#ifndef TRITON_CONVERSION_TRITONNVIDIAGPU_TO_LLVM_UTILITY_H
#define TRITON_CONVERSION_TRITONNVIDIAGPU_TO_LLVM_UTILITY_H

#include <cstdint>
#include <optional>

#include "nvidia/include/TritonNVIDIAGPUToLLVM/PTXAsmFormat.h"

#include "triton/Conversion/TritonGPUToLLVM/Utility.h"

#include "TargetInfo.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "third_party/nvidia/include/Dialect/NVGPU/IR/Dialect.h"
#include "triton/Analysis/Utility.h"
#include "triton/Conversion/MLIRTypes.h"
#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"

#define DEBUG_TYPE "ttgpu_to_llvm"

using namespace mlir;
using namespace mlir::triton;

// Shortcuts for some commonly used LLVM ops to keep code simple and intuitive
// Operators

namespace mlir::triton::gpu {
class MemDescType;
}

namespace mlir {
namespace LLVM {

namespace NVIDIA {
class TargetInfo;

Value shuffleXor(Location loc, RewriterBase &rewriter, Value val, int i);
Value shuffleUp(Location loc, RewriterBase &rewriter, Value val, int i);
Value shuffleIdx(Location loc, RewriterBase &rewriter, Value val, int i);
Value shuffleIdx(Location loc, RewriterBase &rewriter, Value val, Value i);
Value permute(Location loc, RewriterBase &rewriter, Value a, Value b,
              Value mask);

Value llGetPid(Location loc, RewriterBase &rewriter, ModuleOp moduleOp,
               ProgramIDDim axis);

/// Create a predicate with just single active thread.
Value createElectPredicate(Location loc, OpBuilder &rewriter);
Value createElectPredicateWarp0(Location loc, OpBuilder &rewriter);

// Create bar.warp.sync
void createSyncWarp(Location loc, OpBuilder &builder);

// Lower ldmatrix and stmatrix
LogicalResult lowerLdStMatrix(
    Location loc, LinearLayout cvt, bool transpose,
    SmallVector<Value> &vals, // Input for stmatrix, output for ldmatrix
    Value smemBase, Value affineOffset, uint64_t maskSpanAffineOffset,
    Type llvmElemTy, ConversionPatternRewriter &rewriter,
    const mlir::triton::NVIDIA::TargetInfo &targetInfo);

// Given a broadcast mask and the number of CTAs, create a mask of ones
// where for ctaId, it sets as 1's the positions that are in the same broadcast
// group
Value createTMAMulticastMask(Location loc, ConversionPatternRewriter &rewriter,
                             uint16_t broadcastBits, Value ctaId = {});

uint32_t getCGABroadcastMask(mlir::triton::gpu::MemDescType barrierTy);

// Returns the lead CTA predicate for this barrier layout when lowering through
// cluster scope. Returns std::nullopt for CTA-local lowering.
std::optional<Value>
getLeaderCTAPredicate(Location loc, ConversionPatternRewriter &rewriter,
                      mlir::triton::gpu::MemDescType barrierTy);

// Returns the lead CTA barrier address for this layout. If there is no
// cross-cluster lowering, returns barrierPtr unchanged.
Value getLeaderAddress(Location loc, ConversionPatternRewriter &rewriter,
                       Value barrierPtr,
                       mlir::triton::gpu::MemDescType barrierTy);

/// Create a predicate where only the lead CTA is active for two CTA mode.
Value createLeadCTAPredicate(Location loc, RewriterBase &rewriter);
} // namespace NVIDIA
} // namespace LLVM

} // namespace mlir

#endif
