Report AD checkpoint contexts

include/slang.h

@@ -852,6 +852,7 @@ extern "C"
  EmitIr, // bool
  ReportDownstreamTime, // bool
  ReportPerfBenchmark, // bool
+ ReportCheckpointIntermediates, // bool
  SkipSPIRVValidation, // bool
  SourceEmbedStyle,
  SourceEmbedName,

source/slang-record-replay/util/emum-to-string.h

@@ -149,6 +149,7 @@ namespace SlangRecord
  CASE(EmitIr);
  CASE(ReportDownstreamTime);
  CASE(ReportPerfBenchmark);
+ CASE(ReportCheckpointIntermediates);
  CASE(SkipSPIRVValidation);
  CASE(SourceEmbedStyle);
  CASE(SourceEmbedName);

source/slang/slang-compiler.cpp

@@ -2451,12 +2451,16 @@ namespace Slang
  return getTargetProgram()->getOptionSet().getBoolOption(CompilerOptionName::DumpIr);
  }
 
+ bool CodeGenContext::shouldReportCheckpointIntermediates()
+ {
+ return getTargetProgram()->getOptionSet().getBoolOption(CompilerOptionName::ReportCheckpointIntermediates);
+ }
+
  bool CodeGenContext::shouldDumpIntermediates()
  {
  return getTargetProgram()->getOptionSet().getBoolOption(CompilerOptionName::DumpIntermediates);
  }
 
-
  bool CodeGenContext::shouldTrackLiveness()
  {
  return getTargetProgram()->getOptionSet().getBoolOption(CompilerOptionName::TrackLiveness);

source/slang/slang-compiler.h

@@ -2728,6 +2728,7 @@ namespace Slang
 
  bool shouldValidateIR();
  bool shouldDumpIR();
+ bool shouldReportCheckpointIntermediates();
 
  bool shouldTrackLiveness();
 

source/slang/slang-diagnostic-defs.h

@@ -894,6 +894,12 @@ DIAGNOSTIC(58002, Error, unhandledGLSLSSBOType, "Unhandled GLSL Shader Storage B
 
 DIAGNOSTIC(58003, Error, inconsistentPointerAddressSpace, "'$0': use of pointer with inconsistent address space.")
 
+// Autodiff checkpoint reporting
+DIAGNOSTIC(-1, Note, reportCheckpointIntermediates, "checkpointing context of $1 bytes associated with function: '$0'")
+DIAGNOSTIC(-1, Note, reportCheckpointVariable, "$0 bytes ($1) used to checkpoint the following item:")
+DIAGNOSTIC(-1, Note, reportCheckpointCounter, "$0 bytes ($1) used for a loop counter here:")
+DIAGNOSTIC(-1, Note, reportCheckpointNone, "no checkpoint contexts to report")
+
 //
 // 8xxxx - Issues specific to a particular library/technology/platform/etc.
 //

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);
 

source/slang/slang-ir-addr-inst-elimination.cpp

@@ -69,30 +69,28 @@ struct AddressInstEliminationContext
  }
  }
 
- void transformLoadAddr(IRUse* use)
+ void transformLoadAddr(IRBuilder& builder, IRUse* use)
  {
  auto addr = use->get();
  auto load = as<IRLoad>(use->getUser());
 
- IRBuilder builder(module);
  builder.setInsertBefore(use->getUser());
  auto value = getValue(builder, addr);
  load->replaceUsesWith(value);
  load->removeAndDeallocate();
  }
 
- void transformStoreAddr(IRUse* use)
+ void transformStoreAddr(IRBuilder& builder, IRUse* use)
  {
  auto addr = use->get();
  auto store = as<IRStore>(use->getUser());
 
- IRBuilder builder(module);
  builder.setInsertBefore(use->getUser());
  storeValue(builder, addr, store->getVal());
  store->removeAndDeallocate();
  }
 
- void transformCallAddr(IRUse* use)
+ void transformCallAddr(IRBuilder& builder, IRUse* use)
  {
  auto addr = use->get();
  auto call = as<IRCall>(use->getUser());
@@ -103,7 +101,6 @@ struct AddressInstEliminationContext
  return;
  }
 
- IRBuilder builder(module);
  builder.setInsertBefore(call);
  auto tempVar = builder.emitVar(cast<IRPtrTypeBase>(addr->getFullType())->getValueType());
 
@@ -155,17 +152,20 @@ struct AddressInstEliminationContext
  use = nextUse;
  continue;
  }
+
+ IRBuilder transformBuilder(module);
+ IRBuilderSourceLocRAII sourceLocationScope(&transformBuilder, use->getUser()->sourceLoc);
 
  switch (use->getUser()->getOp())
  {
  case kIROp_Load:
- transformLoadAddr(use);
+ transformLoadAddr(transformBuilder, use);
  break;
  case kIROp_Store:
- transformStoreAddr(use);
+ transformStoreAddr(transformBuilder, use);
  break;
  case kIROp_Call:
- transformCallAddr(use);
+ transformCallAddr(transformBuilder, use);
  break;
  case kIROp_GetElementPtr:
  case kIROp_FieldAddress:

source/slang/slang-ir-autodiff-primal-hoist.cpp

@@ -3,8 +3,9 @@
 #include "slang-ir-autodiff-region.h"
 #include "slang-ir-simplify-cfg.h"
 #include "slang-ir-util.h"
-#include "../core/slang-func-ptr.h"
+#include "slang-ir-insts.h"
 #include "slang-ir.h"
+#include "../core/slang-func-ptr.h"
 
 namespace Slang
 {
@@ -1092,7 +1093,8 @@ IRType* getTypeForLocalStorage(
 IRVar* emitIndexedLocalVar(
  IRBlock* varBlock,
  IRType* baseType,
- const List<IndexTrackingInfo>& defBlockIndices)
+ const List<IndexTrackingInfo>& defBlockIndices,
+ SourceLoc location)
 {
  // Cannot store pointers. Case should have been handled by now.
  SLANG_RELEASE_ASSERT(!as<IRPtrTypeBase>(baseType));
@@ -1101,6 +1103,8 @@ IRVar* emitIndexedLocalVar(
  SLANG_RELEASE_ASSERT(!as<IRTypeType>(baseType));
 
  IRBuilder varBuilder(varBlock->getModule());
+ IRBuilderSourceLocRAII sourceLocationScope(&varBuilder, location);
+
  varBuilder.setInsertBefore(varBlock->getFirstOrdinaryInst());
 
  IRType* varType = getTypeForLocalStorage(&varBuilder, baseType, defBlockIndices);
@@ -1179,9 +1183,14 @@ IRVar* storeIndexedValue(
  IRInst* instToStore,
  const List<IndexTrackingInfo>& defBlockIndices)
 {
- IRVar* localVar = emitIndexedLocalVar(defaultVarBlock, instToStore->getDataType(), defBlockIndices);
+ IRVar* localVar = emitIndexedLocalVar(defaultVarBlock,
+ instToStore->getDataType(),
+ defBlockIndices,
+ instToStore->sourceLoc);
 
- IRInst* addr = emitIndexedStoreAddressForVar(builder, localVar, defBlockIndices);
+ IRInst* addr = emitIndexedStoreAddressForVar(builder,
+ localVar,
+ defBlockIndices);
 
  builder->emitStore(addr, instToStore);
 
@@ -1574,12 +1583,16 @@ RefPtr<HoistedPrimalsInfo> ensurePrimalAvailability(
  // region, that means there's no need to allocate a fully indexed var.
  // 
  defBlockIndices = maybeTrimIndices(defBlockIndices, indexedBlockInfo, outOfScopeUses);
-
- IRVar* localVar = storeIndexedValue(
- &builder,
- varBlock,
- builder.emitLoad(varToStore),
- defBlockIndices);
+
+ IRVar* localVar = nullptr;
+ {
+ IRBuilderSourceLocRAII sourceLocationScope(&builder, varToStore->sourceLoc);
+ localVar = storeIndexedValue(
+ &builder,
+ varBlock,
+ builder.emitLoad(varToStore),
+ defBlockIndices);
+ }
 
  for (auto use : outOfScopeUses)
  {
@@ -1626,6 +1639,8 @@ RefPtr<HoistedPrimalsInfo> ensurePrimalAvailability(
  }
  else
  {
+ IRBuilderSourceLocRAII sourceLocationScope(&builder, instToStore->sourceLoc);
+
  // Handle the special case of loop counters.
  // The only case where there will be a reference of primal loop counter from rev blocks
  // is the start of a loop in the reverse code. Since loop counters are not considered a
@@ -1643,6 +1658,8 @@ RefPtr<HoistedPrimalsInfo> ensurePrimalAvailability(
 
  setInsertAfterOrdinaryInst(&builder, instToStore);
  auto localVar = storeIndexedValue(&builder, varBlock, instToStore, defBlockIndices);
+ if (isLoopCounter)
+ builder.addLoopCounterDecoration(localVar);
 
  for (auto use : outOfScopeUses)
  {
@@ -1728,6 +1745,8 @@ static IRBlock* getUpdateBlock(IRLoop* loop)
 void lowerIndexedRegion(IRLoop*& primalLoop, IRLoop*& diffLoop, IRInst*& primalCountParam, IRInst*& diffCountParam)
 {
  IRBuilder builder(primalLoop);
+ IRBuilderSourceLocRAII sourceLocationScope(&builder, primalLoop->sourceLoc);
+
  primalCountParam = nullptr;
 
  // Grab first primal block.
@@ -1899,8 +1918,7 @@ RefPtr<HoistedPrimalsInfo> applyCheckpointPolicy(IRGlobalValueWithCode* func)
  // Legalize the primal inst accesses by introducing local variables / arrays and emitting
  // necessary load/store logic.
  //
- primalsInfo = ensurePrimalAvailability(primalsInfo, func, indexedBlockInfo);
- return primalsInfo;
+ return ensurePrimalAvailability(primalsInfo, func, indexedBlockInfo);
 }
 
 void DefaultCheckpointPolicy::preparePolicy(IRGlobalValueWithCode* func)