5 files changed, 53 insertions, 5 deletions
diff --git a/source/slang/slang-check-expr.cpp b/source/slang/slang-check-expr.cpp
index 6b050aa89..9af8fd867 100644
--- a/source/slang/slang-check-expr.cpp
+++ b/source/slang/slang-check-expr.cpp
@@ -2712,7 +2712,7 @@ namespace Slang
         {
             innerExpr = parenExpr->base;
         }
-        if (!as<InvokeExpr>(innerExpr))
+        if (!as<InvokeExpr>(innerExpr) && !as<IndexExpr>(innerExpr))
         {
             getSink()->diagnose(expr, Diagnostics::invalidUseOfNoDiff);
         }
diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
index 1832a3b46..463c6f525 100644
--- a/source/slang/slang-diagnostic-defs.h
+++ b/source/slang/slang-diagnostic-defs.h
@@ -535,7 +535,7 @@ DIAGNOSTIC(38026, Error, globalTypeArgumentDoesNotConformToInterface, "type argu
 DIAGNOSTIC(38027, Error, mismatchExistentialSlotArgCount, "expected $0 existential slot arguments ($1 provided)")
 DIAGNOSTIC(38029, Error, typeArgumentDoesNotConformToInterface, "type argument '$0' does not conform to the required interface '$1'")
 
-DIAGNOSTIC(38031, Error, invalidUseOfNoDiff, "'no_diff' can only be used to decorate a call.")
+DIAGNOSTIC(38031, Error, invalidUseOfNoDiff, "'no_diff' can only be used to decorate a call or a subscript operation")
 DIAGNOSTIC(38032, Error, useOfNoDiffOnDifferentiableFunc, "use 'no_diff' on a call to a differentiable function has no meaning.")
 DIAGNOSTIC(38033, Error, cannotUseNoDiffInNonDifferentiableFunc, "cannot use 'no_diff' in a non-differentiable function.")
 
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index 4eb3982d3..a4563c254 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -3069,6 +3069,7 @@ public:
     IRInst* emitBackwardDifferentiatePrimalInst(IRType* type, IRInst* baseFn);
     IRInst* emitBackwardDifferentiatePropagateInst(IRType* type, IRInst* baseFn);
     IRInst* emitPrimalSubstituteInst(IRType* type, IRInst* baseFn);
+    IRInst* emitDetachDerivative(IRType* type, IRInst* value);
 
     IRInst* emitDispatchKernelInst(IRType* type, IRInst* baseFn, IRInst* threadGroupSize, IRInst* dispatchSize, Int argCount, IRInst* const* inArgs);
     IRInst* emitCudaKernelLaunch(IRInst* baseFn, IRInst* gridDim, IRInst* blockDim, IRInst* argsArray, IRInst* cudaStream);
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index 84730c913..a44667a79 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -3236,6 +3236,17 @@ namespace Slang
         return inst;
     }
 
+    IRInst *IRBuilder::emitDetachDerivative(IRType *type, IRInst *value)
+    {
+        auto inst = createInst<IRDetachDerivative>(
+            this,
+            kIROp_DetachDerivative,
+            type,
+            value);
+        addInst(inst);
+        return inst;
+    }
+
     IRInst* IRBuilder::emitBackwardDifferentiateInst(IRType* type, IRInst* baseFn)
     {
         auto inst = createInst<IRBackwardDifferentiate>(
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index d869bf60e..a1c7a2b8e 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -3333,9 +3333,45 @@ struct ExprLoweringVisitorBase : ExprVisitor<Derived, LoweredValInfo>
     LoweredValInfo visitTreatAsDifferentiableExpr(TreatAsDifferentiableExpr* expr)
     {
         auto baseVal = lowerSubExpr(expr->innerExpr);
-        SLANG_ASSERT(baseVal.flavor == LoweredValInfo::Flavor::Simple);
-        getBuilder()->addDecoration(baseVal.val, kIROp_TreatAsDifferentiableDecoration);
-        return baseVal;
+
+        IRInst* innerInst = nullptr;
+        if (baseVal.flavor != LoweredValInfo::Flavor::Simple)
+        {
+            if (!isLValueContext())
+            {
+                auto materializedVal = materialize(context, baseVal);
+
+                // TODO(Sai): We might be missing the case where a single materialize could create
+                // multiple calls (multiple index operations?). Not quite sure what the right way
+                // to handle that case might be.
+                // 
+                if (as<IRCall>(materializedVal.val))
+                    getBuilder()->addDecoration(materializedVal.val, kIROp_TreatAsDifferentiableDecoration);
+
+                innerInst = getSimpleVal(context, materializedVal);
+                
+                // We'll special case handle 'loads' here in order to allow TreatAsDifferentiable to be 
+                // used on array index operations. (This is to avoid a discrepancy between using no_diff
+                // on local variable indexing vs. resource indexing.)
+                // 
+                if (as<IRLoad>(innerInst))
+                    innerInst = getBuilder()->emitDetachDerivative(innerInst->getDataType(), innerInst);
+            }
+            else
+            {
+                SLANG_ASSERT("TreatAsDifferentiableExpr on non-simple l-values not properly defined.");
+            }
+        }
+        else
+        {
+            if (as<IRCall>(baseVal.val))
+                getBuilder()->addDecoration(baseVal.val, kIROp_TreatAsDifferentiableDecoration);
+            innerInst = baseVal.val;
+        }
+
+        SLANG_ASSERT(innerInst);
+        
+        return LoweredValInfo::simple(innerInst);
     }
 
     // Emit IR to denote the forward-mode derivative