From b1b76f06ca5bdfc4b688d99095dfb7d561a04d80 Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Sat, 25 Feb 2023 02:23:41 -0500
Subject: Added support for `kIROp_Select` (#2679)

---
 source/slang/slang-ir-autodiff-fwd.cpp     | 40 ++++++++++++++++++++++++++++++
 source/slang/slang-ir-autodiff-fwd.h       |  2 ++
 source/slang/slang-ir-autodiff-transpose.h | 34 +++++++++++++++++++++++++
 3 files changed, 76 insertions(+)

(limited to 'source')

diff --git a/source/slang/slang-ir-autodiff-fwd.cpp b/source/slang/slang-ir-autodiff-fwd.cpp
index b2f420730..72bbe4530 100644
--- a/source/slang/slang-ir-autodiff-fwd.cpp
+++ b/source/slang/slang-ir-autodiff-fwd.cpp
@@ -177,6 +177,43 @@ InstPair ForwardDiffTranscriber::transcribeBinaryLogic(IRBuilder* builder, IRIns
     return InstPair(primalLogic, nullptr);
 }
 
+InstPair ForwardDiffTranscriber::transcribeSelect(IRBuilder* builder, IRInst* origSelect)
+{
+    auto primalCondition = lookupPrimalInst(builder, origSelect->getOperand(0));
+
+    auto origLeft = origSelect->getOperand(1);
+    auto origRight = origSelect->getOperand(2);
+
+    auto primalLeft = findOrTranscribePrimalInst(builder, origLeft);
+    auto primalRight = findOrTranscribePrimalInst(builder, origRight);
+
+    auto diffLeft = findOrTranscribeDiffInst(builder, origLeft);
+    auto diffRight = findOrTranscribeDiffInst(builder, origRight);
+
+    auto primalSelect = maybeCloneForPrimalInst(builder, origSelect);
+
+    auto resultType = primalCondition->getDataType();
+
+    // If both sides have no differential, skip
+    if (diffLeft || diffRight)
+    {
+        diffLeft = diffLeft ? diffLeft : getDifferentialZeroOfType(builder, primalLeft->getDataType());
+        diffRight = diffRight ? diffRight : getDifferentialZeroOfType(builder, primalRight->getDataType());
+
+        auto diffType = (IRType*) differentiableTypeConformanceContext.getDifferentialForType(builder, resultType);
+
+        return InstPair(
+            primalSelect,
+            builder->emitIntrinsicInst(
+                diffType,
+                kIROp_Select,
+                3,
+                List<IRInst*>(primalCondition, diffLeft, diffRight).getBuffer()));
+    }
+    
+    return InstPair(primalSelect, nullptr);
+}
+
 InstPair ForwardDiffTranscriber::transcribeLoad(IRBuilder* builder, IRLoad* origLoad)
 {
     auto origPtr = origLoad->getPtr();
@@ -1365,6 +1402,9 @@ InstPair ForwardDiffTranscriber::transcribeInstImpl(IRBuilder* builder, IRInst*
     case kIROp_Eql:
     case kIROp_Neq:
         return transcribeBinaryLogic(builder, origInst);
+    
+    case kIROp_Select:
+        return transcribeSelect(builder, origInst);
 
     case kIROp_MakeVector:
     case kIROp_MakeMatrix:
diff --git a/source/slang/slang-ir-autodiff-fwd.h b/source/slang/slang-ir-autodiff-fwd.h
index 9089c4412..5b79a6c54 100644
--- a/source/slang/slang-ir-autodiff-fwd.h
+++ b/source/slang/slang-ir-autodiff-fwd.h
@@ -35,6 +35,8 @@ struct ForwardDiffTranscriber : AutoDiffTranscriberBase
 
     InstPair transcribeBinaryLogic(IRBuilder* builder, IRInst* origLogic);
 
+    InstPair transcribeSelect(IRBuilder* builder, IRInst* origSelect);
+
     InstPair transcribeLoad(IRBuilder* builder, IRLoad* origLoad);
 
     InstPair transcribeStore(IRBuilder* builder, IRStore* origStore);
diff --git a/source/slang/slang-ir-autodiff-transpose.h b/source/slang/slang-ir-autodiff-transpose.h
index e1832b9eb..23f57032d 100644
--- a/source/slang/slang-ir-autodiff-transpose.h
+++ b/source/slang/slang-ir-autodiff-transpose.h
@@ -1674,6 +1674,9 @@ struct DiffTransposePass
             case kIROp_Div: 
             case kIROp_Neg:
                 return transposeArithmetic(builder, fwdInst, revValue);
+            
+            case kIROp_Select:
+                return transposeSelect(builder, fwdInst, revValue);
 
             case kIROp_Call:
                 return transposeCall(builder, as<IRCall>(fwdInst), revValue);
@@ -2281,6 +2284,37 @@ struct DiffTransposePass
         }
     }
 
+    
+    TranspositionResult transposeSelect(IRBuilder* builder, IRInst* fwdInst, IRInst* revValue)
+    {
+        auto primalCondition = fwdInst->getOperand(0);
+
+        auto leftZero = emitDZeroOfDiffInstType(builder, tryGetPrimalTypeFromDiffInst(fwdInst->getOperand(1)));
+        auto leftGradientInst = builder->emitIntrinsicInst(
+            fwdInst->getOperand(1)->getDataType(),
+            kIROp_Select,
+            3,
+            List<IRInst*>(primalCondition, revValue, leftZero).getBuffer());
+        
+        auto rightZero = emitDZeroOfDiffInstType(builder, tryGetPrimalTypeFromDiffInst(fwdInst->getOperand(2)));
+        auto rightGradientInst = builder->emitIntrinsicInst(
+            fwdInst->getOperand(2)->getDataType(),
+            kIROp_Select,
+            3,
+            List<IRInst*>(primalCondition, rightZero, revValue).getBuffer());
+        
+        return TranspositionResult(
+                        List<RevGradient>(
+                            RevGradient(
+                                fwdInst->getOperand(1),
+                                leftGradientInst,
+                                fwdInst),
+                            RevGradient(
+                                fwdInst->getOperand(2),
+                                rightGradientInst,
+                                fwdInst)));
+    }
+
     TranspositionResult transposeArithmetic(IRBuilder* builder, IRInst* fwdInst, IRInst* revValue)
     {
         
-- 
cgit v1.2.3