llvm/lib/IR/NVVMIntrinsicUtils.cpp

//===----------------------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements functions associated with NVVM Intrinsics.
//
//===----------------------------------------------------------------------===//

#include "llvm/IR/NVVMIntrinsicUtils.h"
#include "llvm/ADT/StringRef.h"

using namespace llvm;
using namespace nvvm;

void nvvm::printEvictPolicyType(raw_ostream &OS, const Constant *ImmArgVal) {
  const auto *CI = dyn_cast<ConstantInt>(ImmArgVal);
  if (!CI ||
      CI->getZExtValue() > static_cast<uint64_t>(EvictPolicyType::EVICT_LAST)) {
    OS << "Unsupported evict policy";
    return;
  }
  OS << getEvictPolicyName(static_cast<EvictPolicyType>(CI->getZExtValue()));
}

void nvvm::printTMAReductionOp(raw_ostream &OS, const Constant *ImmArgVal) {
  const auto *CI = dyn_cast<ConstantInt>(ImmArgVal);
  if (!CI || CI->getZExtValue() > static_cast<uint64_t>(TMAReductionOp::XOR))
    llvm_unreachable(
        "printTMAReductionOp called with invalid value for immediate argument");

  OS << getTMATensorReductionOpName(
      static_cast<TMAReductionOp>(CI->getZExtValue()));
}

void nvvm::printTMAValidateDataPattern(raw_ostream &OS,
                                       const Constant *ImmArgVal) {
  const auto *CI = dyn_cast<ConstantInt>(ImmArgVal);
  if (!CI || CI->getZExtValue() > static_cast<uint64_t>(
                                      TMAValidateDataPattern::PER_ELEMENT_FF)) {
    OS << "Unknown validate data pattern";
    return;
  }

  OS << getTMAValidateDataPatternName(
      static_cast<TMAValidateDataPattern>(CI->getZExtValue()));
}

void nvvm::printMBarrierLayout(raw_ostream &OS, const Constant *ImmArgVal) {
  if (const auto *CI = dyn_cast<ConstantInt>(ImmArgVal)) {
    switch (static_cast<MBarrierLayout>(CI->getZExtValue())) {
    case MBarrierLayout::V0:
      OS << "v0";
      return;
    case MBarrierLayout::V1:
      OS << "v1";
      return;
    }
  }
}

void nvvm::printMemScope(raw_ostream &OS, const Constant *ImmArgVal) {
  const auto *CI = dyn_cast<ConstantInt>(ImmArgVal);
  if (!CI || CI->getZExtValue() > static_cast<uint64_t>(MemScope::SYS)) {
    OS << "Unknown memory scope";
    return;
  }

  OS << getMemScopeName(static_cast<MemScope>(CI->getZExtValue()));
}

void nvvm::printTcgen05MMAKind(raw_ostream &OS, const Constant *ImmArgVal) {
  if (const auto *CI = dyn_cast<ConstantInt>(ImmArgVal)) {
    uint64_t Val = CI->getZExtValue();
    switch (static_cast<Tcgen05MMAKind>(Val)) {
    case Tcgen05MMAKind::F16:
      OS << "f16";
      return;
    case Tcgen05MMAKind::TF32:
      OS << "tf32";
      return;
    case Tcgen05MMAKind::F8F6F4:
      OS << "f8f6f4";
      return;
    case Tcgen05MMAKind::I8:
      OS << "i8";
      return;
    case Tcgen05MMAKind::TI16:
      OS << "ti16";
      return;
    }
  }
}

void nvvm::printTcgen05CollectorUsageOp(raw_ostream &OS,
                                        const Constant *ImmArgVal) {
  if (const auto *CI = dyn_cast<ConstantInt>(ImmArgVal)) {
    uint64_t Val = CI->getZExtValue();
    switch (static_cast<Tcgen05CollectorUsageOp>(Val)) {
    case Tcgen05CollectorUsageOp::DISCARD:
      OS << "discard";
      return;
    case Tcgen05CollectorUsageOp::LASTUSE:
      OS << "lastuse";
      return;
    case Tcgen05CollectorUsageOp::FILL:
      OS << "fill";
      return;
    case Tcgen05CollectorUsageOp::USE:
      OS << "use";
      return;
    }
  }
}

void nvvm::printTcgen05MMACollectorBBuffer(raw_ostream &OS,
                                           const Constant *ImmArgVal) {
  if (const auto *CI = dyn_cast<ConstantInt>(ImmArgVal)) {
    uint64_t Val = CI->getZExtValue();
    switch (static_cast<Tcgen05MMACollectorBBuffer>(Val)) {
    case Tcgen05MMACollectorBBuffer::B0:
      OS << "b0";
      return;
    case Tcgen05MMACollectorBBuffer::B1:
      OS << "b1";
      return;
    case Tcgen05MMACollectorBBuffer::B2:
      OS << "b2";
      return;
    case Tcgen05MMACollectorBBuffer::B3:
      OS << "b3";
      return;
    }
  }
}

void nvvm::printTensormapElemType(raw_ostream &OS, const Constant *ImmArgVal) {
  static constexpr StringRef TensormapElemTypes[] = {
      "u8",       "u16",   "u32",       "s32",      "u64",  "s64",
      "f16",      "f32",   "f32.ftz",   "f64",      "bf16", "tf32",
      "tf32.ftz", "b4x16", "b4x16_p64", "b6x16_p32"};
  if (const auto *CI = dyn_cast<ConstantInt>(ImmArgVal)) {
    uint64_t Val = CI->getZExtValue();
    if (Val <= static_cast<uint64_t>(nvvm::TensormapElemType::B6x16_p32)) {
      OS << TensormapElemTypes[Val];
      return;
    }
  }
}

void nvvm::printTensormapInterleaveLayout(raw_ostream &OS,
                                          const Constant *ImmArgVal) {
  if (const auto *CI = dyn_cast<ConstantInt>(ImmArgVal)) {
    uint64_t Val = CI->getZExtValue();
    switch (static_cast<TensormapInterleaveLayout>(Val)) {
    case TensormapInterleaveLayout::NO_INTERLEAVE:
      OS << "No interleave";
      return;
    case TensormapInterleaveLayout::INTERLEAVE_16B:
      OS << "16B interleave";
      return;
    case TensormapInterleaveLayout::INTERLEAVE_32B:
      OS << "32B interleave";
      return;
    }
  }
}

void nvvm::printTensormapSwizzleMode(raw_ostream &OS,
                                     const Constant *ImmArgVal) {
  static constexpr StringRef TensormapSwizzleModes[] = {
      "No swizzling", "32B swizzling", "64B swizzling", "128B swizzling",
      "96B swizzling"};
  if (const auto *CI = dyn_cast<ConstantInt>(ImmArgVal)) {
    uint64_t Val = CI->getZExtValue();
    if (Val <= static_cast<uint64_t>(nvvm::TensormapSwizzleMode::SWIZZLE_96B)) {
      OS << TensormapSwizzleModes[Val];
      return;
    }
  }
}

void nvvm::printTensormapSwizzleAtomicity(raw_ostream &OS,
                                          const Constant *ImmArgVal) {
  static constexpr StringRef TensormapSwizzleAtomicities[] = {
      "16B", "32B", "32B + 8B flip", "64B"};
  if (const auto *CI = dyn_cast<ConstantInt>(ImmArgVal)) {
    uint64_t Val = CI->getZExtValue();
    if (Val <= static_cast<uint64_t>(
                   nvvm::TensormapSwizzleAtomicity::SWIZZLE_ATOMICITY_64B)) {
      OS << TensormapSwizzleAtomicities[Val];
      return;
    }
  }
}

void nvvm::printFPRoundingMode(raw_ostream &OS, const Constant *ImmArgVal) {
  if (isa<ConstantInt>(ImmArgVal))
    OS << nvvm::GetRoundingModeName(nvvm::GetRoundingModeFromImmArg(ImmArgVal));
}

void nvvm::printTensormapFillMode(raw_ostream &OS, const Constant *ImmArgVal) {
  if (const auto *CI = dyn_cast<ConstantInt>(ImmArgVal)) {
    uint64_t Val = CI->getZExtValue();
    OS << (Val == static_cast<uint64_t>(TensormapFillMode::ZERO_FILL)
               ? "Zero fill"
               : "OOB-NaN fill");
    return;
  }
}