DXILIntrinsicExpansion.cpp

Go to the documentation of this file.

//===- DXILIntrinsicExpansion.cpp - Prepare LLVM Module for DXIL encoding--===//
//
// 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
//
//===----------------------------------------------------------------------===//
///
/// \file This file contains DXIL intrinsic expansions for those that don't have
//  opcodes in DirectX Intermediate Language (DXIL).
//===----------------------------------------------------------------------===//
 
#include "DXILIntrinsicExpansion.h"
#include "DirectX.h"
#include "llvm/ADT/APInt.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/CodeGen/Passes.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/IR/IntrinsicsDirectX.h"
#include "llvm/IR/MatrixBuilder.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/Type.h"
#include "llvm/Pass.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/MathExtras.h"
 
#define DEBUG_TYPE "dxil-intrinsic-expansion"
 
using namespace llvm;
 
class DXILIntrinsicExpansionLegacy : public ModulePass {
 
public:
  bool runOnModule(Module &M) override;
  DXILIntrinsicExpansionLegacy() : ModulePass(ID) {}
 
  static char ID; // Pass identification.
};
 
static bool resourceAccessNeeds64BitExpansion(Module *M, Type *OverloadTy,
                                              bool IsRaw) {
  if (IsRaw && M->getTargetTriple().getDXILVersion() > VersionTuple(1, 2))
    return false;
 
  Type *ScalarTy = OverloadTy->getScalarType();
  return ScalarTy->isDoubleTy() || ScalarTy->isIntegerTy(64);
}
 
static Value *expand16BitIsInf(CallInst *Orig) {
  Module *M = Orig->getModule();
  if (M->getTargetTriple().getDXILVersion() >= VersionTuple(1, 9))
    return nullptr;
 
  Value *Val = Orig->getOperand(0);
  Type *ValTy = Val->getType();
  if (!ValTy->getScalarType()->isHalfTy())
    return nullptr;
 
  IRBuilder<> Builder(Orig);
  Type *IType = Type::getInt16Ty(M->getContext());
  Constant *PosInf =
      ValTy->isVectorTy()
          ? ConstantVector::getSplat(
                ElementCount::getFixed(
                    cast<FixedVectorType>(ValTy)->getNumElements()),
                ConstantInt::get(IType, 0x7c00))
          : ConstantInt::get(IType, 0x7c00);
 
  Constant *NegInf =
      ValTy->isVectorTy()
          ? ConstantVector::getSplat(
                ElementCount::getFixed(
                    cast<FixedVectorType>(ValTy)->getNumElements()),
                ConstantInt::get(IType, 0xfc00))
          : ConstantInt::get(IType, 0xfc00);
 
  Value *IVal = Builder.CreateBitCast(Val, PosInf->getType());
  Value *B1 = Builder.CreateICmpEQ(IVal, PosInf);
  Value *B2 = Builder.CreateICmpEQ(IVal, NegInf);
  Value *B3 = Builder.CreateOr(B1, B2);
  return B3;
}
 
static Value *expand16BitIsNaN(CallInst *Orig) {
  Module *M = Orig->getModule();
  if (M->getTargetTriple().getDXILVersion() >= VersionTuple(1, 9))
    return nullptr;
 
  Value *Val = Orig->getOperand(0);
  Type *ValTy = Val->getType();
  if (!ValTy->getScalarType()->isHalfTy())
    return nullptr;
 
  IRBuilder<> Builder(Orig);
  Type *IType = Type::getInt16Ty(M->getContext());
 
  Constant *ExpBitMask =
      ValTy->isVectorTy()
          ? ConstantVector::getSplat(
                ElementCount::getFixed(
                    cast<FixedVectorType>(ValTy)->getNumElements()),
                ConstantInt::get(IType, 0x7c00))
          : ConstantInt::get(IType, 0x7c00);
  Constant *SigBitMask =
      ValTy->isVectorTy()
          ? ConstantVector::getSplat(
                ElementCount::getFixed(
                    cast<FixedVectorType>(ValTy)->getNumElements()),
                ConstantInt::get(IType, 0x3ff))
          : ConstantInt::get(IType, 0x3ff);
 
  Constant *Zero =
      ValTy->isVectorTy()
          ? ConstantVector::getSplat(
                ElementCount::getFixed(
                    cast<FixedVectorType>(ValTy)->getNumElements()),
                ConstantInt::get(IType, 0))
          : ConstantInt::get(IType, 0);
 
  Value *IVal = Builder.CreateBitCast(Val, ExpBitMask->getType());
  Value *Exp = Builder.CreateAnd(IVal, ExpBitMask);
  Value *B1 = Builder.CreateICmpEQ(Exp, ExpBitMask);
 
  Value *Sig = Builder.CreateAnd(IVal, SigBitMask);
  Value *B2 = Builder.CreateICmpNE(Sig, Zero);
  Value *B3 = Builder.CreateAnd(B1, B2);
  return B3;
}
 
static Value *expand16BitIsFinite(CallInst *Orig) {
  Module *M = Orig->getModule();
  if (M->getTargetTriple().getDXILVersion() >= VersionTuple(1, 9))
    return nullptr;
 
  Value *Val = Orig->getOperand(0);
  Type *ValTy = Val->getType();
  if (!ValTy->getScalarType()->isHalfTy())
    return nullptr;
 
  IRBuilder<> Builder(Orig);
  Type *IType = Type::getInt16Ty(M->getContext());
 
  Constant *ExpBitMask =
      ValTy->isVectorTy()
          ? ConstantVector::getSplat(
                ElementCount::getFixed(
                    cast<FixedVectorType>(ValTy)->getNumElements()),
                ConstantInt::get(IType, 0x7c00))
          : ConstantInt::get(IType, 0x7c00);
 
  Value *IVal = Builder.CreateBitCast(Val, ExpBitMask->getType());
  Value *Exp = Builder.CreateAnd(IVal, ExpBitMask);
  Value *B1 = Builder.CreateICmpNE(Exp, ExpBitMask);
  return B1;
}
 
static Value *expand16BitIsNormal(CallInst *Orig) {
  Module *M = Orig->getModule();
  if (M->getTargetTriple().getDXILVersion() >= VersionTuple(1, 9))
    return nullptr;
 
  Value *Val = Orig->getOperand(0);
  Type *ValTy = Val->getType();
  if (!ValTy->getScalarType()->isHalfTy())
    return nullptr;
 
  IRBuilder<> Builder(Orig);
  Type *IType = Type::getInt16Ty(M->getContext());
 
  Constant *ExpBitMask =
      ValTy->isVectorTy()
          ? ConstantVector::getSplat(
                ElementCount::getFixed(
                    cast<FixedVectorType>(ValTy)->getNumElements()),
                ConstantInt::get(IType, 0x7c00))
          : ConstantInt::get(IType, 0x7c00);
  Constant *Zero =
      ValTy->isVectorTy()
          ? ConstantVector::getSplat(
                ElementCount::getFixed(
                    cast<FixedVectorType>(ValTy)->getNumElements()),
                ConstantInt::get(IType, 0))
          : ConstantInt::get(IType, 0);
 
  Value *IVal = Builder.CreateBitCast(Val, ExpBitMask->getType());
  Value *Exp = Builder.CreateAnd(IVal, ExpBitMask);
  Value *NotAllZeroes = Builder.CreateICmpNE(Exp, Zero);
  Value *NotAllOnes = Builder.CreateICmpNE(Exp, ExpBitMask);
  Value *B1 = Builder.CreateAnd(NotAllZeroes, NotAllOnes);
  return B1;
}
 
static bool shouldExpandFloatDotIntrinsic(Function &F) {
  assert(F.getIntrinsicID() == Intrinsic::dx_fdot &&
         "Function is not a dx.fdot intrinsic");
  auto *ParamTy = cast<FixedVectorType>(F.getFunctionType()->getParamType(0));
  return ParamTy->getNumElements() <= 4 ||
         F.getParent()->getTargetTriple().getOSVersion() < VersionTuple(6, 9);
}
 
static bool isIntrinsicExpansion(Function &F) {
  switch (F.getIntrinsicID()) {
  case Intrinsic::assume:
  case Intrinsic::abs:
  case Intrinsic::atan2:
  case Intrinsic::copysign:
  case Intrinsic::fshl:
  case Intrinsic::fshr:
  case Intrinsic::exp:
  case Intrinsic::is_fpclass:
  case Intrinsic::log:
  case Intrinsic::log10:
  case Intrinsic::pow:
  case Intrinsic::powi:
  case Intrinsic::dx_all:
  case Intrinsic::dx_any:
  case Intrinsic::dx_uclamp:
  case Intrinsic::dx_sclamp:
  case Intrinsic::dx_nclamp:
  case Intrinsic::dx_isinf:
  case Intrinsic::dx_isnan:
  case Intrinsic::dx_sdot:
  case Intrinsic::dx_udot:
  case Intrinsic::dx_sign:
  case Intrinsic::usub_sat:
  case Intrinsic::vector_reduce_add:
  case Intrinsic::vector_reduce_fadd:
  case Intrinsic::matrix_multiply:
  case Intrinsic::matrix_transpose:
  case Intrinsic::umul_with_overflow:
  case Intrinsic::smul_with_overflow:
  case Intrinsic::dx_load_input:
  case Intrinsic::dx_store_output:
    return true;
  case Intrinsic::dx_fdot:
    return shouldExpandFloatDotIntrinsic(F);
  case Intrinsic::dx_resource_load_rawbuffer:
    return resourceAccessNeeds64BitExpansion(
        F.getParent(), F.getReturnType()->getStructElementType(0),
        /*IsRaw*/ true);
  case Intrinsic::dx_resource_load_typedbuffer:
    return resourceAccessNeeds64BitExpansion(
        F.getParent(), F.getReturnType()->getStructElementType(0),
        /*IsRaw*/ false);
  case Intrinsic::dx_resource_store_rawbuffer:
    return resourceAccessNeeds64BitExpansion(
        F.getParent(), F.getFunctionType()->getParamType(3), /*IsRaw*/ true);
  case Intrinsic::dx_resource_store_typedbuffer:
    return resourceAccessNeeds64BitExpansion(
        F.getParent(), F.getFunctionType()->getParamType(2), /*IsRaw*/ false);
  }
  return false;
}
 
static Value *expandUsubSat(CallInst *Orig) {
  Value *A = Orig->getArgOperand(0);
  Value *B = Orig->getArgOperand(1);
  Type *Ty = A->getType();
 
  IRBuilder<> Builder(Orig);
 
  Value *Cmp = Builder.CreateICmpULT(A, B, "usub.cmp");
  Value *Sub = Builder.CreateSub(A, B, "usub.sub");
  Value *Zero = ConstantInt::get(Ty, 0);
  return Builder.CreateSelect(Cmp, Zero, Sub, "usub.sat");
}
 
// Compute the high N bits of the 2N-bit unsigned product of two N-bit values
// using only N-bit arithmetic, so we don't introduce a wider integer type that
// may be unsupported in DXIL.
static Value *createMulHighUnsigned(IRBuilder<> &Builder, Value *A, Value *B,
                                    Type *Ty, unsigned BW) {
  assert(BW % 2 == 0 && "high-half split needs symmetric halves");
  unsigned Half = BW / 2;
  Value *HalfShift = ConstantInt::get(Ty, Half);
  Value *LoMask = ConstantInt::get(Ty, APInt::getLowBitsSet(BW, Half));
 
  Value *U0 = Builder.CreateAnd(A, LoMask);
  Value *U1 = Builder.CreateLShr(A, HalfShift);
  Value *V0 = Builder.CreateAnd(B, LoMask);
  Value *V1 = Builder.CreateLShr(B, HalfShift);
 
  Value *W0 = Builder.CreateMul(U0, V0);
  Value *T = Builder.CreateAdd(Builder.CreateMul(U1, V0),
                               Builder.CreateLShr(W0, HalfShift));
  Value *W1 = Builder.CreateAnd(T, LoMask);
  Value *W2 = Builder.CreateLShr(T, HalfShift);
  W1 = Builder.CreateAdd(Builder.CreateMul(U0, V1), W1);
  return Builder.CreateAdd(Builder.CreateAdd(Builder.CreateMul(U1, V1), W2),
                           Builder.CreateLShr(W1, HalfShift));
}
 
// Expand a {u,s}mul.with.overflow intrinsic. The low half of the result is a
// plain multiply; overflow is derived from the high half of the double-width
// product.
static Value *expandMulWithOverflow(CallInst *Orig, bool Signed) {
  IRBuilder<> Builder(Orig);
  Value *A = Orig->getArgOperand(0);
  Value *B = Orig->getArgOperand(1);
  Type *Ty = A->getType();
  unsigned BW = Ty->getScalarSizeInBits();
 
  Value *Lo;
  Value *Ov;
 
  // A plain double-width multiply is simplest, but we avoid it once it would
  // introduce a 64-bit (or wider) integer, which DXIL does not always support.
  // For i32 we use the native DXIL IMul/UMul ops, which return the full product
  // as two i32s; wider types fall back to a same-width high-half computation.
  if (2 * BW <= 32) {
    Lo = Builder.CreateMul(A, B);
    Type *WideTy = Ty->getWithNewBitWidth(2 * BW);
    Value *WideA =
        Signed ? Builder.CreateSExt(A, WideTy) : Builder.CreateZExt(A, WideTy);
    Value *WideB =
        Signed ? Builder.CreateSExt(B, WideTy) : Builder.CreateZExt(B, WideTy);
    Value *Wide = Builder.CreateMul(WideA, WideB);
    if (Signed) {
      // Overflow when the full product doesn't fit back into BW signed bits.
      Ov = Builder.CreateICmpNE(Wide, Builder.CreateSExt(Lo, WideTy));
    } else {
      Value *Hi = Builder.CreateLShr(Wide, ConstantInt::get(WideTy, BW));
      Ov = Builder.CreateICmpNE(Hi, ConstantInt::get(WideTy, 0));
    }
  } else if (BW == 32) {
    // IMul/UMul return {high, low}; index 0 is the high 32 bits.
    Type *ResTy = StructType::get(Ty, Ty);
    Intrinsic::ID IntrinsicID =
        Signed ? Intrinsic::dx_imul : Intrinsic::dx_umul;
    Value *Mul = Builder.CreateIntrinsic(ResTy, IntrinsicID, {A, B});
    Value *Hi = Builder.CreateExtractValue(Mul, 0);
    Lo = Builder.CreateExtractValue(Mul, 1);
    if (Signed)
      Ov = Builder.CreateICmpNE(
          Hi, Builder.CreateAShr(Lo, ConstantInt::get(Ty, BW - 1)));
    else
      Ov = Builder.CreateICmpNE(Hi, ConstantInt::get(Ty, 0));
  } else {
    Lo = Builder.CreateMul(A, B);
    Value *Hi = createMulHighUnsigned(Builder, A, B, Ty, BW);
    if (Signed) {
      // Turn the unsigned high half into the signed one, then overflow means it
      // isn't the sign extension of the low half.
      Value *SignShift = ConstantInt::get(Ty, BW - 1);
      Value *ASign = Builder.CreateAShr(A, SignShift);
      Value *BSign = Builder.CreateAShr(B, SignShift);
      Hi = Builder.CreateSub(Hi, Builder.CreateAnd(ASign, B));
      Hi = Builder.CreateSub(Hi, Builder.CreateAnd(BSign, A));
      Ov = Builder.CreateICmpNE(Hi, Builder.CreateAShr(Lo, SignShift));
    } else {
      Ov = Builder.CreateICmpNE(Hi, ConstantInt::get(Ty, 0));
    }
  }
 
  Value *Agg = PoisonValue::get(Orig->getType());
  Agg = Builder.CreateInsertValue(Agg, Lo, 0);
  return Builder.CreateInsertValue(Agg, Ov, 1);
}
 
static Value *expandVecReduceAdd(CallInst *Orig, Intrinsic::ID IntrinsicId) {
  assert(IntrinsicId == Intrinsic::vector_reduce_add ||
         IntrinsicId == Intrinsic::vector_reduce_fadd);
 
  IRBuilder<> Builder(Orig);
  bool IsFAdd = (IntrinsicId == Intrinsic::vector_reduce_fadd);
 
  Value *X = Orig->getOperand(IsFAdd ? 1 : 0);
  Type *Ty = X->getType();
  auto *XVec = dyn_cast<FixedVectorType>(Ty);
  unsigned XVecSize = XVec->getNumElements();
  Value *Sum = Builder.CreateExtractElement(X, static_cast<uint64_t>(0));
 
  // Handle the initial start value for floating-point addition.
  if (IsFAdd) {
    Constant *StartValue = dyn_cast<Constant>(Orig->getOperand(0));
    if (StartValue && !StartValue->isNullValue())
      Sum = Builder.CreateFAdd(Sum, StartValue);
  }
 
  // Accumulate the remaining vector elements.
  for (unsigned I = 1; I < XVecSize; I++) {
    Value *Elt = Builder.CreateExtractElement(X, I);
    if (IsFAdd)
      Sum = Builder.CreateFAdd(Sum, Elt);
    else
      Sum = Builder.CreateAdd(Sum, Elt);
  }
 
  return Sum;
}
 
static Value *expandAbs(CallInst *Orig) {
  Value *X = Orig->getOperand(0);
  IRBuilder<> Builder(Orig);
  Type *Ty = X->getType();
  Type *EltTy = Ty->getScalarType();
  Constant *Zero = Ty->isVectorTy()
                       ? ConstantVector::getSplat(
                             ElementCount::getFixed(
                                 cast<FixedVectorType>(Ty)->getNumElements()),
                             ConstantInt::get(EltTy, 0))
                       : ConstantInt::get(EltTy, 0);
  auto *V = Builder.CreateSub(Zero, X);
  return Builder.CreateIntrinsic(Ty, Intrinsic::smax, {X, V}, nullptr,
                                 "dx.max");
}
 
// Create a DXIL dot2, dot3, or dot4 for the given operands.
static Value *expandFloatDotChunk(CallInst *Orig, Value *A, Value *B) {
  Type *ATy = A->getType();
  [[maybe_unused]] Type *BTy = B->getType();
  assert(ATy->isVectorTy() && BTy->isVectorTy());
 
  IRBuilder<> Builder(Orig);
 
  auto *AVec = dyn_cast<FixedVectorType>(ATy);
 
  assert(ATy->getScalarType()->isFloatingPointTy());
 
  unsigned NumElts = AVec->getNumElements();
  Intrinsic::ID DotIntrinsic;
  switch (NumElts) {
  case 2:
    DotIntrinsic = Intrinsic::dx_dot2;
    break;
  case 3:
    DotIntrinsic = Intrinsic::dx_dot3;
    break;
  case 4:
    DotIntrinsic = Intrinsic::dx_dot4;
    break;
  default:
    reportFatalUsageError(
        "Invalid dot product input vector: length is outside 2-4");
  }
 
  SmallVector<Value *> Args;
  for (unsigned I = 0; I < NumElts; ++I)
    Args.push_back(Builder.CreateExtractElement(A, Builder.getInt32(I)));
  for (unsigned I = 0; I < NumElts; ++I)
    Args.push_back(Builder.CreateExtractElement(B, Builder.getInt32(I)));
  return Builder.CreateIntrinsic(ATy->getScalarType(), DotIntrinsic, Args,
                                 nullptr, "dot");
}
 
// Expand an arbitrary-width float dot into the minimum number of legal DXIL
// dot2, dot3, and dot4 operations.
static Value *expandFloatDotIntrinsic(CallInst *Orig) {
  Value *A = Orig->getOperand(0);
  Value *B = Orig->getOperand(1);
  unsigned NumElts = cast<FixedVectorType>(A->getType())->getNumElements();
 
  //  We return early here to avoid constructing unnecessary identity shuffles.
  if (NumElts <= 4)
    return expandFloatDotChunk(Orig, A, B);
 
  assert(Orig->getModule()->getTargetTriple().getOSVersion() <
             VersionTuple(6, 9) &&
         "long fdot must not be expanded for shader model 6.9 or later");
 
  IRBuilder<> Builder(Orig);
  Value *Result = nullptr;
  for (unsigned Offset = 0; Offset < NumElts;) {
    unsigned Remaining = NumElts - Offset;
    // Taking four is optimal unless it would leave an illegal one-element
    // tail. In that case, take three and finish with dot2.
    unsigned ChunkSize = Remaining == 5 ? 3 : std::min(Remaining, 4u);
    SmallVector<int, 4> Mask;
    for (unsigned I = 0; I < ChunkSize; ++I)
      Mask.push_back(Offset + I);
    Value *AChunk = Builder.CreateShuffleVector(A, Mask);
    Value *BChunk = Builder.CreateShuffleVector(B, Mask);
    Value *Chunk = expandFloatDotChunk(Orig, AChunk, BChunk);
    Result = Result ? Builder.CreateFAdd(Result, Chunk, "dot.add") : Chunk;
    Offset += ChunkSize;
  }
  return Result;
}
 
// Expand integer dot product to multiply and add ops
static Value *expandIntegerDotIntrinsic(CallInst *Orig,
                                        Intrinsic::ID DotIntrinsic) {
  assert(DotIntrinsic == Intrinsic::dx_sdot ||
         DotIntrinsic == Intrinsic::dx_udot);
  Value *A = Orig->getOperand(0);
  Value *B = Orig->getOperand(1);
  Type *ATy = A->getType();
  [[maybe_unused]] Type *BTy = B->getType();
  assert(ATy->isVectorTy() && BTy->isVectorTy());
 
  IRBuilder<> Builder(Orig);
 
  auto *AVec = dyn_cast<FixedVectorType>(ATy);
 
  assert(ATy->getScalarType()->isIntegerTy());
 
  Value *Result;
  Intrinsic::ID MadIntrinsic = DotIntrinsic == Intrinsic::dx_sdot
                                   ? Intrinsic::dx_imad
                                   : Intrinsic::dx_umad;
  Value *Elt0 = Builder.CreateExtractElement(A, (uint64_t)0);
  Value *Elt1 = Builder.CreateExtractElement(B, (uint64_t)0);
  Result = Builder.CreateMul(Elt0, Elt1);
  for (unsigned I = 1; I < AVec->getNumElements(); I++) {
    Elt0 = Builder.CreateExtractElement(A, I);
    Elt1 = Builder.CreateExtractElement(B, I);
    Result = Builder.CreateIntrinsic(Result->getType(), MadIntrinsic,
                                     ArrayRef<Value *>{Elt0, Elt1, Result},
                                     nullptr, "dx.mad");
  }
  return Result;
}
 
static Value *expandExpIntrinsic(CallInst *Orig) {
  Value *X = Orig->getOperand(0);
  IRBuilder<> Builder(Orig);
  Type *Ty = X->getType();
  Type *EltTy = Ty->getScalarType();
  Constant *Log2eConst =
      Ty->isVectorTy() ? ConstantVector::getSplat(
                             ElementCount::getFixed(
                                 cast<FixedVectorType>(Ty)->getNumElements()),
                             ConstantFP::get(EltTy, numbers::log2ef))
                       : ConstantFP::get(EltTy, numbers::log2ef);
  Value *NewX = Builder.CreateFMul(Log2eConst, X);
  CallInst *Exp2Call = Builder.CreateIntrinsicWithoutFolding(
      Ty, Intrinsic::exp2, {NewX}, nullptr, "dx.exp2");
  Exp2Call->setTailCall(Orig->isTailCall());
  Exp2Call->setAttributes(Orig->getAttributes());
  return Exp2Call;
}
 
static Value *expandIsFPClass(CallInst *Orig) {
  Value *T = Orig->getArgOperand(1);
  auto *TCI = dyn_cast<ConstantInt>(T);
 
  // These FPClassTest cases have DXIL opcodes, so they will be handled in
  // DXIL Op Lowering instead for all non f16 cases.
  switch (TCI->getZExtValue()) {
  case FPClassTest::fcInf:
    return expand16BitIsInf(Orig);
  case FPClassTest::fcNan:
    return expand16BitIsNaN(Orig);
  case FPClassTest::fcNormal:
    return expand16BitIsNormal(Orig);
  case FPClassTest::fcFinite:
    return expand16BitIsFinite(Orig);
  }
 
  IRBuilder<> Builder(Orig);
 
  Value *F = Orig->getArgOperand(0);
  Type *FTy = F->getType();
  unsigned FNumElem = 0; // 0 => F is not a vector
 
  unsigned BitWidth; // Bit width of F or the ElemTy of F
  Type *BitCastTy;   // An IntNTy of the same bitwidth as F or ElemTy of F
 
  if (auto *FVecTy = dyn_cast<FixedVectorType>(FTy)) {
    Type *ElemTy = FVecTy->getElementType();
    FNumElem = FVecTy->getNumElements();
    BitWidth = ElemTy->getPrimitiveSizeInBits();
    BitCastTy = FixedVectorType::get(Builder.getIntNTy(BitWidth), FNumElem);
  } else {
    BitWidth = FTy->getPrimitiveSizeInBits();
    BitCastTy = Builder.getIntNTy(BitWidth);
  }
 
  Value *FBitCast = Builder.CreateBitCast(F, BitCastTy);
  switch (TCI->getZExtValue()) {
  case FPClassTest::fcNegZero: {
    Value *NegZero =
        ConstantInt::get(Builder.getIntNTy(BitWidth), 1 << (BitWidth - 1),
                         /*IsSigned=*/true);
    Value *RetVal;
    if (FNumElem) {
      Value *NegZeroSplat = Builder.CreateVectorSplat(FNumElem, NegZero);
      RetVal =
          Builder.CreateICmpEQ(FBitCast, NegZeroSplat, "is.fpclass.negzero");
    } else
      RetVal = Builder.CreateICmpEQ(FBitCast, NegZero, "is.fpclass.negzero");
    return RetVal;
  }
  default:
    reportFatalUsageError("Unsupported FPClassTest");
  }
}
 
static Value *expandAnyOrAllIntrinsic(CallInst *Orig,
                                      Intrinsic::ID IntrinsicId) {
  Value *X = Orig->getOperand(0);
  IRBuilder<> Builder(Orig);
  Type *Ty = X->getType();
  Type *EltTy = Ty->getScalarType();
 
  auto ApplyOp = [&Builder](Intrinsic::ID IntrinsicId, Value *Result,
                            Value *Elt) {
    if (IntrinsicId == Intrinsic::dx_any)
      return Builder.CreateOr(Result, Elt);
    assert(IntrinsicId == Intrinsic::dx_all);
    return Builder.CreateAnd(Result, Elt);
  };
 
  Value *Result = nullptr;
  if (!Ty->isVectorTy()) {
    Result = EltTy->isFloatingPointTy()
                 ? Builder.CreateFCmpUNE(X, ConstantFP::get(EltTy, 0))
                 : Builder.CreateICmpNE(X, ConstantInt::get(EltTy, 0));
  } else {
    auto *XVec = dyn_cast<FixedVectorType>(Ty);
    Value *Cond =
        EltTy->isFloatingPointTy()
            ? Builder.CreateFCmpUNE(
                  X, ConstantVector::getSplat(
                         ElementCount::getFixed(XVec->getNumElements()),
                         ConstantFP::get(EltTy, 0)))
            : Builder.CreateICmpNE(
                  X, ConstantVector::getSplat(
                         ElementCount::getFixed(XVec->getNumElements()),
                         ConstantInt::get(EltTy, 0)));
    Result = Builder.CreateExtractElement(Cond, (uint64_t)0);
    for (unsigned I = 1; I < XVec->getNumElements(); I++) {
      Value *Elt = Builder.CreateExtractElement(Cond, I);
      Result = ApplyOp(IntrinsicId, Result, Elt);
    }
  }
  return Result;
}
 
static Value *expandLogIntrinsic(CallInst *Orig,
                                 float LogConstVal = numbers::ln2f) {
  Value *X = Orig->getOperand(0);
  IRBuilder<> Builder(Orig);
  Type *Ty = X->getType();
  Type *EltTy = Ty->getScalarType();
  Constant *Ln2Const =
      Ty->isVectorTy() ? ConstantVector::getSplat(
                             ElementCount::getFixed(
                                 cast<FixedVectorType>(Ty)->getNumElements()),
                             ConstantFP::get(EltTy, LogConstVal))
                       : ConstantFP::get(EltTy, LogConstVal);
  CallInst *Log2Call = Builder.CreateIntrinsicWithoutFolding(
      Ty, Intrinsic::log2, {X}, nullptr, "elt.log2");
  Log2Call->setTailCall(Orig->isTailCall());
  Log2Call->setAttributes(Orig->getAttributes());
  return Builder.CreateFMul(Ln2Const, Log2Call);
}
static Value *expandLog10Intrinsic(CallInst *Orig) {
  return expandLogIntrinsic(Orig, numbers::ln2f / numbers::ln10f);
}
 
static Value *expandAtan2Intrinsic(CallInst *Orig) {
  Value *Y = Orig->getOperand(0);
  Value *X = Orig->getOperand(1);
  Type *Ty = X->getType();
  IRBuilder<> Builder(Orig);
  Builder.setFastMathFlags(Orig->getFastMathFlags());
 
  Value *Tan = Builder.CreateFDiv(Y, X);
 
  CallInst *Atan = Builder.CreateIntrinsicWithoutFolding(
      Ty, Intrinsic::atan, {Tan}, nullptr, "Elt.Atan");
  Atan->setTailCall(Orig->isTailCall());
  Atan->setAttributes(Orig->getAttributes());
 
  // Modify atan result based on https://en.wikipedia.org/wiki/Atan2.
  Constant *Pi = ConstantFP::get(Ty, llvm::numbers::pi);
  Constant *HalfPi = ConstantFP::get(Ty, llvm::numbers::pi / 2);
  Constant *NegHalfPi = ConstantFP::get(Ty, -llvm::numbers::pi / 2);
  Constant *Zero = ConstantFP::get(Ty, 0);
  Value *AtanAddPi = Builder.CreateFAdd(Atan, Pi);
  Value *AtanSubPi = Builder.CreateFSub(Atan, Pi);
 
  // x > 0 -> atan.
  Value *Result = Atan;
  Value *XLt0 = Builder.CreateFCmpOLT(X, Zero);
  Value *XEq0 = Builder.CreateFCmpOEQ(X, Zero);
  Value *YGe0 = Builder.CreateFCmpOGE(Y, Zero);
  Value *YLt0 = Builder.CreateFCmpOLT(Y, Zero);
 
  // x < 0, y >= 0 -> atan + pi.
  Value *XLt0AndYGe0 = Builder.CreateAnd(XLt0, YGe0);
  Result = Builder.CreateSelect(XLt0AndYGe0, AtanAddPi, Result);
 
  // x < 0, y < 0 -> atan - pi.
  Value *XLt0AndYLt0 = Builder.CreateAnd(XLt0, YLt0);
  Result = Builder.CreateSelect(XLt0AndYLt0, AtanSubPi, Result);
 
  // x == 0, y < 0 -> -pi/2
  Value *XEq0AndYLt0 = Builder.CreateAnd(XEq0, YLt0);
  Result = Builder.CreateSelect(XEq0AndYLt0, NegHalfPi, Result);
 
  // x == 0, y > 0 -> pi/2
  Value *XEq0AndYGe0 = Builder.CreateAnd(XEq0, YGe0);
  Result = Builder.CreateSelect(XEq0AndYGe0, HalfPi, Result);
 
  return Result;
}
 
template <bool LeftFunnel>
static Value *expandFunnelShiftIntrinsic(CallInst *Orig) {
  Type *Ty = Orig->getType();
  Value *A = Orig->getOperand(0);
  Value *B = Orig->getOperand(1);
  Value *Shift = Orig->getOperand(2);
 
  IRBuilder<> Builder(Orig);
 
  unsigned BitWidth = Ty->getScalarSizeInBits();
  assert(llvm::isPowerOf2_32(BitWidth) &&
         "Can't use Mask to compute modulo and inverse");
 
  // Note: if (Shift % BitWidth) == 0 then (BitWidth - Shift) == BitWidth,
  // shifting by the bitwidth for shl/lshr returns a poisoned result. As such,
  // we implement the same formula as LegalizerHelper::lowerFunnelShiftAsShifts.
  //
  // The funnel shift is expanded like so:
  // fshl
  //  -> msb_extract((concat(A, B) << (Shift % BitWidth)), BitWidth)
  //  -> A << (Shift % BitWidth) | B >> 1 >> (BitWidth - 1 - (Shift % BitWidth))
  // fshr
  //  -> lsb_extract((concat(A, B) >> (Shift % BitWidth), BitWidth))
  //  -> A << 1 << (BitWidth - 1 - (Shift % BitWidth)) | B >> (Shift % BitWidth)
 
  // (BitWidth - 1) -> Mask
  Constant *Mask = ConstantInt::get(Ty, Ty->getScalarSizeInBits() - 1);
 
  // Shift % BitWidth
  //  -> Shift & (BitWidth - 1)
  //  -> Shift & Mask
  Value *MaskedShift = Builder.CreateAnd(Shift, Mask);
 
  // (BitWidth - 1) - (Shift % BitWidth)
  //  -> ~Shift & (BitWidth - 1)
  //  -> ~Shift & Mask
  Value *NotShift = Builder.CreateNot(Shift);
  Value *InverseShift = Builder.CreateAnd(NotShift, Mask);
 
  Constant *One = ConstantInt::get(Ty, 1);
  Value *ShiftedA;
  Value *ShiftedB;
 
  if (LeftFunnel) {
    ShiftedA = Builder.CreateShl(A, MaskedShift);
    Value *ShiftB1 = Builder.CreateLShr(B, One);
    ShiftedB = Builder.CreateLShr(ShiftB1, InverseShift);
  } else {
    Value *ShiftA1 = Builder.CreateShl(A, One);
    ShiftedA = Builder.CreateShl(ShiftA1, InverseShift);
    ShiftedB = Builder.CreateLShr(B, MaskedShift);
  }
 
  Value *Result = Builder.CreateOr(ShiftedA, ShiftedB);
  return Result;
}
 
static Value *expandPowIntrinsic(CallInst *Orig, Intrinsic::ID IntrinsicId) {
 
  Value *X = Orig->getOperand(0);
  Value *Y = Orig->getOperand(1);
  Type *Ty = X->getType();
  IRBuilder<> Builder(Orig);
 
  if (IntrinsicId == Intrinsic::powi)
    Y = Builder.CreateSIToFP(Y, Ty);
 
  Value *Log2Call =
      Builder.CreateIntrinsic(Ty, Intrinsic::log2, {X}, nullptr, "elt.log2");
  auto *Mul = Builder.CreateFMul(Log2Call, Y);
  CallInst *Exp2Call = Builder.CreateIntrinsicWithoutFolding(
      Ty, Intrinsic::exp2, {Mul}, nullptr, "elt.exp2");
  Exp2Call->setTailCall(Orig->isTailCall());
  Exp2Call->setAttributes(Orig->getAttributes());
  return Exp2Call;
}
 
static bool expandBufferLoadIntrinsic(CallInst *Orig, bool IsRaw) {
  IRBuilder<> Builder(Orig);
 
  Type *BufferTy = Orig->getType()->getStructElementType(0);
  Type *ScalarTy = BufferTy->getScalarType();
  bool IsDouble = ScalarTy->isDoubleTy();
  assert(IsDouble || ScalarTy->isIntegerTy(64) &&
                         "Only expand double or int64 scalars or vectors");
  bool IsVector = false;
  unsigned ExtractNum = 2;
  if (auto *VT = dyn_cast<FixedVectorType>(BufferTy)) {
    ExtractNum = 2 * VT->getNumElements();
    IsVector = true;
    assert(IsRaw || ExtractNum == 4 && "TypedBufferLoad vector must be size 2");
  }
 
  SmallVector<Value *, 2> Loads;
  Value *Result = PoisonValue::get(BufferTy);
  unsigned Base = 0;
  // If we need to extract more than 4 i32; we need to break it up into
  // more than one load. LoadNum tells us how many i32s we are loading in
  // each load
  while (ExtractNum > 0) {
    unsigned LoadNum = std::min(ExtractNum, 4u);
    Type *Ty = VectorType::get(Builder.getInt32Ty(), LoadNum, false);
 
    Type *LoadType = StructType::get(Ty, Builder.getInt1Ty());
    Intrinsic::ID LoadIntrinsic = Intrinsic::dx_resource_load_typedbuffer;
    SmallVector<Value *, 3> Args = {Orig->getOperand(0), Orig->getOperand(1)};
    if (IsRaw) {
      LoadIntrinsic = Intrinsic::dx_resource_load_rawbuffer;
      Value *Tmp = Builder.getInt32(4 * Base * 2);
      Args.push_back(Builder.CreateAdd(Orig->getOperand(2), Tmp));
    }
 
    Value *Load = Builder.CreateIntrinsic(LoadType, LoadIntrinsic, Args);
    Loads.push_back(Load);
 
    // extract the buffer load's result
    Value *Extract = Builder.CreateExtractValue(Load, {0});
 
    SmallVector<Value *> ExtractElements;
    for (unsigned I = 0; I < LoadNum; ++I)
      ExtractElements.push_back(
          Builder.CreateExtractElement(Extract, Builder.getInt32(I)));
 
    // combine into double(s) or int64(s)
    for (unsigned I = 0; I < LoadNum; I += 2) {
      Value *Combined = nullptr;
      if (IsDouble)
        // For doubles, use dx_asdouble intrinsic
        Combined = Builder.CreateIntrinsic(
            Builder.getDoubleTy(), Intrinsic::dx_asdouble,
            {ExtractElements[I], ExtractElements[I + 1]});
      else {
        // For int64, manually combine two int32s
        // First, zero-extend both values to i64
        Value *Lo =
            Builder.CreateZExt(ExtractElements[I], Builder.getInt64Ty());
        Value *Hi =
            Builder.CreateZExt(ExtractElements[I + 1], Builder.getInt64Ty());
        // Shift the high bits left by 32 bits
        Value *ShiftedHi = Builder.CreateShl(Hi, Builder.getInt64(32));
        // OR the high and low bits together
        Combined = Builder.CreateOr(Lo, ShiftedHi);
      }
 
      if (IsVector)
        Result = Builder.CreateInsertElement(Result, Combined,
                                             Builder.getInt32((I / 2) + Base));
      else
        Result = Combined;
    }
 
    ExtractNum -= LoadNum;
    Base += LoadNum / 2;
  }
 
  Value *CheckBit = nullptr;
  for (User *U : make_early_inc_range(Orig->users())) {
    // If it's not a ExtractValueInst, we don't know how to
    // handle it
    auto *EVI = dyn_cast<ExtractValueInst>(U);
    if (!EVI)
      llvm_unreachable("Unexpected user of typedbufferload");
 
    ArrayRef<unsigned> Indices = EVI->getIndices();
    assert(Indices.size() == 1);
 
    if (Indices[0] == 0) {
      // Use of the value(s)
      EVI->replaceAllUsesWith(Result);
    } else {
      // Use of the check bit
      assert(Indices[0] == 1 && "Unexpected type for typedbufferload");
      // Note: This does not always match the historical behaviour of DXC.
      // See https://github.com/microsoft/DirectXShaderCompiler/issues/7622
      if (!CheckBit) {
        SmallVector<Value *, 2> CheckBits;
        for (Value *L : Loads)
          CheckBits.push_back(Builder.CreateExtractValue(L, {1}));
        CheckBit = Builder.CreateAnd(CheckBits);
      }
      EVI->replaceAllUsesWith(CheckBit);
    }
    EVI->eraseFromParent();
  }
  Orig->eraseFromParent();
  return true;
}
 
static bool expandBufferStoreIntrinsic(CallInst *Orig, bool IsRaw) {
  IRBuilder<> Builder(Orig);
 
  unsigned ValIndex = IsRaw ? 3 : 2;
  Type *BufferTy = Orig->getFunctionType()->getParamType(ValIndex);
  Type *ScalarTy = BufferTy->getScalarType();
  bool IsDouble = ScalarTy->isDoubleTy();
  assert((IsDouble || ScalarTy->isIntegerTy(64)) &&
         "Only expand double or int64 scalars or vectors");
 
  // Determine if we're dealing with a vector or scalar
  bool IsVector = false;
  unsigned ExtractNum = 2;
  unsigned VecLen = 0;
  if (auto *VT = dyn_cast<FixedVectorType>(BufferTy)) {
    VecLen = VT->getNumElements();
    assert(IsRaw || VecLen == 2 && "TypedBufferStore vector must be size 2");
    ExtractNum = VecLen * 2;
    IsVector = true;
  }
 
  // Create the appropriate vector type for the result
  Type *Int32Ty = Builder.getInt32Ty();
  Type *ResultTy = VectorType::get(Int32Ty, ExtractNum, false);
  Value *Val = PoisonValue::get(ResultTy);
 
  Type *SplitElementTy = Int32Ty;
  if (IsVector)
    SplitElementTy = VectorType::get(SplitElementTy, VecLen, false);
 
  Value *LowBits = nullptr;
  Value *HighBits = nullptr;
  // Split the 64-bit values into 32-bit components
  if (IsDouble) {
    auto *SplitTy = llvm::StructType::get(SplitElementTy, SplitElementTy);
    Value *Split = Builder.CreateIntrinsic(SplitTy, Intrinsic::dx_splitdouble,
                                           {Orig->getOperand(ValIndex)});
    LowBits = Builder.CreateExtractValue(Split, 0);
    HighBits = Builder.CreateExtractValue(Split, 1);
  } else {
    // Handle int64 type(s)
    Value *InputVal = Orig->getOperand(ValIndex);
    Constant *ShiftAmt = Builder.getInt64(32);
    if (IsVector)
      ShiftAmt =
          ConstantVector::getSplat(ElementCount::getFixed(VecLen), ShiftAmt);
 
    // Split into low and high 32-bit parts
    LowBits = Builder.CreateTrunc(InputVal, SplitElementTy);
    Value *ShiftedVal = Builder.CreateLShr(InputVal, ShiftAmt);
    HighBits = Builder.CreateTrunc(ShiftedVal, SplitElementTy);
  }
 
  if (IsVector) {
    SmallVector<int, 8> Mask;
    for (unsigned I = 0; I < VecLen; ++I) {
      Mask.push_back(I);
      Mask.push_back(I + VecLen);
    }
    Val = Builder.CreateShuffleVector(LowBits, HighBits, Mask);
  } else {
    Val = Builder.CreateInsertElement(Val, LowBits, Builder.getInt32(0));
    Val = Builder.CreateInsertElement(Val, HighBits, Builder.getInt32(1));
  }
 
  // If we need to extract more than 4 i32; we need to break it up into
  // more than one store. StoreNum tells us how many i32s we are storing in
  // each store
  unsigned Base = 0;
  while (ExtractNum > 0) {
    unsigned StoreNum = std::min(ExtractNum, 4u);
 
    Intrinsic::ID StoreIntrinsic = Intrinsic::dx_resource_store_typedbuffer;
    SmallVector<Value *, 4> Args = {Orig->getOperand(0), Orig->getOperand(1)};
    if (IsRaw) {
      StoreIntrinsic = Intrinsic::dx_resource_store_rawbuffer;
      Value *Tmp = Builder.getInt32(4 * Base);
      Args.push_back(Builder.CreateAdd(Orig->getOperand(2), Tmp));
    }
 
    SmallVector<int, 4> Mask;
    for (unsigned I = 0; I < StoreNum; ++I) {
      Mask.push_back(Base + I);
    }
 
    Value *SubVal = Val;
    if (VecLen > 2)
      SubVal = Builder.CreateShuffleVector(Val, Mask);
 
    Args.push_back(SubVal);
    // Create the final intrinsic call
    Builder.CreateIntrinsic(Builder.getVoidTy(), StoreIntrinsic, Args);
 
    ExtractNum -= StoreNum;
    Base += StoreNum;
  }
  Orig->eraseFromParent();
  return true;
}
 
static Intrinsic::ID getMaxForClamp(Intrinsic::ID ClampIntrinsic) {
  if (ClampIntrinsic == Intrinsic::dx_uclamp)
    return Intrinsic::umax;
  if (ClampIntrinsic == Intrinsic::dx_sclamp)
    return Intrinsic::smax;
  assert(ClampIntrinsic == Intrinsic::dx_nclamp);
  return Intrinsic::maxnum;
}
 
static Intrinsic::ID getMinForClamp(Intrinsic::ID ClampIntrinsic) {
  if (ClampIntrinsic == Intrinsic::dx_uclamp)
    return Intrinsic::umin;
  if (ClampIntrinsic == Intrinsic::dx_sclamp)
    return Intrinsic::smin;
  assert(ClampIntrinsic == Intrinsic::dx_nclamp);
  return Intrinsic::minnum;
}
 
static Value *expandClampIntrinsic(CallInst *Orig,
                                   Intrinsic::ID ClampIntrinsic) {
  Value *X = Orig->getOperand(0);
  Value *Min = Orig->getOperand(1);
  Value *Max = Orig->getOperand(2);
  Type *Ty = X->getType();
  IRBuilder<> Builder(Orig);
  auto *MaxCall = Builder.CreateIntrinsic(Ty, getMaxForClamp(ClampIntrinsic),
                                          {X, Min}, nullptr, "dx.max");
  return Builder.CreateIntrinsic(Ty, getMinForClamp(ClampIntrinsic),
                                 {MaxCall, Max}, nullptr, "dx.min");
}
 
static Value *expandSignIntrinsic(CallInst *Orig) {
  Value *X = Orig->getOperand(0);
  Type *Ty = X->getType();
  Type *ScalarTy = Ty->getScalarType();
  Type *RetTy = Orig->getType();
  Constant *Zero = Constant::getNullValue(Ty);
 
  IRBuilder<> Builder(Orig);
 
  Value *GT;
  Value *LT;
  if (ScalarTy->isFloatingPointTy()) {
    GT = Builder.CreateFCmpOLT(Zero, X);
    LT = Builder.CreateFCmpOLT(X, Zero);
  } else {
    assert(ScalarTy->isIntegerTy());
    GT = Builder.CreateICmpSLT(Zero, X);
    LT = Builder.CreateICmpSLT(X, Zero);
  }
 
  Value *ZextGT = Builder.CreateZExt(GT, RetTy);
  Value *ZextLT = Builder.CreateZExt(LT, RetTy);
 
  return Builder.CreateSub(ZextGT, ZextLT);
}
 
// Expand llvm.copysign by combining the sign bit with the magnitude bits using
// bitwise operations.
static Value *expandCopySignIntrinsic(CallInst *Orig) {
  Value *Magnitude = Orig->getOperand(0);
  Value *Sign = Orig->getOperand(1);
  Type *Ty = Orig->getType();
 
  IRBuilder<> Builder(Orig);
 
  bool IsDouble = Ty->getScalarType()->isDoubleTy();
  unsigned BitWidth = IsDouble ? 32 : Ty->getScalarSizeInBits();
  Type *IntTy = Ty->getWithNewType(Builder.getIntNTy(BitWidth));
 
  auto CopySignBit = [&](Value *MagnitudeInt, Value *SignInt) {
    APInt SignMaskVal = APInt::getSignMask(BitWidth);
    // `ConstantInt::get` broadcasts to a splat when `IntTy` is a vector.
    Constant *SignMask = ConstantInt::get(IntTy, SignMaskVal);
    Constant *NotSignMask = ConstantInt::get(IntTy, ~SignMaskVal);
 
    Value *MagnitudeBits = Builder.CreateAnd(MagnitudeInt, NotSignMask);
    Value *SignBits = Builder.CreateAnd(SignInt, SignMask);
    return Builder.CreateOr(MagnitudeBits, SignBits);
  };
 
  // Avoid i64 bitwise ops, which require the Int64Ops shader feature.
  if (IsDouble) {
    auto *SplitTy = StructType::get(IntTy, IntTy);
    Value *MagnitudeHalves = Builder.CreateIntrinsic(
        SplitTy, Intrinsic::dx_splitdouble, {Magnitude});
    Value *SignHalves =
        Builder.CreateIntrinsic(SplitTy, Intrinsic::dx_splitdouble, {Sign});
    Value *MagnitudeLow = Builder.CreateExtractValue(MagnitudeHalves, 0);
    Value *MagnitudeHigh = Builder.CreateExtractValue(MagnitudeHalves, 1);
    Value *SignHigh = Builder.CreateExtractValue(SignHalves, 1);
 
    Value *CombinedHigh = CopySignBit(MagnitudeHigh, SignHigh);
    return Builder.CreateIntrinsic(Ty, Intrinsic::dx_asdouble,
                                   {MagnitudeLow, CombinedHigh});
  }
 
  Value *MagnitudeInt = Builder.CreateBitCast(Magnitude, IntTy);
  Value *SignInt = Builder.CreateBitCast(Sign, IntTy);
  Value *CombinedInt = CopySignBit(MagnitudeInt, SignInt);
  return Builder.CreateBitCast(CombinedInt, Ty);
}
 
// Expand llvm.matrix.multiply by extracting row/column vectors and computing
// dot products.
// Result[r,c] = dot(row_r(LHS), col_c(RHS))
// Element (r,c) is at index c*NumRows + r (column-major).
static Value *expandMatrixMultiply(CallInst *Orig) {
  Value *LHS = Orig->getArgOperand(0);
  Value *RHS = Orig->getArgOperand(1);
  unsigned LHSRows = cast<ConstantInt>(Orig->getArgOperand(2))->getZExtValue();
  unsigned LHSCols = cast<ConstantInt>(Orig->getArgOperand(3))->getZExtValue();
  unsigned RHSCols = cast<ConstantInt>(Orig->getArgOperand(4))->getZExtValue();
 
  auto *RetTy = cast<FixedVectorType>(Orig->getType());
  Type *EltTy = RetTy->getElementType();
  bool IsFP = EltTy->isFloatingPointTy();
 
  IRBuilder<> Builder(Orig);
 
  // Column-major indexing:
  // LHS row R, element K: index = K * LHSRows + R
  // RHS col C, element K: index = C * LHSCols + K
  Value *Result = PoisonValue::get(RetTy);
 
  // Extract all scalar elements from LHS and RHS once, then reuse them.
  unsigned LHSSize = LHSRows * LHSCols;
  unsigned RHSSize = LHSCols * RHSCols;
  SmallVector<Value *, 16> LHSElts(LHSSize);
  SmallVector<Value *, 16> RHSElts(RHSSize);
  for (unsigned I = 0; I < LHSSize; ++I)
    LHSElts[I] = Builder.CreateExtractElement(LHS, I);
  for (unsigned I = 0; I < RHSSize; ++I)
    RHSElts[I] = Builder.CreateExtractElement(RHS, I);
 
  // Choose the appropriate scalar-arg dot intrinsic for floats.
  // K=1 and double types use scalar expansion instead.
  Intrinsic::ID FloatDotID = Intrinsic::not_intrinsic;
  bool UseScalarFP = IsFP && (EltTy->isDoubleTy() || LHSCols == 1);
  if (IsFP && !UseScalarFP) {
    switch (LHSCols) {
    case 2:
      FloatDotID = Intrinsic::dx_dot2;
      break;
    case 3:
      FloatDotID = Intrinsic::dx_dot3;
      break;
    case 4:
      FloatDotID = Intrinsic::dx_dot4;
      break;
    default:
      reportFatalUsageError(
          "Invalid matrix inner dimension for dot product: must be 2-4");
      return nullptr;
    }
  }
 
  for (unsigned C = 0; C < RHSCols; ++C) {
    for (unsigned R = 0; R < LHSRows; ++R) {
      // Gather row R from LHS and column C from RHS.
      SmallVector<Value *, 4> RowElts, ColElts;
      for (unsigned K = 0; K < LHSCols; ++K) {
        RowElts.push_back(LHSElts[K * LHSRows + R]);
        ColElts.push_back(RHSElts[C * LHSCols + K]);
      }
 
      Value *Dot;
      if (UseScalarFP) {
        // Scalar fmul+fmuladd expansion for double types and K=1.
        Dot = Builder.CreateFMul(RowElts[0], ColElts[0]);
        for (unsigned K = 1; K < LHSCols; ++K)
          Dot = Builder.CreateIntrinsic(EltTy, Intrinsic::fmuladd,
                                        {RowElts[K], ColElts[K], Dot});
      } else if (IsFP) {
        // Emit scalar-arg DXIL dot directly (dx.dot2/dx.dot3/dx.dot4).
        SmallVector<Value *, 8> Args;
        Args.append(RowElts.begin(), RowElts.end());
        Args.append(ColElts.begin(), ColElts.end());
        Dot = Builder.CreateIntrinsic(EltTy, FloatDotID, Args);
      } else {
        // Integer: emit multiply + imad chain.
        Dot = Builder.CreateMul(RowElts[0], ColElts[0]);
        for (unsigned K = 1; K < LHSCols; ++K)
          Dot = Builder.CreateIntrinsic(EltTy, Intrinsic::dx_imad,
                                        {RowElts[K], ColElts[K], Dot});
      }
      unsigned ResIdx = C * LHSRows + R;
      Result = Builder.CreateInsertElement(Result, Dot, ResIdx);
    }
  }
  return Result;
}
 
// Expand llvm.matrix.transpose as a shufflevector that permutes elements
// from column-major source to column-major transposed layout.
// Element (r,c) at index c*Rows + r moves to index r*Cols + c.
static Value *expandMatrixTranspose(CallInst *Orig) {
  Value *Mat = Orig->getArgOperand(0);
  unsigned Rows = cast<ConstantInt>(Orig->getArgOperand(1))->getZExtValue();
  unsigned Cols = cast<ConstantInt>(Orig->getArgOperand(2))->getZExtValue();
 
  unsigned NumElts = Rows * Cols;
  SmallVector<int, 16> Mask(NumElts);
  for (unsigned I = 0; I < NumElts; ++I)
    Mask[I] = (I % Cols) * Rows + (I / Cols);
 
  IRBuilder<> Builder(Orig);
  return Builder.CreateShuffleVector(Mat, Mask);
}
 
// Scalarize a vector int_dx_store_output call into per-component scalar calls.
// The DXIL StoreOutput op is per-component; vector intrinsics are split here
// so that DXILOpLowering sees only scalar variants.
static bool expandStoreOutput(CallInst *Orig) {
  auto *VT = dyn_cast<FixedVectorType>(Orig->getArgOperand(3)->getType());
  if (!VT)
    return false; // already scalar, nothing to expand
 
  IRBuilder<> Builder(Orig);
  Module *M = Orig->getModule();
  Type *Int8Ty = Builder.getInt8Ty();
  Type *Int32Ty = Builder.getInt32Ty();
  Type *ScalarTy = VT->getElementType();
  unsigned NumElems = VT->getNumElements();
 
  Value *SigElementId = Orig->getArgOperand(0);
  Value *RowIndex = Orig->getArgOperand(1);
  Value *StartCol = Orig->getArgOperand(2); // i8
  Value *Data = Orig->getArgOperand(3);
  Value *StartColI32 = Builder.CreateZExt(StartCol, Int32Ty);
 
  Function *ScalarFn = Intrinsic::getOrInsertDeclaration(
      M, Intrinsic::dx_store_output, {ScalarTy});
 
  for (unsigned I = 0; I < NumElems; ++I) {
    Value *Scalar =
        Builder.CreateExtractElement(Data, ConstantInt::get(Int32Ty, I));
    Value *ColIdx =
        Builder.CreateAdd(StartColI32, ConstantInt::get(Int32Ty, I));
    Value *ColI8 = Builder.CreateTrunc(ColIdx, Int8Ty);
    Builder.CreateCall(ScalarFn, {SigElementId, RowIndex, ColI8, Scalar});
  }
 
  Orig->eraseFromParent();
  return true;
}
 
// Scalarize a vector int_dx_load_input call into per-component scalar calls
// and reassemble the vector. The DXIL LoadInput op is per-component.
static Value *expandLoadInput(CallInst *Orig) {
  auto *VT = dyn_cast<FixedVectorType>(Orig->getType());
  if (!VT)
    return nullptr; // already scalar, nothing to expand
 
  IRBuilder<> Builder(Orig);
  Module *M = Orig->getModule();
  Type *Int8Ty = Builder.getInt8Ty();
  Type *Int32Ty = Builder.getInt32Ty();
  Type *ScalarTy = VT->getElementType();
  unsigned NumElems = VT->getNumElements();
 
  Value *SigElementId = Orig->getArgOperand(0);
  Value *RowIndex = Orig->getArgOperand(1);
  Value *StartCol = Orig->getArgOperand(2); // i8
  Value *GsVertexOrPrimIndex = Orig->getArgOperand(3);
  Value *StartColI32 = Builder.CreateZExt(StartCol, Int32Ty);
 
  Function *ScalarFn = Intrinsic::getOrInsertDeclaration(
      M, Intrinsic::dx_load_input, {ScalarTy});
 
  Value *Vec = PoisonValue::get(VT);
  for (unsigned I = 0; I < NumElems; ++I) {
    Value *ColIdx =
        Builder.CreateAdd(StartColI32, ConstantInt::get(Int32Ty, I));
    Value *ColI8 = Builder.CreateTrunc(ColIdx, Int8Ty);
    Value *Scalar = Builder.CreateCall(
        ScalarFn, {SigElementId, RowIndex, ColI8, GsVertexOrPrimIndex});
    Vec =
        Builder.CreateInsertElement(Vec, Scalar, ConstantInt::get(Int32Ty, I));
  }
 
  return Vec;
}
 
static bool expandIntrinsic(Function &F, CallInst *Orig) {
  Value *Result = nullptr;
  Intrinsic::ID IntrinsicId = F.getIntrinsicID();
  switch (IntrinsicId) {
  case Intrinsic::abs:
    Result = expandAbs(Orig);
    break;
  case Intrinsic::assume:
    Orig->eraseFromParent();
    return true;
  case Intrinsic::atan2:
    Result = expandAtan2Intrinsic(Orig);
    break;
  case Intrinsic::copysign:
    Result = expandCopySignIntrinsic(Orig);
    break;
  case Intrinsic::fshl:
    Result = expandFunnelShiftIntrinsic<true>(Orig);
    break;
  case Intrinsic::fshr:
    Result = expandFunnelShiftIntrinsic<false>(Orig);
    break;
  case Intrinsic::exp:
    Result = expandExpIntrinsic(Orig);
    break;
  case Intrinsic::is_fpclass:
    Result = expandIsFPClass(Orig);
    break;
  case Intrinsic::log:
    Result = expandLogIntrinsic(Orig);
    break;
  case Intrinsic::log10:
    Result = expandLog10Intrinsic(Orig);
    break;
  case Intrinsic::pow:
  case Intrinsic::powi:
    Result = expandPowIntrinsic(Orig, IntrinsicId);
    break;
  case Intrinsic::dx_all:
  case Intrinsic::dx_any:
    Result = expandAnyOrAllIntrinsic(Orig, IntrinsicId);
    break;
  case Intrinsic::dx_uclamp:
  case Intrinsic::dx_sclamp:
  case Intrinsic::dx_nclamp:
    Result = expandClampIntrinsic(Orig, IntrinsicId);
    break;
  case Intrinsic::dx_isinf:
    Result = expand16BitIsInf(Orig);
    break;
  case Intrinsic::dx_isnan:
    Result = expand16BitIsNaN(Orig);
    break;
  case Intrinsic::dx_fdot:
    Result = expandFloatDotIntrinsic(Orig);
    break;
  case Intrinsic::dx_sdot:
  case Intrinsic::dx_udot:
    Result = expandIntegerDotIntrinsic(Orig, IntrinsicId);
    break;
  case Intrinsic::dx_sign:
    Result = expandSignIntrinsic(Orig);
    break;
  case Intrinsic::dx_load_input:
    Result = expandLoadInput(Orig);
    break;
  case Intrinsic::dx_store_output:
    if (expandStoreOutput(Orig))
      return true;
    break;
  case Intrinsic::dx_resource_load_rawbuffer:
    if (expandBufferLoadIntrinsic(Orig, /*IsRaw*/ true))
      return true;
    break;
  case Intrinsic::dx_resource_store_rawbuffer:
    if (expandBufferStoreIntrinsic(Orig, /*IsRaw*/ true))
      return true;
    break;
  case Intrinsic::dx_resource_load_typedbuffer:
    if (expandBufferLoadIntrinsic(Orig, /*IsRaw*/ false))
      return true;
    break;
  case Intrinsic::dx_resource_store_typedbuffer:
    if (expandBufferStoreIntrinsic(Orig, /*IsRaw*/ false))
      return true;
    break;
  case Intrinsic::usub_sat:
    Result = expandUsubSat(Orig);
    break;
  case Intrinsic::umul_with_overflow:
  case Intrinsic::smul_with_overflow:
    Result = expandMulWithOverflow(Orig, /*Signed=*/IntrinsicId ==
                                             Intrinsic::smul_with_overflow);
    break;
  case Intrinsic::vector_reduce_add:
  case Intrinsic::vector_reduce_fadd:
    Result = expandVecReduceAdd(Orig, IntrinsicId);
    break;
  case Intrinsic::matrix_multiply:
    Result = expandMatrixMultiply(Orig);
    break;
  case Intrinsic::matrix_transpose:
    Result = expandMatrixTranspose(Orig);
    break;
  }
  if (Result) {
    Orig->replaceAllUsesWith(Result);
    Orig->eraseFromParent();
    return true;
  }
  return false;
}
 
static bool expansionIntrinsics(Module &M) {
  for (auto &F : make_early_inc_range(M.functions())) {
    if (!isIntrinsicExpansion(F))
      continue;
    bool IntrinsicExpanded = false;
    for (User *U : make_early_inc_range(F.users())) {
      auto *IntrinsicCall = dyn_cast<CallInst>(U);
      if (!IntrinsicCall)
        continue;
      IntrinsicExpanded = expandIntrinsic(F, IntrinsicCall);
    }
    if (F.user_empty() && IntrinsicExpanded)
      F.eraseFromParent();
  }
  return true;
}
 
PreservedAnalyses DXILIntrinsicExpansion::run(Module &M,
                                              ModuleAnalysisManager &) {
  if (expansionIntrinsics(M))
    return PreservedAnalyses::none();
  return PreservedAnalyses::all();
}
 
bool DXILIntrinsicExpansionLegacy::runOnModule(Module &M) {
  return expansionIntrinsics(M);
}
 
char DXILIntrinsicExpansionLegacy::ID = 0;
 
INITIALIZE_PASS_BEGIN(DXILIntrinsicExpansionLegacy, DEBUG_TYPE,
                      "DXIL Intrinsic Expansion", false, false)
INITIALIZE_PASS_END(DXILIntrinsicExpansionLegacy, DEBUG_TYPE,
                    "DXIL Intrinsic Expansion", false, false)
 
ModulePass *llvm::createDXILIntrinsicExpansionLegacyPass() {
  return new DXILIntrinsicExpansionLegacy();
}