diff options
| -rw-r--r-- | source/slang/slang-ir-autodiff-fwd.cpp | 40 | ||||
| -rw-r--r-- | source/slang/slang-ir-autodiff-fwd.h | 2 | ||||
| -rw-r--r-- | source/slang/slang-ir-autodiff-transpose.h | 34 | ||||
| -rw-r--r-- | tests/autodiff/select.slang | 38 | ||||
| -rw-r--r-- | tests/autodiff/select.slang.expected.txt | 6 |
5 files changed, 120 insertions, 0 deletions
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) { diff --git a/tests/autodiff/select.slang b/tests/autodiff/select.slang new file mode 100644 index 000000000..261c170db --- /dev/null +++ b/tests/autodiff/select.slang @@ -0,0 +1,38 @@ +//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type +//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type + +//TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer +RWStructuredBuffer<float> outputBuffer; + +typedef DifferentialPair<float> dpfloat; +typedef float.Differential dfloat; + +[BackwardDifferentiable] +float f(float x, float y) +{ + return x > 0.5 ? x : y; +} + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) +{ + { + dpfloat dpa = dpfloat(2.0, 1.0); + dpfloat dpb = dpfloat(0.5, 1.0); + + __bwd_diff(f)(dpa, dpb, 1.0); + + outputBuffer[0] = dpa.d; // Expect: 1.0 + outputBuffer[1] = dpb.d; // Expect: 0.0 + } + + { + dpfloat dpa = dpfloat(0.3, 1.0); + dpfloat dpb = dpfloat(0.3, 1.0); + + __bwd_diff(f)(dpa, dpb, 1.0); + + outputBuffer[2] = dpa.d; // Expect: 0.0 + outputBuffer[3] = dpb.d; // Expect: 1.0 + } +} diff --git a/tests/autodiff/select.slang.expected.txt b/tests/autodiff/select.slang.expected.txt new file mode 100644 index 000000000..52dd2d9ca --- /dev/null +++ b/tests/autodiff/select.slang.expected.txt @@ -0,0 +1,6 @@ +type: float +1.000000 +0.000000 +0.000000 +1.000000 +0.000000 |