From b808aa4df50d46eaa569561f7e464c55c1c2d72a Mon Sep 17 00:00:00 2001
From: venkataram-nv <vedavamadath@nvidia.com>
Date: Wed, 18 Sep 2024 20:42:07 -0700
Subject: Report AD checkpoint contexts (#5058)

* Transferring source locations when creating phi instructions

* Tracking for simple variables

* Deriving source locations for loop counters

* Printing checkpoint structure breakdown

* More readable output format

* Special behavior for loop counters

* Writing report to file

* Add slangc option to enable checkpoint reports

* Display types of checkpointed fields

* Message in case there are no checkpointing contexts

* Catch source locations for function calls

* Source cleanup

* Fix compilation warnings

* Remove stray dump()

* Provide the report through diagnostic notes

* Add missing path for sourceLoc during unzip pass

* Add tests for reporting intermediates

* Include more transfer cases for source locations

* Fix ordering in address elimination

* Fill in more holes with source location transfer

* Remove debugging line

* Reverting changes to diagnostic sink

* Simplify address elimination using source location RAII contexts

* Eliminating manual source loc transfers in forward transcription

* Fix local var adaptation to use RAII location setter

* Simplify primal hoisting logic for source location transfer

* Simplify unzipping with RAII location scopes

* Simplify transpose logic

* Cleaning up for rev.cpp

* Reverting spacing changes

* Fix mistake with source loc RAII instantiation

* Fix formatting issues
---
 source/slang/slang-emit.cpp | 67 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)

(limited to 'source/slang/slang-emit.cpp')

diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index cdd2ca5b6..6e3556064 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -34,6 +34,7 @@
 #include "slang-ir-wgsl-legalize.h"
 #include "slang-ir-insts.h"
 #include "slang-ir-inline.h"
+#include "slang-ir-layout.h"
 #include "slang-ir-legalize-array-return-type.h"
 #include "slang-ir-legalize-mesh-outputs.h"
 #include "slang-ir-legalize-varying-params.h"
@@ -214,6 +215,68 @@ static void dumpIRIfEnabled(
     }
 }
 
+static void reportCheckpointIntermediates(CodeGenContext* codeGenContext, DiagnosticSink* sink, IRModule* irModule)
+{
+    // Report checkpointing information
+    CompilerOptionSet& optionSet = codeGenContext->getTargetProgram()->getOptionSet();
+    SourceManager* sourceManager = sink->getSourceManager();
+    
+    SourceWriter typeWriter(sourceManager, LineDirectiveMode::None, nullptr);
+
+    CLikeSourceEmitter::Desc description;
+    description.codeGenContext = codeGenContext;
+    description.sourceWriter = &typeWriter;
+
+    CPPSourceEmitter emitter(description);
+
+    int nonEmptyStructs = 0;
+    for (auto inst : irModule->getGlobalInsts())
+    {
+        IRStructType *structType = as<IRStructType>(inst);
+        if (!structType)
+            continue;
+
+        auto checkpointDecoration = structType->findDecoration<IRCheckpointIntermediateDecoration>();
+        if (!checkpointDecoration)
+            continue;
+
+        IRSizeAndAlignment structSize;
+        getNaturalSizeAndAlignment(optionSet, structType, &structSize);
+
+        // Reporting happens before empty structs are optimized out
+        // and we still want to keep the checkpointing decorations,
+        // so we end up needing to check for non-zero-ness
+        if (structSize.size == 0)
+            continue;
+
+        auto func = checkpointDecoration->getSourceFunction();
+        sink->diagnose(structType, Diagnostics::reportCheckpointIntermediates, func, structSize.size);
+        nonEmptyStructs++;
+
+        for (auto field : structType->getFields())
+        {
+            IRType *fieldType = field->getFieldType();
+            IRSizeAndAlignment fieldSize;
+            getNaturalSizeAndAlignment(optionSet, fieldType, &fieldSize);
+            if (fieldSize.size == 0)
+                continue;
+
+            typeWriter.clearContent();
+            emitter.emitType(fieldType);
+
+            sink->diagnose(field->sourceLoc,
+                field->findDecoration<IRLoopCounterDecoration>()
+                    ? Diagnostics::reportCheckpointCounter
+                    : Diagnostics::reportCheckpointVariable,
+                fieldSize.size,
+                typeWriter.getContent());
+        }
+    }
+
+    if (nonEmptyStructs == 0)
+        sink->diagnose(SourceLoc(), Diagnostics::reportCheckpointNone);
+}
+
 struct LinkingAndOptimizationOptions
 {
     bool shouldLegalizeExistentialAndResourceTypes = true;
@@ -767,6 +830,10 @@ Result linkAndOptimizeIR(
             break;
     }
 
+    // Report checkpointing information
+    if (codeGenContext->shouldReportCheckpointIntermediates())
+        reportCheckpointIntermediates(codeGenContext, sink, irModule);
+
     if (requiredLoweringPassSet.autodiff)
         finalizeAutoDiffPass(targetProgram, irModule);
 
-- 
cgit v1.2.3