#ifndef TRITON_CONVERSION_TRITONGPU_TO_LLVM_TARGETINFONVIDIA_H
#define TRITON_CONVERSION_TRITONGPU_TO_LLVM_TARGETINFONVIDIA_H

#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h"
#include "triton/Dialect/TritonNvidiaGPU/IR/TargetFeatures.h"

namespace mlir::triton::NVIDIA {

class TargetInfo : public mlir::triton::TargetInfoBase {
public:
  explicit TargetInfo(int computeCapability)
      : TargetInfo(computeCapability, /*ptxVersion=*/0) {}
  TargetInfo(int computeCapability, int ptxVersion)
      : targetFeatures(computeCapability), ptxVersion(ptxVersion) {}

  bool supportMaximumMinimum() const override;

  Value getClusterCTAId(RewriterBase &rewriter, Location loc) const override;

  Value ballot(RewriterBase &rewriter, Location loc, Type type,
               Value cmp) const override;

  Value getGlobalTimer(RewriterBase &rewriter, Location loc) const override;

  StringRef getAtomicSyncScope(MemSyncScope scope) const override;

  void barrier(Location loc, RewriterBase &rewriter,
               triton::gpu::AddrSpace targets) const override;
  void clusterBarrier(Location loc, RewriterBase &rewriter,
                      Operation *sourceOp) const override;

  void warpSync(Location loc, RewriterBase &rewriter) const override;

  void storeDShared(RewriterBase &rewriter, Location loc, Value ptr,
                    Value ctaId, Value val, Value pred) const override;
  Value loadDShared(RewriterBase &rewriter, Location loc, Value ptr,
                    Value ctaId, Type elemTy, Value pred,
                    Operation *localLoadOp = nullptr) const override;
  Value mapDShared(RewriterBase &rewriter, Location loc, Value ptr, Value ctaId,
                   Value pred) const;

  bool supportLdMatrix() const override {
    return targetFeatures.supportLdMatrix();
  }
  bool supportStMatrix() const override {
    return targetFeatures.supportStMatrix();
  }
  bool supportLdStMatrixB8() const override {
    return targetFeatures.supportLdStMatrixB8();
  }
  bool supportBitwidth16Elementwise() const override {
    return targetFeatures.supportBitwidth16Elementwise();
  }
  bool supportBitwidth32Elementwise() const override {
    return targetFeatures.supportBitwidth32Elementwise();
  }

  Value shuffleXor(RewriterBase &rewriter, Location loc, Value val,
                   int i) const override;
  Value shuffleUp(RewriterBase &rewriter, Location loc, Value val,
                  int i) const override;
  Value shuffleIdx(RewriterBase &rewriter, Location loc, Value val,
                   int i) const override;
  Value shuffleIdx(RewriterBase &rewriter, Location loc, Value val,
                   Value i) const override;

  Value permute(RewriterBase &rewriter, Location loc, Value a, Value b,
                Value selector) const override;

  Value programId(RewriterBase &rewriter, Location loc, ModuleOp moduleOp,
                  ProgramIDDim axis) const override;

  bool warpReduce(RewriterBase &rewriter, Location loc, SmallVector<Value> &acc,
                  triton::ReduceOp op,
                  unsigned reduceLaneIdMask) const override;

  std::string getMulhiFuncName(Type resultElementTy) const override;

  void printf(RewriterBase &rewriter, Value formatStrStart,
              int formatStrByteCount, ValueRange args,
              ArrayRef<bool> isSigned = {}) const override;

  void printf(RewriterBase &rewriter, StringRef msg, ValueRange args,

              ArrayRef<bool> isSigned = {}) const override;

  void assertFail(RewriterBase &rewriter, Location loc, StringRef message,
                  StringRef file, StringRef func, int line) const override;

  int getSharedAddressSpace() const override;

  int getAddressSpace(Attribute addressSpace) const override;

  bool supportVectorizedAtomics() const override;

  int getPtxVersion() const { return ptxVersion; }
  int getComputeCapability() const {
    return targetFeatures.getComputeCapability();
  }
  const triton::nvidia_gpu::TargetFeatures &getTargetFeatures() const {
    return targetFeatures;
  }

  bool isCuda() const override { return true; }

private:
  triton::nvidia_gpu::TargetFeatures targetFeatures;
  int ptxVersion;
};

} // namespace mlir::triton::NVIDIA

#endif // TRITON_CONVERSION_TRITONGPU_TO_LLVM_TARGETINFONVIDIA_H
