clang-p2996/llvm/lib/Target/DirectX/DXILOpBuilder.cpp

//===- DXILOpBuilder.cpp - Helper class for build DIXLOp functions --------===//
//
// 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 class to help build DXIL op functions.
//===----------------------------------------------------------------------===//

#include "DXILOpBuilder.h"
#include "DXILConstants.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Module.h"
#include "llvm/Support/DXILABI.h"
#include "llvm/Support/ErrorHandling.h"
#include <optional>

using namespace llvm;
using namespace llvm::dxil;

constexpr StringLiteral DXILOpNamePrefix = "dx.op.";

namespace {
enum OverloadKind : uint16_t {
  UNDEFINED = 0,
  VOID = 1,
  HALF = 1 << 1,
  FLOAT = 1 << 2,
  DOUBLE = 1 << 3,
  I1 = 1 << 4,
  I8 = 1 << 5,
  I16 = 1 << 6,
  I32 = 1 << 7,
  I64 = 1 << 8,
  UserDefineType = 1 << 9,
  ObjectType = 1 << 10,
};
struct Version {
  unsigned Major = 0;
  unsigned Minor = 0;
};

struct OpOverload {
  Version DXILVersion;
  uint16_t ValidTys;
};
} // namespace

struct OpStage {
  Version DXILVersion;
  uint32_t ValidStages;
};

struct OpAttribute {
  Version DXILVersion;
  uint32_t ValidAttrs;
};

static const char *getOverloadTypeName(OverloadKind Kind) {
  switch (Kind) {
  case OverloadKind::HALF:
    return "f16";
  case OverloadKind::FLOAT:
    return "f32";
  case OverloadKind::DOUBLE:
    return "f64";
  case OverloadKind::I1:
    return "i1";
  case OverloadKind::I8:
    return "i8";
  case OverloadKind::I16:
    return "i16";
  case OverloadKind::I32:
    return "i32";
  case OverloadKind::I64:
    return "i64";
  case OverloadKind::VOID:
  case OverloadKind::UNDEFINED:
    return "void";
  case OverloadKind::ObjectType:
  case OverloadKind::UserDefineType:
    break;
  }
  llvm_unreachable("invalid overload type for name");
}

static OverloadKind getOverloadKind(Type *Ty) {
  Type::TypeID T = Ty->getTypeID();
  switch (T) {
  case Type::VoidTyID:
    return OverloadKind::VOID;
  case Type::HalfTyID:
    return OverloadKind::HALF;
  case Type::FloatTyID:
    return OverloadKind::FLOAT;
  case Type::DoubleTyID:
    return OverloadKind::DOUBLE;
  case Type::IntegerTyID: {
    IntegerType *ITy = cast<IntegerType>(Ty);
    unsigned Bits = ITy->getBitWidth();
    switch (Bits) {
    case 1:
      return OverloadKind::I1;
    case 8:
      return OverloadKind::I8;
    case 16:
      return OverloadKind::I16;
    case 32:
      return OverloadKind::I32;
    case 64:
      return OverloadKind::I64;
    default:
      llvm_unreachable("invalid overload type");
      return OverloadKind::VOID;
    }
  }
  case Type::PointerTyID:
    return OverloadKind::UserDefineType;
  case Type::StructTyID:
    return OverloadKind::ObjectType;
  default:
    llvm_unreachable("invalid overload type");
    return OverloadKind::VOID;
  }
}

static std::string getTypeName(OverloadKind Kind, Type *Ty) {
  if (Kind < OverloadKind::UserDefineType) {
    return getOverloadTypeName(Kind);
  } else if (Kind == OverloadKind::UserDefineType) {
    StructType *ST = cast<StructType>(Ty);
    return ST->getStructName().str();
  } else if (Kind == OverloadKind::ObjectType) {
    StructType *ST = cast<StructType>(Ty);
    return ST->getStructName().str();
  } else {
    std::string Str;
    raw_string_ostream OS(Str);
    Ty->print(OS);
    return OS.str();
  }
}

// Static properties.
struct OpCodeProperty {
  dxil::OpCode OpCode;
  // Offset in DXILOpCodeNameTable.
  unsigned OpCodeNameOffset;
  dxil::OpCodeClass OpCodeClass;
  // Offset in DXILOpCodeClassNameTable.
  unsigned OpCodeClassNameOffset;
  llvm::SmallVector<OpOverload> Overloads;
  llvm::SmallVector<OpStage> Stages;
  llvm::SmallVector<OpAttribute> Attributes;
  int OverloadParamIndex;        // parameter index which control the overload.
                                 // When < 0, should be only 1 overload type.
  unsigned NumOfParameters;      // Number of parameters include return value.
  unsigned ParameterTableOffset; // Offset in ParameterTable.
};

// Include getOpCodeClassName getOpCodeProperty, getOpCodeName and
// getOpCodeParameterKind which generated by tableGen.
#define DXIL_OP_OPERATION_TABLE
#include "DXILOperation.inc"
#undef DXIL_OP_OPERATION_TABLE

static std::string constructOverloadName(OverloadKind Kind, Type *Ty,
                                         const OpCodeProperty &Prop) {
  if (Kind == OverloadKind::VOID) {
    return (Twine(DXILOpNamePrefix) + getOpCodeClassName(Prop)).str();
  }
  return (Twine(DXILOpNamePrefix) + getOpCodeClassName(Prop) + "." +
          getTypeName(Kind, Ty))
      .str();
}

static std::string constructOverloadTypeName(OverloadKind Kind,
                                             StringRef TypeName) {
  if (Kind == OverloadKind::VOID)
    return TypeName.str();

  assert(Kind < OverloadKind::UserDefineType && "invalid overload kind");
  return (Twine(TypeName) + getOverloadTypeName(Kind)).str();
}

static StructType *getOrCreateStructType(StringRef Name,
                                         ArrayRef<Type *> EltTys,
                                         LLVMContext &Ctx) {
  StructType *ST = StructType::getTypeByName(Ctx, Name);
  if (ST)
    return ST;

  return StructType::create(Ctx, EltTys, Name);
}

static StructType *getResRetType(Type *OverloadTy, LLVMContext &Ctx) {
  OverloadKind Kind = getOverloadKind(OverloadTy);
  std::string TypeName = constructOverloadTypeName(Kind, "dx.types.ResRet.");
  Type *FieldTypes[5] = {OverloadTy, OverloadTy, OverloadTy, OverloadTy,
                         Type::getInt32Ty(Ctx)};
  return getOrCreateStructType(TypeName, FieldTypes, Ctx);
}

static StructType *getHandleType(LLVMContext &Ctx) {
  return getOrCreateStructType("dx.types.Handle", PointerType::getUnqual(Ctx),
                               Ctx);
}

static Type *getTypeFromParameterKind(ParameterKind Kind, Type *OverloadTy) {
  auto &Ctx = OverloadTy->getContext();
  switch (Kind) {
  case ParameterKind::Void:
    return Type::getVoidTy(Ctx);
  case ParameterKind::Half:
    return Type::getHalfTy(Ctx);
  case ParameterKind::Float:
    return Type::getFloatTy(Ctx);
  case ParameterKind::Double:
    return Type::getDoubleTy(Ctx);
  case ParameterKind::I1:
    return Type::getInt1Ty(Ctx);
  case ParameterKind::I8:
    return Type::getInt8Ty(Ctx);
  case ParameterKind::I16:
    return Type::getInt16Ty(Ctx);
  case ParameterKind::I32:
    return Type::getInt32Ty(Ctx);
  case ParameterKind::I64:
    return Type::getInt64Ty(Ctx);
  case ParameterKind::Overload:
    return OverloadTy;
  case ParameterKind::ResourceRet:
    return getResRetType(OverloadTy, Ctx);
  case ParameterKind::DXILHandle:
    return getHandleType(Ctx);
  default:
    break;
  }
  llvm_unreachable("Invalid parameter kind");
  return nullptr;
}

static ShaderKind getShaderKindEnum(Triple::EnvironmentType EnvType) {
  switch (EnvType) {
  case Triple::Pixel:
    return ShaderKind::pixel;
  case Triple::Vertex:
    return ShaderKind::vertex;
  case Triple::Geometry:
    return ShaderKind::geometry;
  case Triple::Hull:
    return ShaderKind::hull;
  case Triple::Domain:
    return ShaderKind::domain;
  case Triple::Compute:
    return ShaderKind::compute;
  case Triple::Library:
    return ShaderKind::library;
  case Triple::RayGeneration:
    return ShaderKind::raygeneration;
  case Triple::Intersection:
    return ShaderKind::intersection;
  case Triple::AnyHit:
    return ShaderKind::anyhit;
  case Triple::ClosestHit:
    return ShaderKind::closesthit;
  case Triple::Miss:
    return ShaderKind::miss;
  case Triple::Callable:
    return ShaderKind::callable;
  case Triple::Mesh:
    return ShaderKind::mesh;
  case Triple::Amplification:
    return ShaderKind::amplification;
  default:
    break;
  }
  llvm_unreachable(
      "Shader Kind Not Found - Invalid DXIL Environment Specified");
}

/// Construct DXIL function type. This is the type of a function with
/// the following prototype
///     OverloadType dx.op.<opclass>.<return-type>(int opcode, <param types>)
/// <param-types> are constructed from types in Prop.
/// \param Prop  Structure containing DXIL Operation properties based on
///               its specification in DXIL.td.
/// \param OverloadTy Return type to be used to construct DXIL function type.
static FunctionType *getDXILOpFunctionType(const OpCodeProperty *Prop,
                                           Type *ReturnTy, Type *OverloadTy) {
  SmallVector<Type *> ArgTys;

  const ParameterKind *ParamKinds = getOpCodeParameterKind(*Prop);

  // Add ReturnTy as return type of the function
  ArgTys.emplace_back(ReturnTy);

  // Add DXIL Opcode value type viz., Int32 as first argument
  ArgTys.emplace_back(Type::getInt32Ty(OverloadTy->getContext()));

  // Add DXIL Operation parameter types as specified in DXIL properties
  for (unsigned I = 0; I < Prop->NumOfParameters; ++I) {
    ParameterKind Kind = ParamKinds[I];
    ArgTys.emplace_back(getTypeFromParameterKind(Kind, OverloadTy));
  }
  return FunctionType::get(
      ArgTys[0], ArrayRef<Type *>(&ArgTys[1], ArgTys.size() - 1), false);
}

/// Get index of the property from PropList valid for the most recent
/// DXIL version not greater than DXILVer.
/// PropList is expected to be sorted in ascending order of DXIL version.
template <typename T>
static std::optional<size_t> getPropIndex(ArrayRef<T> PropList,
                                          const VersionTuple DXILVer) {
  size_t Index = PropList.size() - 1;
  for (auto Iter = PropList.rbegin(); Iter != PropList.rend();
       Iter++, Index--) {
    const T &Prop = *Iter;
    if (VersionTuple(Prop.DXILVersion.Major, Prop.DXILVersion.Minor) <=
        DXILVer) {
      return Index;
    }
  }
  return std::nullopt;
}

namespace llvm {
namespace dxil {

// No extra checks on TargetTriple need be performed to verify that the
// Triple is well-formed or that the target is supported since these checks
// would have been done at the time the module M is constructed in the earlier
// stages of compilation.
DXILOpBuilder::DXILOpBuilder(Module &M, IRBuilderBase &B) : M(M), B(B) {
  Triple TT(Triple(M.getTargetTriple()));
  DXILVersion = TT.getDXILVersion();
  ShaderStage = TT.getEnvironment();
  // Ensure Environment type is known
  if (ShaderStage == Triple::UnknownEnvironment) {
    report_fatal_error(
        Twine(DXILVersion.getAsString()) +
            ": Unknown Compilation Target Shader Stage specified ",
        /*gen_crash_diag*/ false);
  }
}

CallInst *DXILOpBuilder::createDXILOpCall(dxil::OpCode OpCode, Type *ReturnTy,
                                          Type *OverloadTy,
                                          SmallVector<Value *> Args) {

  const OpCodeProperty *Prop = getOpCodeProperty(OpCode);
  std::optional<size_t> OlIndexOrErr =
      getPropIndex(ArrayRef(Prop->Overloads), DXILVersion);
  if (!OlIndexOrErr.has_value()) {
    report_fatal_error(Twine(getOpCodeName(OpCode)) +
                           ": No valid overloads found for DXIL Version - " +
                           DXILVersion.getAsString(),
                       /*gen_crash_diag*/ false);
  }
  uint16_t ValidTyMask = Prop->Overloads[*OlIndexOrErr].ValidTys;

  OverloadKind Kind = getOverloadKind(OverloadTy);

  // Check if the operation supports overload types and OverloadTy is valid
  // per the specified types for the operation
  if ((ValidTyMask != OverloadKind::UNDEFINED) &&
      (ValidTyMask & (uint16_t)Kind) == 0) {
    report_fatal_error(Twine("Invalid Overload Type for DXIL operation - ") +
                           getOpCodeName(OpCode),
                       /* gen_crash_diag=*/false);
  }

  // Perform necessary checks to ensure Opcode is valid in the targeted shader
  // kind
  std::optional<size_t> StIndexOrErr =
      getPropIndex(ArrayRef(Prop->Stages), DXILVersion);
  if (!StIndexOrErr.has_value()) {
    report_fatal_error(Twine(getOpCodeName(OpCode)) +
                           ": No valid stages found for DXIL Version - " +
                           DXILVersion.getAsString(),
                       /*gen_crash_diag*/ false);
  }
  uint16_t ValidShaderKindMask = Prop->Stages[*StIndexOrErr].ValidStages;

  // Ensure valid shader stage properties are specified
  if (ValidShaderKindMask == ShaderKind::removed) {
    report_fatal_error(
        Twine(DXILVersion.getAsString()) +
            ": Unsupported Target Shader Stage for DXIL operation - " +
            getOpCodeName(OpCode),
        /*gen_crash_diag*/ false);
  }

  // Shader stage need not be validated since getShaderKindEnum() fails
  // for unknown shader stage.

  // Verify the target shader stage is valid for the DXIL operation
  ShaderKind ModuleStagekind = getShaderKindEnum(ShaderStage);
  if (!(ValidShaderKindMask & ModuleStagekind)) {
    auto ShaderEnvStr = Triple::getEnvironmentTypeName(ShaderStage);
    report_fatal_error(Twine(ShaderEnvStr) +
                           " : Invalid Shader Stage for DXIL operation - " +
                           getOpCodeName(OpCode) + " for DXIL Version " +
                           DXILVersion.getAsString(),
                       /*gen_crash_diag*/ false);
  }

  std::string DXILFnName = constructOverloadName(Kind, OverloadTy, *Prop);
  FunctionCallee DXILFn;
  // Get the function with name DXILFnName, if one exists
  if (auto *Func = M.getFunction(DXILFnName)) {
    DXILFn = FunctionCallee(Func);
  } else {
    // Construct and add a function with name DXILFnName
    FunctionType *DXILOpFT = getDXILOpFunctionType(Prop, ReturnTy, OverloadTy);
    DXILFn = M.getOrInsertFunction(DXILFnName, DXILOpFT);
  }

  return B.CreateCall(DXILFn, Args);
}

Type *DXILOpBuilder::getOverloadTy(dxil::OpCode OpCode, FunctionType *FT) {

  const OpCodeProperty *Prop = getOpCodeProperty(OpCode);
  // If DXIL Op has no overload parameter, just return the
  // precise return type specified.
  if (Prop->OverloadParamIndex < 0) {
    return FT->getReturnType();
  }

  // Consider FT->getReturnType() as default overload type, unless
  // Prop->OverloadParamIndex != 0.
  Type *OverloadType = FT->getReturnType();
  if (Prop->OverloadParamIndex != 0) {
    // Skip Return Type.
    OverloadType = FT->getParamType(Prop->OverloadParamIndex - 1);
  }

  const ParameterKind *ParamKinds = getOpCodeParameterKind(*Prop);
  auto Kind = ParamKinds[Prop->OverloadParamIndex];
  // For ResRet and CBufferRet, OverloadTy is in field of StructType.
  if (Kind == ParameterKind::CBufferRet ||
      Kind == ParameterKind::ResourceRet) {
    auto *ST = cast<StructType>(OverloadType);
    OverloadType = ST->getElementType(0);
  }
  return OverloadType;
}

const char *DXILOpBuilder::getOpCodeName(dxil::OpCode DXILOp) {
  return ::getOpCodeName(DXILOp);
}
} // namespace dxil
} // namespace llvm