Support `IDifferentiablePtrType` #5031

saipraveenb25 · 2024-09-06T19:14:38Z

This patch adds initial support for IDifferentiablePtrType interface, which implements a part of the new diff type system proposal

Concretely, we add a new interface: IDifferentiablePtrType and a new pair type IDifferentialPtrPair<T : IDifferentiablePtrType> that can be used to properly represent differentiable reference, pointer, buffer & resource types that need to be transformed into pairs during auto-diff, but should not be used with the transposition logic.

Our forward-mode pass now treats types conforming to IDifferentiable as differentiable 'value' types and types conforming to IDifferentiablePtrType as differentiable 'ptr' types. Upon differentiation, the differential part of differentiable-ptr-typed insts are marked as 'primal' to instruct the later passes to move them to the primal part and not transpose them during the reverse-mode steps.

We also alter the primal context function's signature rules to translate T : IDifferentiablePtrType to the appropriate DifferentialPtrPair<T>. This is because we treat the full pair as a primal value that may be modified in the function, and so we need to run the context function to record it (if necessary).

Effectively, the semantics for our primal context functions have changed from "run primal computation and record intermediate values" to "run primal computation for differentiable value types & ptr-pair computations for differentiable ptr types and record necessary intermediate values"

Fixes: #4998

csyonghe · 2024-09-16T23:44:29Z

source/slang/slang-ast-support-types.h

@@ -1606,6 +1606,7 @@ namespace Slang
        DefaultInitializableConstructor, ///< The `IDefaultInitializable.__init()` method

        DifferentialType, ///< The `IDifferentiable.Differential` associated type requirement 
+        DifferentialPtrType, ///< The `IDifferentiable.DifferentialPtr` associated type requirement


Can remove DMulFunc down below?

I'm doing dmul removal in a separate patch

csyonghe · 2024-09-16T23:47:24Z

source/slang/slang-check-expr.cpp

        // Check if the provided type inherits from IDifferentiable.
        // If not, return the original type.
-        if (conformanceWitness)
+        if (auto conformanceValWitness = as<Witness>(


we should make isTypeDifferentiable function return the witness instead so we don't have to call isSubtype directly.
This should help improve readability.

csyonghe · 2024-09-16T23:48:07Z

source/slang/slang-ir-autodiff-fwd.cpp

@@ -16,6 +16,86 @@

 namespace Slang
 {
+
+IRInst* emitMakeDifferentialPair(IRBuilder* builder, IRType* pairType, IRInst* primalVal, IRInst* diffVal)


Can these functions be inside IRBuilder?

I didn't want to put this one in IRBuilder because it has fairly non-trivial logic.. IRBuilder has its own emitMakeDifferentialPair & emitMakeDifferentialPtrPair which emit the specific op.
This helper function decides which one to emit based on the conformance of the primal type. We need the DiffTypeConformanceContext to be around, so we can't easily move this into IRBuilder

This function isn't using DiffTypeConformanceContext, and feels like it fits well in IRBuilder.

We should move into the direction where IRBuilder can do more folding and instruction selection instead of emitting exactly the opcode being requested.

csyonghe · 2024-09-16T23:54:48Z

source/slang/slang-ir-autodiff-transcriber-base.cpp

@@ -445,8 +457,17 @@ IRType* AutoDiffTranscriberBase::differentiateExtractExistentialType(IRBuilder*
    auto interfaceType = as<IRInterfaceType>(unwrapAttributedType(origType->getOperand(0)->getDataType()));
    if (!interfaceType)
        return nullptr;
-    List<IRInterfaceRequirementEntry*> lookupKeyPath = differentiableTypeConformanceContext.findDifferentiableInterfaceLookupPath(
+    List<IRInterfaceRequirementEntry*> lookupPathValueType = differentiableTypeConformanceContext.findInterfaceLookupPath(


This seems not very efficient in that it attempts to lookup twice and then drop one of them. Can we at least not do the lookup for ptr conformance if value conformance is found?

I wonder why there is the need for findInterfaceLookupPath in the first place. For any interfaceType, if we want to know if interfaceType:IDifferentiable, there should be a SubtypeWitness that we can construct in the frontend and register into the dictionary, so we don't have to rediscover them in the backend.

This might be from before the inheritance flattening patch.. you can have interfaceType:middleInterface:IDifferentiable so this was for emitting lookup insts into the current block, to recurse into the sub-type witness, if the witness is a run-time variable.

Post-flattening, is there no need for this?

csyonghe · 2024-09-16T23:56:48Z

source/slang/slang-ir-autodiff-transcriber-base.cpp

@@ -1075,4 +1121,64 @@ InstPair AutoDiffTranscriberBase::transcribeInst(IRBuilder* builder, IRInst* ori
    return result;
 }

+
+void AutoDiffTranscriberBase::markDiffTypeInst(IRBuilder* builder, IRInst* diffInst, IRType* primalType)


Is all this function saying that anything other than a DifferentiablePtr type should be marked as differential?

It marks insts of T : IDifferentiable as differential (to move into the derivative function & transpose) and T : IDifferentiablePtrPair as primal (to move it to the primal blocks, and not consider it during transposition)

line 1147 is what I don't understand.
if diffInst isn't an IRType, then it can only be one of lookupWitness, specialization, extractExistentialType. How is this related to the subject of this function?

Ah, I think the misunderstanding may be due to what diffInst is supposed to be.
Here diffInst is any inst (not just a type) that is produced during forward-mode as the derivative counter-part for the primal inst. It is not the type of the inst, but the inst itself, which can be an IRType for cases where we construct types within the method body).
This method is responsible for deciding if this inst requires further processing or not by marking it appropriately.

Perhaps markDiffInst is a better thing to do here.

csyonghe · 2024-09-16T23:57:30Z

source/slang/slang-ir-autodiff.cpp

@@ -514,7 +665,8 @@ IRInst* DifferentiableTypeConformanceContext::getDifferentialTypeFromDiffPairTyp

 IRInst* DifferentiableTypeConformanceContext::getDiffTypeFromPairType(IRBuilder* builder, IRDifferentialPairTypeBase* type)
 {
-    return _getDiffTypeFromPairType(sharedContext, builder, type);
+    return this->differentiateType(builder, type->getValueType());
+    //return _getDiffTypeFromPairType(sharedContext, builder, type);


delete commented code.

csyonghe · 2024-09-16T23:57:41Z

source/slang/slang-ir-autodiff.cpp

@@ -722,7 +895,8 @@ void DifferentiableTypeConformanceContext::buildGlobalWitnessDictionary()
    {
        if (auto pairType = as<IRDifferentialPairTypeBase>(globalInst))
        {
-            differentiableWitnessDictionary.addIfNotExists(pairType->getValueType(), pairType->getWitness());
+            addTypeToDictionary(pairType->getValueType(), pairType->getWitness());
+            //differentiableWitnessDictionary.addIfNotExists(pairType->getValueType(), pairType->getWitness());


ccommented code.

…iablePtrInterfaceType`

csyonghe

I think a lot of duplication can be avoided if we consolidate the two dictionaries in type conformance context.

csyonghe · 2024-09-18T19:24:03Z

source/slang/slang-ir-autodiff-transcriber-base.cpp

 {
-    return differentiableTypeConformanceContext.tryGetDifferentiableWitness(builder, originalType);
+    if (kind == DiffConformanceKind::Any)


should differentiableTypeConformanceContext.tryGetDifferentiableWitness handle the Any case as well?

csyonghe · 2024-09-18T19:29:59Z

source/slang/slang-ir-autodiff-transcriber-base.cpp

@@ -1075,4 +1121,64 @@ InstPair AutoDiffTranscriberBase::transcribeInst(IRBuilder* builder, IRInst* ori
    return result;
 }

+
+void AutoDiffTranscriberBase::markDiffTypeInst(IRBuilder* builder, IRInst* diffInst, IRType* primalType)


line 1147 is what I don't understand.
if diffInst isn't an IRType, then it can only be one of lookupWitness, specialization, extractExistentialType. How is this related to the subject of this function?

csyonghe · 2024-09-18T19:31:38Z

source/slang/slang-ir-autodiff-fwd.cpp

@@ -16,6 +16,86 @@

 namespace Slang
 {
+
+IRInst* emitMakeDifferentialPair(IRBuilder* builder, IRType* pairType, IRInst* primalVal, IRInst* diffVal)


This function isn't using DiffTypeConformanceContext, and feels like it fits well in IRBuilder.

csyonghe · 2024-09-18T19:32:43Z

source/slang/slang-ir-autodiff-fwd.cpp

@@ -16,6 +16,86 @@

 namespace Slang
 {
+
+IRInst* emitMakeDifferentialPair(IRBuilder* builder, IRType* pairType, IRInst* primalVal, IRInst* diffVal)


We should move into the direction where IRBuilder can do more folding and instruction selection instead of emitting exactly the opcode being requested.

csyonghe · 2024-09-18T19:37:03Z

source/slang/slang-ir-autodiff.cpp

 {
    IRInst* foundResult = nullptr;
-    differentiableWitnessDictionary.tryGetValue(type, foundResult);
+


I still wonder why do we need two separate dictionaries. A type can either conform to valueType witness or reftype witness but not both. Why can't we consolidate the dictionaries so all lookups can be done in a single run?

csyonghe · 2024-09-18T19:37:18Z

source/slang/slang-ir-autodiff.cpp

 {
    if (isNoDiffType((IRType*)primalType))
        return nullptr;
-
-    IRInst* witness = lookUpConformanceForType((IRType*)primalType);
+


trailing space.

saipraveenb25 added 2 commits September 6, 2024 15:03

initial diff-ref-type interface

4ea736e

Initial support for IDifferentiablePtrType

ceaacab

saipraveenb25 added pr: new feature pr: non-breaking PRs without breaking changes labels Sep 6, 2024

saipraveenb25 and others added 10 commits September 6, 2024 15:52

Fix unused vars

53b9cfb

More tests + fix switch case fallthrough.

2cbf5f5

Update slang-ir-autodiff.cpp

96fa52d

Update diff-ptr-type-loop.slang

7c93b9d

Add optimization to allow more complex pair types

96a2791

Update slang-ir-autodiff-primal-hoist.cpp

6919026

Update diff-ptr-type-loop.slang

65c9f2c

Merge branch 'master' into idiff-ptr-types-3

04b338b

Update slang-ir-autodiff-primal-hoist.cpp

0d4a615

Merge branch 'master' into idiff-ptr-types-3

746f446

saipraveenb25 mentioned this pull request Sep 16, 2024

Update auto-diff user guide #5071

Merged

csyonghe reviewed Sep 17, 2024

View reviewed changes

saipraveenb25 and others added 5 commits September 17, 2024 14:23

More fixes to address reviews

7e79d00

Update slang-check-expr.cpp

4fc4af9

Merge branch 'master' into idiff-ptr-types-3

7587019

Merge branch 'master' into idiff-ptr-types-3

9361fcf

Optimizations + rename differentiableRefInterfaceType -> `different…

66507ec

…iablePtrInterfaceType`

saipraveenb25 requested a review from csyonghe September 18, 2024 18:13

csyonghe reviewed Sep 18, 2024

View reviewed changes

Move pair logic to ir-builder, unify the type dictionaries.

b12dbfc

csyonghe approved these changes Sep 19, 2024

View reviewed changes

Merge branch 'master' into idiff-ptr-types-3

6a3ca8f

csyonghe merged commit ccc310f into shader-slang:master Sep 19, 2024
12 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support `IDifferentiablePtrType` #5031

Support `IDifferentiablePtrType` #5031

saipraveenb25 commented Sep 6, 2024 •

edited

Loading

csyonghe Sep 16, 2024

saipraveenb25 Sep 17, 2024

csyonghe Sep 16, 2024

csyonghe Sep 16, 2024

saipraveenb25 Sep 17, 2024 •

edited

Loading

csyonghe Sep 18, 2024

csyonghe Sep 18, 2024

csyonghe Sep 16, 2024

saipraveenb25 Sep 17, 2024

csyonghe Sep 16, 2024

saipraveenb25 Sep 17, 2024

csyonghe Sep 18, 2024

saipraveenb25 Sep 18, 2024

csyonghe Sep 16, 2024

csyonghe Sep 16, 2024

csyonghe left a comment

csyonghe Sep 18, 2024

csyonghe Sep 18, 2024

csyonghe Sep 18, 2024

csyonghe Sep 18, 2024

csyonghe Sep 18, 2024

csyonghe Sep 18, 2024

Support IDifferentiablePtrType #5031

Support IDifferentiablePtrType #5031

Conversation

saipraveenb25 commented Sep 6, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

saipraveenb25 Sep 17, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

csyonghe left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Support `IDifferentiablePtrType` #5031

Support `IDifferentiablePtrType` #5031

saipraveenb25 commented Sep 6, 2024 •

edited

Loading

saipraveenb25 Sep 17, 2024 •

edited

Loading