Refactor udt intrinsic arg copy to before SROA

tex3d · tex3d · commit 39f80672e71d · 2025-05-07T23:40:13.000-07:00
Intrinsics that take UDT arguments need copy-in/copy-out.
Other aggregate args are flattened for intrinsic calls.
Previously, these operations were intermingled, driven by SROA on alloca/GV values.

There were RayDesc arguments that weren't treated consistently, leading to problems.
They should be flattened into the intrinsic arguments, but TraceRay calls didn't do this.
Doing this for TraceRay would cause an issue where the subsequent Payload UDT arg could be in a different position depending on whether it was processed before or after flattening the RayDesc.

This change is in preparation for flattening RayDesc args.
It separates the UDT copy-in/copy-out into a separate operation before SROA.

When doing the copy-in/copy-out separately, we don't expand the memcpy operations until they are handled separately during SROA.
This causes a change in some IR shape, which had to be adjusted for in some tests.
There are a couple remaining PIX pass tests that still need adjustment.
These will likely need further changes when flattening RayDesc, so maybe they should be revisited after that.

TODO:
- Fix the remaining tests
- Flatten the RayDesc parameters
- Adjust tests for accordingly
diff --git a/include/dxc/HLSL/HLOperations.h b/include/dxc/HLSL/HLOperations.h
@@ -396,7 +396,10 @@ const unsigned kAnnotateHandleResourceTypeOpIdx = 3;
 
 // TraceRay.
 const unsigned kTraceRayRayDescOpIdx = 7;
-const unsigned kTraceRayPayLoadOpIdx = 8;
+// kTraceRayPayloadPreOpIdx is before flattening the RayDesc
+const unsigned kTraceRayPayloadPreOpIdx = 8;
+// kTraceRayPayloadOpIdx is after flattening the RayDesc
+const unsigned kTraceRayPayloadOpIdx = 8; // TODO: Update after flattening
 
 // AllocateRayQuery
 const unsigned kAllocateRayQueryRayFlagsIdx = 1;
@@ -434,12 +437,22 @@ const unsigned kAnnotateNodeHandleNodePropIdx = 2;
 const unsigned kAnnotateNodeRecordHandleNodeRecordPropIdx = 2;
 
 // HitObject::MakeMiss
-const unsigned kHitObjectMakeMiss_NumOp = 8;
+const unsigned kHitObjectMakeMiss_PreNumOp = 8;
+const unsigned kHitObjectMakeMiss_NumOp = 8; // TODO: Update after flattening
 const unsigned kHitObjectMakeMissRayDescOpIdx = 4;
 
 // HitObject::TraceRay
 const unsigned kHitObjectTraceRay_RayDescOpIdx = 8;
-const unsigned kHitObjectTraceRay_NumOp = 10;
+// kHitObjectTraceRay_PayloadPreOpIdx is before flattening the RayDesc
+const unsigned kHitObjectTraceRay_PayloadPreOpIdx = 9;
+// kHitObjectTraceRay_PayloadOpIdx is after flattening the RayDesc
+const unsigned kHitObjectTraceRay_PayloadOpIdx = 9; // TODO: Update after
+                                                    // flattening
+const unsigned kHitObjectTraceRay_PreNumOp = 10;
+const unsigned kHitObjectTraceRay_NumOp = 10; // TODO: Update after flattening
+
+// HitObject::Invoke
+const unsigned kHitObjectInvoke_PayloadOpIdx = 2;
 
 // HitObject::FromRayQuery
 const unsigned kHitObjectFromRayQuery_WithAttrs_AttributeOpIdx = 4;
diff --git a/lib/HLSL/HLOperationLower.cpp b/lib/HLSL/HLOperationLower.cpp
@@ -5760,7 +5760,7 @@ Value *TranslateTraceRay(CallInst *CI, IntrinsicOp IOP, OP::OpCode OpCode,
   hlsl::OP *OP = &Helper.hlslOP;
 
   Value *RayDesc = CI->getArgOperand(HLOperandIndex::kTraceRayRayDescOpIdx);
-  Value *PayLoad = CI->getArgOperand(HLOperandIndex::kTraceRayPayLoadOpIdx);
+  Value *PayLoad = CI->getArgOperand(HLOperandIndex::kTraceRayPayloadOpIdx);
 
   Value *Args[DXIL::OperandIndex::kTraceRayNumOp];
   Args[0] = OP->GetU32Const(static_cast<unsigned>(OpCode));
diff --git a/lib/Transforms/Scalar/ScalarReplAggregatesHLSL.cpp b/lib/Transforms/Scalar/ScalarReplAggregatesHLSL.cpp
@@ -129,7 +129,6 @@ class SROA_Helper {
   void RewriteMemIntrin(MemIntrinsic *MI, Value *OldV);
   void RewriteCall(CallInst *CI);
   void RewriteBitCast(BitCastInst *BCI);
-  void RewriteCallArg(CallInst *CI, unsigned ArgIdx, bool bIn, bool bOut);
 };
 
 } // namespace
@@ -1478,6 +1477,46 @@ void isSafePHISelectUseForScalarRepl(Instruction *I, uint64_t Offset,
   }
 }
 
+static bool isUDTIntrinsicArg(CallInst *CI, unsigned OpIdx) {
+  if (HLOpcodeGroup::HLIntrinsic != GetHLOpcodeGroup(CI->getCalledFunction()))
+    return false;
+  switch (static_cast<IntrinsicOp>(GetHLOpcode(CI))) {
+  case IntrinsicOp::IOP_TraceRay:
+    if (OpIdx == HLOperandIndex::kTraceRayRayDescOpIdx)
+      return true;
+    if (OpIdx == HLOperandIndex::kTraceRayPayloadPreOpIdx)
+      return true;
+    break;
+  case IntrinsicOp::IOP_ReportHit:
+    if (OpIdx == HLOperandIndex::kReportIntersectionAttributeOpIdx)
+      return true;
+    break;
+  case IntrinsicOp::IOP_CallShader:
+    if (OpIdx == HLOperandIndex::kCallShaderPayloadOpIdx)
+      return true;
+    break;
+  case IntrinsicOp::MOP_DxHitObject_FromRayQuery:
+    if (OpIdx ==
+        HLOperandIndex::kHitObjectFromRayQuery_WithAttrs_AttributeOpIdx)
+      return true;
+    break;
+  case IntrinsicOp::MOP_DxHitObject_TraceRay:
+    // TODO: Remove RayDesc for flattening
+    if (OpIdx == HLOperandIndex::kHitObjectTraceRay_RayDescOpIdx)
+      return true;
+    if (OpIdx == HLOperandIndex::kHitObjectTraceRay_PayloadPreOpIdx)
+      return true;
+    break;
+  case IntrinsicOp::MOP_DxHitObject_Invoke:
+    if (OpIdx == HLOperandIndex::kHitObjectInvoke_PayloadOpIdx)
+      return true;
+    break;
+  default:
+    break;
+  }
+  return false;
+}
+
 /// isSafeForScalarRepl - Check if instruction I is a safe use with regard to
 /// performing scalar replacement of alloca AI.  The results are flagged in
 /// the Info parameter.  Offset indicates the position within AI that is
@@ -1536,15 +1575,9 @@ void isSafeForScalarRepl(Instruction *I, uint64_t Offset, AllocaInfo &Info) {
       if (HLOpcodeGroup::NotHL == group)
         return MarkUnsafe(Info, User);
       else if (HLOpcodeGroup::HLIntrinsic == group) {
-        // TODO: should we check HL parameter type for UDT overload instead of
-        // basing on IOP?
-        IntrinsicOp opcode = static_cast<IntrinsicOp>(GetHLOpcode(CI));
-        if (IntrinsicOp::IOP_TraceRay == opcode ||
-            IntrinsicOp::MOP_DxHitObject_TraceRay == opcode ||
-            IntrinsicOp::MOP_DxHitObject_Invoke == opcode ||
-            IntrinsicOp::IOP_ReportHit == opcode ||
-            IntrinsicOp::IOP_CallShader == opcode) {
-          return MarkUnsafe(Info, User);
+        for (unsigned OpIdx = 0; OpIdx < CI->getNumArgOperands(); OpIdx++) {
+          if (CI->getArgOperand(OpIdx) == I && isUDTIntrinsicArg(CI, OpIdx))
+            return MarkUnsafe(Info, User);
         }
       }
     } else {
@@ -2666,8 +2699,7 @@ void SROA_Helper::RewriteBitCast(BitCastInst *BCI) {
 ///                  replace OldVal with alloca and
 ///                  copy in copy out data between alloca and flattened NewElts
 ///                  in CallInst.
-void SROA_Helper::RewriteCallArg(CallInst *CI, unsigned ArgIdx, bool bIn,
-                                 bool bOut) {
+static void RewriteCallArg(CallInst *CI, unsigned ArgIdx, bool bIn, bool bOut) {
   Function *F = CI->getParent()->getParent();
   IRBuilder<> AllocaBuilder(dxilutil::FindAllocaInsertionPt(F));
   const DataLayout &DL = F->getParent()->getDataLayout();
@@ -2678,16 +2710,60 @@ void SROA_Helper::RewriteCallArg(CallInst *CI, unsigned ArgIdx, bool bIn,
   Value *Alloca = AllocaBuilder.CreateAlloca(userTyElt);
   IRBuilder<> Builder(CI);
   if (bIn) {
-    MemCpyInst *cpy = cast<MemCpyInst>(Builder.CreateMemCpy(
-        Alloca, userTyV, DL.getTypeAllocSize(userTyElt), false));
-    RewriteMemIntrin(cpy, cpy->getRawSource());
+    Builder.CreateMemCpy(Alloca, userTyV, DL.getTypeAllocSize(userTyElt),
+                         false);
   }
   CI->setArgOperand(ArgIdx, Alloca);
   if (bOut) {
     Builder.SetInsertPoint(CI->getNextNode());
-    MemCpyInst *cpy = cast<MemCpyInst>(Builder.CreateMemCpy(
-        userTyV, Alloca, DL.getTypeAllocSize(userTyElt), false));
-    RewriteMemIntrin(cpy, cpy->getRawSource());
+    Builder.CreateMemCpy(userTyV, Alloca, DL.getTypeAllocSize(userTyElt),
+                         false);
+  }
+}
+
+static void copyIntrinsicUDTArgs(HLModule &HLM) {
+  // Iterate HLIntrinsic function users
+  // For specific intrinsics, use RewriteCallArg on UDT args
+  for (Function &F : HLM.GetModule()->functions()) {
+    if (F.isIntrinsic() || !F.isDeclaration())
+      continue;
+    if (GetHLOpcodeGroup(&F) != HLOpcodeGroup::HLIntrinsic)
+      continue;
+    // Iterate users
+    for (User *U : F.users()) {
+      if (CallInst *CI = dyn_cast<CallInst>(U)) {
+        switch (static_cast<IntrinsicOp>(GetHLOpcode(CI))) {
+        case IntrinsicOp::IOP_TraceRay:
+          // TODO: Remove RayDesc for flattening
+          RewriteCallArg(CI, HLOperandIndex::kTraceRayRayDescOpIdx,
+                         /*bIn*/ true, /*bOut*/ false);
+          RewriteCallArg(CI, HLOperandIndex::kTraceRayPayloadPreOpIdx,
+                         /*bIn*/ true, /*bOut*/ true);
+          break;
+        case IntrinsicOp::IOP_ReportHit:
+          RewriteCallArg(CI, HLOperandIndex::kReportIntersectionAttributeOpIdx,
+                         /*bIn*/ true, /*bOut*/ false);
+          break;
+        case IntrinsicOp::IOP_CallShader:
+          RewriteCallArg(CI, HLOperandIndex::kCallShaderPayloadOpIdx,
+                         /*bIn*/ true, /*bOut*/ true);
+          break;
+        case IntrinsicOp::MOP_DxHitObject_FromRayQuery:
+          if (CI->getNumArgOperands() ==
+              HLOperandIndex::kHitObjectFromRayQuery_WithAttrs_NumOp) {
+            RewriteCallArg(
+                CI,
+                HLOperandIndex::kHitObjectFromRayQuery_WithAttrs_AttributeOpIdx,
+                /*bIn*/ true, /*bOut*/ false);
+          }
+          break;
+        case IntrinsicOp::MOP_DxHitObject_TraceRay:
+          RewriteCallArg(CI, HLOperandIndex::kHitObjectTraceRay_PayloadPreOpIdx,
+                         /*bIn*/ true, /*bOut*/ true);
+          break;
+        }
+      }
+    }
   }
 }
 
@@ -2741,10 +2817,23 @@ static CallInst *RewriteWithFlattenedHLIntrinsicCall(CallInst *CI,
 
 /// RewriteCall - Replace OldVal with flattened NewElts in CallInst.
 void SROA_Helper::RewriteCall(CallInst *CI) {
-  HLOpcodeGroup group = GetHLOpcodeGroupByName(CI->getCalledFunction());
-  if (group != HLOpcodeGroup::NotHL) {
+  HLOpcodeGroup Group = GetHLOpcodeGroupByName(CI->getCalledFunction());
+  if (Group != HLOpcodeGroup::NotHL) {
     unsigned opcode = GetHLOpcode(CI);
-    if (group == HLOpcodeGroup::HLIntrinsic) {
+    if (Group == HLOpcodeGroup::HLIntrinsic) {
+      // RayQuery this pointer replacement.
+      if (OldVal->getType()->isPointerTy() &&
+          dxilutil::IsHLSLRayQueryType(
+              OldVal->getType()->getPointerElementType())) {
+        // For RayQuery methods, we want to replace the RayQuery this pointer
+        // with a load and use of the underlying handle value.
+        // This will allow elimination of RayQuery types earlier.
+        RewriteWithFlattenedHLIntrinsicCall(CI, OldVal, NewElts,
+                                            /*loadElts*/ true);
+        DeadInsts.push_back(CI);
+        return;
+      }
+
       IntrinsicOp IOP = static_cast<IntrinsicOp>(opcode);
       switch (IOP) {
       case IntrinsicOp::MOP_Append: {
@@ -2756,84 +2845,42 @@ void SROA_Helper::RewriteCall(CallInst *CI) {
                                             /*loadElts*/ false);
         DeadInsts.push_back(CI);
       } break;
-      case IntrinsicOp::IOP_TraceRay: {
-        if (OldVal ==
-            CI->getArgOperand(HLOperandIndex::kTraceRayRayDescOpIdx)) {
-          RewriteCallArg(CI, HLOperandIndex::kTraceRayRayDescOpIdx,
-                         /*bIn*/ true, /*bOut*/ false);
-        } else {
-          DXASSERT(OldVal ==
-                       CI->getArgOperand(HLOperandIndex::kTraceRayPayLoadOpIdx),
-                   "else invalid TraceRay");
-          RewriteCallArg(CI, HLOperandIndex::kTraceRayPayLoadOpIdx,
-                         /*bIn*/ true, /*bOut*/ true);
-        }
-      } break;
-      case IntrinsicOp::IOP_ReportHit: {
-        RewriteCallArg(CI, HLOperandIndex::kReportIntersectionAttributeOpIdx,
-                       /*bIn*/ true, /*bOut*/ false);
-      } break;
-      case IntrinsicOp::IOP_CallShader: {
-        RewriteCallArg(CI, HLOperandIndex::kCallShaderPayloadOpIdx,
-                       /*bIn*/ true, /*bOut*/ true);
-      } break;
-      case IntrinsicOp::MOP_DxHitObject_MakeMiss: {
+      //case IntrinsicOp::IOP_TraceRay:
+      //  if (OldVal ==
+      //      CI->getArgOperand(HLOperandIndex::kTraceRayRayDescOpIdx)) {
+      //    // TODO: flatten RayDesc
+      //     RewriteWithFlattenedHLIntrinsicCall(CI, OldVal, NewElts,
+      //                                        /*loadElts*/ true);
+      //     DeadInsts.push_back(CI);
+      //  }
+      //break;
+      //case IntrinsicOp::MOP_DxHitObject_TraceRay:
+      //  if (OldVal ==
+      //      CI->getArgOperand(HLOperandIndex::kHitObjectTraceRay_RayDescOpIdx)) {
+      //    // TODO: flatten RayDesc
+      //     RewriteWithFlattenedHLIntrinsicCall(CI, OldVal, NewElts,
+      //                                        /*loadElts*/ true);
+      //     DeadInsts.push_back(CI);
+      //  }
+      //  break;
+      case IntrinsicOp::MOP_DxHitObject_MakeMiss:
         if (OldVal ==
             CI->getArgOperand(HLOperandIndex::kHitObjectMakeMissRayDescOpIdx)) {
           RewriteWithFlattenedHLIntrinsicCall(CI, OldVal, NewElts,
                                               /*loadElts*/ true);
           DeadInsts.push_back(CI);
         }
-      } break;
-      case IntrinsicOp::MOP_TraceRayInline: {
+        break;
+      case IntrinsicOp::MOP_TraceRayInline:
         if (OldVal ==
             CI->getArgOperand(HLOperandIndex::kTraceRayInlineRayDescOpIdx)) {
           RewriteWithFlattenedHLIntrinsicCall(CI, OldVal, NewElts,
                                               /*loadElts*/ true);
           DeadInsts.push_back(CI);
           break;
         }
-      }
         LLVM_FALLTHROUGH;
-      case IntrinsicOp::MOP_DxHitObject_FromRayQuery: {
-        const bool IsWithAttrs =
-            CI->getNumArgOperands() ==
-            HLOperandIndex::kHitObjectFromRayQuery_WithAttrs_NumOp;
-        if (IsWithAttrs &&
-            (OldVal ==
-             CI->getArgOperand(
-                 HLOperandIndex::
-                     kHitObjectFromRayQuery_WithAttrs_AttributeOpIdx))) {
-          RewriteCallArg(
-              CI,
-              HLOperandIndex::kHitObjectFromRayQuery_WithAttrs_AttributeOpIdx,
-              /*bIn*/ true, /*bOut*/ false);
-          break;
-        }
-
-        // For RayQuery methods, we want to replace the RayQuery this pointer
-        // with a load and use of the underlying handle value.
-        // This will allow elimination of RayQuery types earlier.
-        RewriteWithFlattenedHLIntrinsicCall(CI, OldVal, NewElts,
-                                            /*loadElts*/ true);
-        DeadInsts.push_back(CI);
-        break;
-      }
       default:
-        // RayQuery this pointer replacement.
-        if (OldVal->getType()->isPointerTy() &&
-            CI->getNumArgOperands() >= HLOperandIndex::kHandleOpIdx &&
-            OldVal == CI->getArgOperand(HLOperandIndex::kHandleOpIdx) &&
-            dxilutil::IsHLSLRayQueryType(
-                OldVal->getType()->getPointerElementType())) {
-          // For RayQuery methods, we want to replace the RayQuery this pointer
-          // with a load and use of the underlying handle value.
-          // This will allow elimination of RayQuery types earlier.
-          RewriteWithFlattenedHLIntrinsicCall(CI, OldVal, NewElts,
-                                              /*loadElts*/ true);
-          DeadInsts.push_back(CI);
-          break;
-        }
         DXASSERT(0, "cannot flatten hlsl intrinsic.");
       }
     }
@@ -4416,6 +4463,9 @@ class SROA_Parameter_HLSL : public ModulePass {
       F->eraseFromParent();
     }
 
+    // Expand flattened copy-in/copy-out for intrinsic UDT args:
+    copyIntrinsicUDTArgs(*m_pHLModule);
+
     // SROA globals and allocas.
     SROAGlobalAndAllocas(*m_pHLModule, m_HasDbgInfo);
 
diff --git a/tools/clang/test/DXC/Passes/ScalarReplHLSL/hitobject_fromrayquery_scalarrepl.ll b/tools/clang/test/DXC/Passes/ScalarReplHLSL/hitobject_fromrayquery_scalarrepl.ll
@@ -95,10 +95,10 @@ target triple = "dxil-ms-dx"
 @"$Globals" = external constant %ConstantBuffer
 
 ; CHECK: %[[RQA:[^ ]+]] = alloca i32
-; CHECK: %[[ATTRA0:[^ ]+]] = alloca %struct.CustomAttrs
-; CHECK: %[[ATTRA1:[^ ]+]] = alloca %struct.CustomAttrs
 ; CHECK: %[[XATTRA:[^ ]+]] = alloca float
 ; CHECK: %[[YATTRA:[^ ]+]] = alloca float
+; CHECK: %[[ATTRA0:[^ ]+]] = alloca %struct.CustomAttrs
+; CHECK: %[[ATTRA1:[^ ]+]] = alloca %struct.CustomAttrs
 
 ; COM: Check same query handle used for TraceRayInline and the FromRayQuery calls
 ; CHECK: %[[RQH:[^ ]+]] = load i32, i32* %[[RQA]]
@@ -122,7 +122,7 @@ target triple = "dxil-ms-dx"
 ; CHECK: store float %[[XF1]], float* %[[XPTR0]]
 ; CHECK: %[[YPTR0:[^ ]+]] = getelementptr inbounds %struct.CustomAttrs, %struct.CustomAttrs* %[[ATTRA0]], i32 0, i32 1
 ; CHECK: %[[YF1:[^ ]+]] = load float, float* %[[YATTRA]]
-; CHECK: store float %[[YF1]], float* %[[YPTR0]], align 4
+; CHECK: store float %[[YF1]], float* %[[YPTR0]]
 ; CHECK: %[[RQH1:[^ ]+]] = load i32, i32* %[[RQA]]
 ; CHECK: call void @"dx.hl.op..void (i32, %dx.types.HitObject*, i32, i32, %struct.CustomAttrs*)"(i32 363, %dx.types.HitObject* %{{[^ ]+}}, i32 %[[RQH1]], i32 16, %struct.CustomAttrs* %[[ATTRA0]])
 
@@ -140,7 +140,7 @@ target triple = "dxil-ms-dx"
 ; CHECK: store float %[[XF2]], float* %[[XPTR1]]
 ; CHECK: %[[YPTR1:[^ ]+]] = getelementptr inbounds %struct.CustomAttrs, %struct.CustomAttrs* %[[ATTRA1]], i32 0, i32 1
 ; CHECK: %[[YF2:[^ ]+]] = load float, float* %[[YATTRA]]
-; CHECK: store float %[[YF2]], float* %[[YPTR1]], align 4
+; CHECK: store float %[[YF2]], float* %[[YPTR1]]
 ; CHECK: %[[RQH2:[^ ]+]] = load i32, i32* %[[RQA]]
 ; CHECK: call void @"dx.hl.op..void (i32, %dx.types.HitObject*, i32, i32, %struct.CustomAttrs*)"(i32 363, %dx.types.HitObject* %{{[^ ]+}}, i32 %[[RQH2]], i32 17, %struct.CustomAttrs* %[[ATTRA1]])
 
diff --git a/tools/clang/test/HLSLFileCheck/shader_targets/raytracing/raytracing_intersection_geometryIndex.hlsl b/tools/clang/test/HLSLFileCheck/shader_targets/raytracing/raytracing_intersection_geometryIndex.hlsl
@@ -1,10 +1,10 @@
 // RUN: %dxc -T lib_6_5 -auto-binding-space 11 %s | FileCheck %s
 
 // CHECK: define void [[intersection1:@"\\01\?intersection1@[^\"]+"]]() #0 {
-// CHECK:   [[rayTCurrent:%[^ ]+]] = call float @dx.op.rayTCurrent.f32(i32 154)
-// CHECK:   [[GeometryIndex:%[^ ]+]] = call i32 @dx.op.geometryIndex.i32(i32 213)
-// CHECK:   icmp eq i32 [[GeometryIndex]], 0
-// CHECK:   call i1 @dx.op.reportHit.struct.MyAttributes(i32 158, float [[rayTCurrent]], i32 0, %struct.MyAttributes* nonnull {{.*}})
+// CHECK-DAG:   [[rayTCurrent:%[^ ]+]] = call float @dx.op.rayTCurrent.f32(i32 154)
+// CHECK-DAG:   [[GeometryIndex:%[^ ]+]] = call i32 @dx.op.geometryIndex.i32(i32 213)
+// CHECK-DAG:   icmp eq i32 [[GeometryIndex]], 0
+// CHECK-DAG:   call i1 @dx.op.reportHit.struct.MyAttributes(i32 158, float [[rayTCurrent]], i32 0, %struct.MyAttributes* nonnull {{.*}})
 // CHECK:   ret void
 
 struct MyAttributes {
