Add cuda-host decoration for the synthesized constructor

Add cuda-host decoration for the synthesized constructor if there is any usage of TorchTensor.
shader-slang · Jan 17, 2025 · adec731 · adec731
1 parent 99acef4
commit adec731
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/source/slang/slang-ir-check-differentiability.cpp b/source/slang/slang-ir-check-differentiability.cpp
@@ -249,13 +249,25 @@ struct CheckDifferentiabilityPassContext : public InstPassBase
         if (outerFuncInst->findDecoration<IRTorchEntryPointDecoration>())
             return;
 
+        bool isSynthesizeConstructor = false;
+
+        if(auto constructor = outerFuncInst->findDecoration<IRConstructorDecorartion>())
+            isSynthesizeConstructor = constructor->getSynthesizedStatus();
+
         // This is a kernel function, we don't allow using TorchTensor type here.
         for (auto b : funcInst->getBlocks())
         {
             for (auto inst : b->getChildren())
             {
                 if (!checkType(inst->getDataType()))
                 {
+                    if (isSynthesizeConstructor)
+                    {
+                        IRBuilder irBuilder(funcInst);
+                        irBuilder.addDecoration(funcInst, kIROp_CudaHostDecoration);
+                        return;
+                    }
+
                     auto loc = inst->sourceLoc;
                     if (!loc.isValid())
                         loc = funcInst->sourceLoc;