[NFC][DXIL] add ability to alias to an intrinsic

   - allows an intrinsic to be aliased for another one and provide
     additional arguments when being replaced

Author: inbelic

llvm/lib/Target/DirectX/DXIL.td

@@ -321,6 +321,9 @@ class DXILOp<int opcode, DXILOpClass opclass> {
 
   // Versioned attributes of operation
   list<Attributes> attributes = [];
+
+  // List of valid aliases
+  list<Intrinsic> aliases = [];
 }
 
 // Concrete definitions of DXIL Operations

llvm/lib/Target/DirectX/DXILOpLowering.cpp

@@ -40,6 +40,19 @@ static bool isVectorArgExpansion(Function &F) {
   return false;
 }
 
+static SmallVector<Value *> getAliasArgs(Function &F, IRBuilder<> &Builder) {
+  switch (F.getIntrinsicID()) {
+  case Intrinsic::dx_wave_active_sum:
+    return getWaveActiveOpArgs<0, true>(F, Builder);
+  case Intrinsic::dx_wave_active_usum:
+    return getWaveActiveOpArgs<0, false>(F, Builder);
+  default:
+    break;
+  }
+  SmallVector<Value *, 0> Args;
+  return Args;
+}
+
 static SmallVector<Value *> populateOperands(Value *Arg, IRBuilder<> &Builder) {
   SmallVector<Value *> ExtractedElements;
   auto *VecArg = dyn_cast<FixedVectorType>(Arg->getType());
@@ -106,7 +119,9 @@ class OpLowerer {
   }
 
   [[nodiscard]]
-  bool replaceFunctionWithOp(Function &F, dxil::OpCode DXILOp) {
+  bool replaceFunctionWithOp(
+      Function &F, dxil::OpCode DXILOp,
+      std::optional<SmallVector<Value *>> AliasArgs = std::nullopt) {
     bool IsVectorArgExpansion = isVectorArgExpansion(F);
     return replaceFunction(F, [&](CallInst *CI) -> Error {
       SmallVector<Value *> Args;
@@ -117,6 +132,11 @@ class OpLowerer {
       } else
         Args.append(CI->arg_begin(), CI->arg_end());
 
+      // Append any given alias arguments
+      if (AliasArgs) {
+        Args.append(AliasArgs->begin(), AliasArgs->end());
+      }
+
       Expected<CallInst *> OpCall =
           OpBuilder.tryCreateOp(DXILOp, Args, CI->getName(), F.getReturnType());
       if (Error E = OpCall.takeError())
@@ -476,6 +496,13 @@ class OpLowerer {
   case Intrin:                                                                 \
     HasErrors |= replaceFunctionWithOp(F, OpCode);                             \
     break;
+#include "DXILOperation.inc"
+#define DXIL_OP_INTRINSIC_ALIAS(OpCode, Intrin)                                \
+  case Intrin: {                                                               \
+    SmallVector<Value *> AliasArgs = getAliasArgs(F, OpBuilder.getIRB());      \
+    HasErrors |= replaceFunctionWithOp(F, OpCode, AliasArgs);                  \
+    break;                                                                     \
+  }
 #include "DXILOperation.inc"
       case Intrinsic::dx_handle_fromBinding:
         HasErrors |= lowerHandleFromBinding(F);

llvm/utils/TableGen/DXILEmitter.cpp

@@ -45,6 +45,7 @@ struct DXILOperationDesc {
   SmallVector<const Record *> AttrRecs;
   StringRef Intrinsic; // The llvm intrinsic map to OpName. Default is "" which
                        // means no map exists
+  SmallVector<StringRef> IntrinsicAliases;
   SmallVector<StringRef, 4>
       ShaderStages; // shader stages to which this applies, empty for all.
   int OverloadParamIndex;             // Index of parameter with overload type.
@@ -168,6 +169,15 @@ DXILOperationDesc::DXILOperationDesc(const Record *R) {
       Intrinsic = DefName.substr(4);
     }
   }
+
+  Recs = R->getValueAsListOfDefs("aliases");
+
+  for (const Record *CR : Recs) {
+    auto DefName = CR->getName();
+    assert(DefName.starts_with("int_") && "invalid intrinsic name");
+    // Remove the int_ from intrinsic name.
+    IntrinsicAliases.push_back(DefName.substr(4));
+  }
 }
 
 /// Return a string representation of OverloadKind enum that maps to
@@ -388,6 +398,26 @@ static void emitDXILIntrinsicMap(ArrayRef<DXILOperationDesc> Ops,
   OS << "#endif\n\n";
 }
 
+/// Emit map of DXIL operation to LLVM or DirectX intrinsic alias
+/// \param A vector of DXIL Ops
+/// \param Output stream
+static void emitDXILIntrinsicAliasMap(ArrayRef<DXILOperationDesc> Ops,
+                                      raw_ostream &OS) {
+  OS << "#ifdef DXIL_OP_INTRINSIC_ALIAS\n";
+  OS << "\n";
+  for (const auto &Op : Ops) {
+    for (const auto &Alias : Op.IntrinsicAliases) {
+      if (Alias.empty())
+        continue;
+      OS << "DXIL_OP_INTRINSIC_ALIAS(dxil::OpCode::" << Op.OpName
+         << ", Intrinsic::" << Alias << ")\n";
+    }
+  }
+  OS << "\n";
+  OS << "#undef DXIL_OP_INTRINSIC_ALIAS\n";
+  OS << "#endif\n\n";
+}
+
 /// Emit DXIL operation table
 /// \param A vector of DXIL Ops
 /// \param Output stream
@@ -529,6 +559,7 @@ static void EmitDXILOperation(const RecordKeeper &Records, raw_ostream &OS) {
   emitDXILOpParamTypes(Records, OS);
   emitDXILOpFunctionTypes(DXILOps, OS);
   emitDXILIntrinsicMap(DXILOps, OS);
+  emitDXILIntrinsicAliasMap(DXILOps, OS);
   OS << "#ifdef DXIL_OP_OPERATION_TABLE\n\n";
   emitDXILOperationTableDataStructs(Records, OS);
   emitDXILOperationTable(DXILOps, OS);