[DirectX][OpLowering] Simplify named struct handling

[bogner]

This removes "replaceFunctionWithNamedStructOp" and folds its functionality into "replaceFunctionWithOp". It turns out we were overcomplicating things and this is trivial to handle generically.

Fixes #113192

[llvmbot]

@llvm/pr-subscribers-backend-directx

Author: Justin Bogner (bogner)

Changes

This removes "replaceFunctionWithNamedStructOp" and folds its functionality into "replaceFunctionWithOp". It turns out we were overcomplicating things and this is trivial to handle generically.

Fixes #113192

---

Full diff: https://github.com/llvm/llvm-project/pull/128247.diff

2 Files Affected:

• (modified) llvm/lib/Target/DirectX/DXIL.td (+1)
• (modified) llvm/lib/Target/DirectX/DXILOpLowering.cpp (+27-34)

diff --git a/llvm/lib/Target/DirectX/DXIL.td b/llvm/lib/Target/DirectX/DXIL.td
index 7cb841d9bd5b5..d59e28c37b91d 100644
--- a/llvm/lib/Target/DirectX/DXIL.td
+++ b/llvm/lib/Target/DirectX/DXIL.td
@@ -930,6 +930,7 @@ def MakeDouble : DXILOp<101, makeDouble> {
 
 def SplitDouble : DXILOp<102, splitDouble> {
   let Doc = "Splits a double into 2 uints";
+  let intrinsics = [IntrinSelect<int_dx_splitdouble>];
   let arguments = [OverloadTy];
   let result = SplitDoubleTy;
   let overloads = [Overloads<DXIL1_0, [DoubleTy]>];
diff --git a/llvm/lib/Target/DirectX/DXILOpLowering.cpp b/llvm/lib/Target/DirectX/DXILOpLowering.cpp
index 0c245c1a43d31..5459695f96863 100644
--- a/llvm/lib/Target/DirectX/DXILOpLowering.cpp
+++ b/llvm/lib/Target/DirectX/DXILOpLowering.cpp
@@ -120,6 +120,28 @@ class OpLowerer {
     int Value;
   };
 
+  /// Replaces uses of a struct with uses of an equivalent named struct.
+  ///
+  /// DXIL operations that return structs give them well known names, so we need
+  /// to update uses when we switch from an LLVM intrinsic to an op.
+  Error replaceNamedStructUses(CallInst *Intrin, CallInst *DXILOp) {
+    auto *IntrinTy = cast<StructType>(Intrin->getType());
+    auto *DXILOpTy = cast<StructType>(DXILOp->getType());
+    if (!IntrinTy->isLayoutIdentical(DXILOpTy))
+      return make_error<StringError>(
+          "Type mismatch between intrinsic and DXIL op",
+          inconvertibleErrorCode());
+
+    for (Use &U : make_early_inc_range(Intrin->uses()))
+      if (auto *EVI = dyn_cast<ExtractValueInst>(U.getUser()))
+        EVI->setOperand(0, DXILOp);
+      else
+        return make_error<StringError>(
+            "DXIL ops that return structs may only be used by extractvalue",
+            inconvertibleErrorCode());
+    return Error::success();
+  }
+
   [[nodiscard]] bool
   replaceFunctionWithOp(Function &F, dxil::OpCode DXILOp,
                         ArrayRef<IntrinArgSelect> ArgSelects) {
@@ -154,32 +176,13 @@ class OpLowerer {
       if (Error E = OpCall.takeError())
         return E;
 
-      CI->replaceAllUsesWith(*OpCall);
-      CI->eraseFromParent();
-      return Error::success();
-    });
-  }
-
-  [[nodiscard]] bool replaceFunctionWithNamedStructOp(
-      Function &F, dxil::OpCode DXILOp, Type *NewRetTy,
-      llvm::function_ref<Error(CallInst *CI, CallInst *Op)> ReplaceUses) {
-    bool IsVectorArgExpansion = isVectorArgExpansion(F);
-    return replaceFunction(F, [&](CallInst *CI) -> Error {
-      SmallVector<Value *> Args;
-      OpBuilder.getIRB().SetInsertPoint(CI);
-      if (IsVectorArgExpansion) {
-        SmallVector<Value *> NewArgs = argVectorFlatten(CI, OpBuilder.getIRB());
-        Args.append(NewArgs.begin(), NewArgs.end());
+      if (isa<StructType>(CI->getType())) {
+        if (Error E = replaceNamedStructUses(CI, *OpCall))
+          return E;
       } else
-        Args.append(CI->arg_begin(), CI->arg_end());
-
-      Expected<CallInst *> OpCall =
-          OpBuilder.tryCreateOp(DXILOp, Args, CI->getName(), NewRetTy);
-      if (Error E = OpCall.takeError())
-        return E;
-      if (Error E = ReplaceUses(CI, *OpCall))
-        return E;
+        CI->replaceAllUsesWith(*OpCall);
 
+      CI->eraseFromParent();
       return Error::success();
     });
   }
@@ -814,16 +817,6 @@ class OpLowerer {
       case Intrinsic::dx_resource_updatecounter:
         HasErrors |= lowerUpdateCounter(F);
         break;
-      // TODO: this can be removed when
-      // https://github.com/llvm/llvm-project/issues/113192 is fixed
-      case Intrinsic::dx_splitdouble:
-        HasErrors |= replaceFunctionWithNamedStructOp(
-            F, OpCode::SplitDouble,
-            OpBuilder.getSplitDoubleType(M.getContext()),
-            [&](CallInst *CI, CallInst *Op) {
-              return replaceSplitDoubleCallUsages(CI, Op);
-            });
-        break;
       case Intrinsic::ctpop:
         HasErrors |= lowerCtpopToCountBits(F);
         break;

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[hekota]

LGTM! There is possibly one dead function to remove.

[hekota]

It looks like this is the only place replaceSplitDoubleCallUsages is called, so you might want to remove its definition as well.