From 6e4eae1050ab9282b460a33a013652c387c1e585 Mon Sep 17 00:00:00 2001 From: Yong He Date: Thu, 23 Mar 2023 17:16:32 -0700 Subject: [PATCH] Hack handling of primal insts that has a function type. (#2728) * Update diff-bwd material test * Minor update * Hack handling of primal insts that has a function type. --------- Co-authored-by: winmad Co-authored-by: Yong He --- .../slang-ir-autodiff-transcriber-base.cpp | 7 ++-- source/slang/slang-ir-autodiff-transpose.h | 27 +++++++++++++-- .../material/GlossyMaterialInstance.slang | 10 +++++- tests/autodiff/material/IBSDF.slang | 3 ++ .../diff-bwd-falcor-material-system.slang | 33 +++++++++++++++++++ ...-falcor-material-system.slang.expected.txt | 7 ++++ .../diff-falcor-material-system.slang | 2 +- 7 files changed, 83 insertions(+), 6 deletions(-) create mode 100644 tests/autodiff/material/diff-bwd-falcor-material-system.slang create mode 100644 tests/autodiff/material/diff-bwd-falcor-material-system.slang.expected.txt diff --git a/source/slang/slang-ir-autodiff-transcriber-base.cpp b/source/slang/slang-ir-autodiff-transcriber-base.cpp index 2bc67e5613..552ac762c7 100644 --- a/source/slang/slang-ir-autodiff-transcriber-base.cpp +++ b/source/slang/slang-ir-autodiff-transcriber-base.cpp @@ -1116,9 +1116,12 @@ IRInst* AutoDiffTranscriberBase::transcribe(IRBuilder* builder, IRInst* origInst } else { - if (!pair.primal->findDecoration() - && !as(pair.differential)) + if (!pair.primal->findDecoration()) { + if (as(pair.differential)) + break; + if (as(pair.differential)) + break; auto mixedType = (IRType*)(pair.primal->getDataType()); builder->markInstAsMixedDifferential(pair.primal, mixedType); } diff --git a/source/slang/slang-ir-autodiff-transpose.h b/source/slang/slang-ir-autodiff-transpose.h index 6535899337..bbdb01290e 100644 --- a/source/slang/slang-ir-autodiff-transpose.h +++ b/source/slang/slang-ir-autodiff-transpose.h @@ -950,8 +950,25 @@ struct DiffTransposePass // Slang doesn't support function values. So if we see a func-typed inst // it's proabably a reference to a function. // - if (as(child->getDataType())) + switch (child->getOp()) + { + /* + TODO: need a better way to move specialize, lookupwitness, extractExistentialType/Value/Witness + insts to a proper location that dominates all their use sites. Create copies of these insts + when necessary. + case kIROp_Specialize: + case kIROp_LookupWitness: + case kIROp_ExtractExistentialType: + case kIROp_ExtractExistentialValue: + case kIROp_ExtractExistentialWitnessTable: + */ + case kIROp_ForwardDifferentiate: + case kIROp_BackwardDifferentiate: + case kIROp_BackwardDifferentiatePrimal: + case kIROp_BackwardDifferentiatePropagate: typeInsts.add(child); + break; + } } for (auto inst : typeInsts) @@ -959,7 +976,6 @@ struct DiffTransposePass inst->insertAtEnd(revBlock); } - // Then, go backwards through the regular instructions, and transpose them into the new // rev block. // Note the 'reverse' traversal here. @@ -2221,6 +2237,10 @@ struct DiffTransposePass case kIROp_ifElse: case kIROp_loop: case kIROp_Switch: + case kIROp_LookupWitness: + case kIROp_ExtractExistentialType: + case kIROp_ExtractExistentialValue: + case kIROp_ExtractExistentialWitnessTable: { // Ignore. transposeBlock() should take care of adding the // appropriate branch instruction. @@ -3474,6 +3494,9 @@ struct DiffTransposePass List primalUsesToHoist; Dictionary mapStoreToDefBlock; + + IRCloneEnv typeInstCloneEnv = {}; + }; diff --git a/tests/autodiff/material/GlossyMaterialInstance.slang b/tests/autodiff/material/GlossyMaterialInstance.slang index b1430e62c7..fbbcc1a268 100644 --- a/tests/autodiff/material/GlossyMaterialInstance.slang +++ b/tests/autodiff/material/GlossyMaterialInstance.slang @@ -7,18 +7,26 @@ struct GlossyBSDF : IBSDF { float3 albedo; - [BackwardDifferentiable] float3 getAlbedo() { return albedo; } [ForwardDerivativeOf(getAlbedo)] + [TreatAsDifferentiable] DifferentialPair __fwd_d_getAlbedo() { return diffPair(albedo, float3(1.f)); } + [BackwardDerivativeOf(getAlbedo)] + [TreatAsDifferentiable] + void __bwd_d_getAlbedo(float3 dOut) + { + [unroll] + for (int j = 0; j < 3; j++) outputBuffer[j + 3] += dOut[j]; + } + [BackwardDifferentiable] float3 eval(const float3 wiLocal, const float3 woLocal) { diff --git a/tests/autodiff/material/IBSDF.slang b/tests/autodiff/material/IBSDF.slang index 41b0957db4..9bdeb91974 100644 --- a/tests/autodiff/material/IBSDF.slang +++ b/tests/autodiff/material/IBSDF.slang @@ -1,5 +1,8 @@ //TEST_IGNORE_FILE: +//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + interface IBSDF { [BackwardDifferentiable] diff --git a/tests/autodiff/material/diff-bwd-falcor-material-system.slang b/tests/autodiff/material/diff-bwd-falcor-material-system.slang new file mode 100644 index 0000000000..6d4101d121 --- /dev/null +++ b/tests/autodiff/material/diff-bwd-falcor-material-system.slang @@ -0,0 +1,33 @@ +//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type +//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type + +// outputBuffer is defined in IBSDF.slang +//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer + +import MaterialSystem; +import DiffuseMaterial; +import DiffuseMaterialInstance; +import GlossyMaterial; +import GlossyMaterialInstance; + +//TEST_INPUT: type_conformance DiffuseMaterial:IMaterial = 0 +//TEST_INPUT: type_conformance GlossyMaterial:IMaterial = 1 + +[BackwardDifferentiable] +float3 evalBSDF(int type) +{ + float3 wi = normalize(float3(0.5, 0.2, 0.8)); + float3 wo = normalize(float3(-0.1, -0.3, 0.9)); + + IMaterial material = createMaterialClassConformance(type, float3(0.9f, 0.6f, 0.2f)); + let mi = material.setupMaterialInstance(); + float3 f = mi.eval(wi, wo); + return f; +} + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID) +{ + __bwd_diff(evalBSDF)(0, float3(1.f)); + __bwd_diff(evalBSDF)(1, float3(1.f)); +} diff --git a/tests/autodiff/material/diff-bwd-falcor-material-system.slang.expected.txt b/tests/autodiff/material/diff-bwd-falcor-material-system.slang.expected.txt new file mode 100644 index 0000000000..e600054772 --- /dev/null +++ b/tests/autodiff/material/diff-bwd-falcor-material-system.slang.expected.txt @@ -0,0 +1,7 @@ +type: float +0.0 +0.0 +0.0 +1.8 +1.2 +0.4 \ No newline at end of file diff --git a/tests/autodiff/material/diff-falcor-material-system.slang b/tests/autodiff/material/diff-falcor-material-system.slang index 1cb8275fae..875432c280 100644 --- a/tests/autodiff/material/diff-falcor-material-system.slang +++ b/tests/autodiff/material/diff-falcor-material-system.slang @@ -1,8 +1,8 @@ //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type +// outputBuffer is defined in IBSDF.slang //TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer -RWStructuredBuffer outputBuffer; import MaterialSystem; import DiffuseMaterial;