module PartitionedRegression.Build.CoordinateDevice import Std.Numeric.Profile import Representation.ParticipantGroup import Learning.Numerics import PartitionedRegression.Coordinates import PartitionedRegression.CoordinateHost import Runtime.OrderedCombine import Data.SHA256Digest import Data.SHA256Schedule import Data.SHA256ScheduleExpand import Data.SHA256Compress import Data.SHA256Core import Compiler.Planning.OrderedCombine import PartitionedRegression.Build.ScalarEvaluation import PartitionedRegression.Build.CoordinateMemory import Platform.Linux.Nvidia.ProgramReadback import Checkpoint.NativeRegions import Compiler.Planning.RegionFrame import Hardware.Nvidia.SM86.Command.WholeProgramPlan import Platform.Linux.Nvidia.DeviceBinding import Platform.Linux.Nvidia.ResourcePlan import Platform.Linux.Nvidia.LinearStep import Platform.Linux.Nvidia.ProgramPlan import Platform.Linux.Nvidia.PlanHost import Platform.Linux.Nvidia.PlanHostRequest import Platform.RunPod.Nvidia.SM86 import Runtime.ArenaCertificate import Runtime.NativeFloatPlan import Runtime.NativeMemoryPlan import Runtime.NativePeerEndpoint import Runtime.NativePhysicalEmbeddedArtifact import Runtime.NativePhysicalNativeELF import Runtime.NativePhysicalProgram import Runtime.NativeLaunchRecipeLoad import Runtime.RegionFrame import Runtime.SocketAddress import Model.Word32 import Model.Word64 import Std.Foundation import Std.Natural -- This policy is part of the frame/checkpoint identity. Each local product -- rounds its operands to binary32 on the explicit GPU realization and widens -- the device result. Residual summation and persistent updates use binary64. def cdNumerics : Bytes = b"f32-products-f64-accumulation-v1" def cdIdentity = (lambda unrestricted kind : (family NativePhysicalEmbeddedComponentKind) . (let unrestricted name = (eliminate NativePhysicalEmbeddedComponentKind (lambda unrestricted current : (family NativePhysicalEmbeddedComponentKind) . Bytes) kind (branch NativePhysicalEmbeddedHostELF . b"coordinate-scalar-sm86-host-v1") (branch NativePhysicalEmbeddedProgramTable . b"coordinate-scalar-sm86-program-v1") (branch NativePhysicalEmbeddedQMDTable . b"coordinate-scalar-sm86-qmd-v1") (branch NativePhysicalEmbeddedPushbuffer . b"coordinate-scalar-sm86-push-v1") (branch NativePhysicalEmbeddedGPFIFO . b"coordinate-scalar-sm86-ring-v1")) in (bytes-append name (phZeros (naturalSaturatingSubtract nativePhysicalEmbeddedIdentityLength (bytes-length name)))))) def cdBindingAdmitted = (lambda unrestricted device : (family NvidiaDeviceBinding) . (naturalAnd (nvidiaDeviceBindingAdmitted device) (naturalEqual (bytes-length (nvidiaDeviceBindingUUID device)) 16))) def cdCounter = (lambda unrestricted end : Nat . (frameAlign end)) def cdScratch = (lambda unrestricted end : Nat . (naturalAdd (cdCounter end) coordinateWordBytes)) def cdExtent = (lambda unrestricted end : Nat . (naturalAdd (cdScratch end) coordinateWordBytes)) -- State addresses are relative to the native interpreter arena, separate from -- the target's independently certified GPU/host virtual mappings. The native -- session prefix and the needed counter/scratch words cannot alias learning state. def cdSessionResidents = (lambda unrestricted end : Nat . (lambda unrestricted extent : Nat . (lambda unrestricted residents : (family ArenaResidents) . (constructor ArenaResidents ArenaResidentsNext (constructor ArenaResident ArenaResidentValue b"native-device-session" 0 phStateExtent coordinateWordBytes) (constructor ArenaResidents ArenaResidentsNext (constructor ArenaResident ArenaResidentValue b"device-local-state" (cdCounter end) (naturalSaturatingSubtract extent (cdCounter end)) coordinateWordBytes) residents))))) -- Training needs only an issue counter; prediction also uses its explicit -- conversion scratch. A pairing retains only the host state it actually uses. def cdTrainingExtent = (lambda unrestricted end : Nat . (naturalAdd (cdCounter end) coordinateWordBytes)) def cdTrainingResidents = (lambda unrestricted end : Nat . (cdSessionResidents end (cdTrainingExtent end))) def cdResidents = (lambda unrestricted end : Nat . (cdSessionResidents end (cdExtent end))) def cdPredictionResidents = (lambda unrestricted layout : (family CoordinatePredictionLayout) . (constructor ArenaResidents ArenaResidentsNext (constructor ArenaResident ArenaResidentValue b"prediction-values" (cpLeft layout) (naturalSaturatingSubtract (cpExtent layout) (cpLeft layout)) coordinateWordBytes) (nativeRegionResidents (cpCheckpoint layout)))) -- A scalar prediction field is a product, suitable for both the forward -- partial and the backward residual-times-feature. Choosing the gradient -- field here would multiply by the input twice and change the learner. def cdProductResult : (family NvidiaLinearScalarResult) = (constructor NvidiaLinearScalarResult NvidiaLinearScalarPrediction) def cdProduct = (lambda unrestricted gpu : (family NvidiaPlanHostLayout) . (lambda unrestricted entry : Nat . (lambda unrestricted data : Nat . (lambda unrestricted inputOffset : Nat . (lambda unrestricted targetOffset : Nat . (lambda unrestricted outputOffset : Nat . (lambda unrestricted end : Nat . (lambda unrestricted weight : Nat . (lambda unrestricted input : (family NvidiaLinearScalarInput) . (lambda unrestricted destination : Nat . (nlsEvaluateScalar b"coordinate-device-arithmetic" gpu entry data inputOffset targetOffset outputOffset (cdCounter end) (cdScratch end) (constructor NvidiaLinearScalarInput NvidiaLinearScalarBinary64 (nativeMemoryLoad weight)) input (constructor NvidiaLinearScalarInput NvidiaLinearScalarNatural 0) cdProductResult destination))))))))))) def cdTrainingAdmitted = (lambda unrestricted worker : Nat . (lambda unrestricted resume : Nat . (lambda unrestricted epoch : Nat . (lambda unrestricted steps : Nat . (naturalAnd (naturalLessOrEqual worker 1) (naturalAnd (naturalLessOrEqual resume 1) (naturalAnd (naturalNonzero epoch) (naturalAnd (naturalLessOrEqual (naturalAdd epoch resume) regionFrameWordLimit) (naturalAnd (naturalNonzero steps) (naturalLessOrEqual (naturalMultiply cdmIssueCount steps) regionFrameWordLimit)))))))))) -- Only this initial transfer writes parameter/accumulator from host state. -- A restored checkpoint follows the same path before the first device issue. def cdSeedMemory = (lambda unrestricted layout : (family CoordinateHostLayout) . (lambda unrestricted data : Nat . (lambda unrestricted tail : (family NativePhysicalCommands) . (phCopyStateToMapped (chWeight layout) (naturalAdd data cdmParameter) coordinateWordBytes (phCopyStateToMapped (chAccumulator layout) (naturalAdd data cdmAccumulator) coordinateWordBytes (phCopyStateToMapped (chFeature layout) (naturalAdd data cdmFeature) coordinateWordBytes (phCopyStateToMapped (chDecay layout) (naturalAdd data cdmDecay) coordinateWordBytes (phCopyStateToMapped (chRate layout) (naturalAdd data cdmRate) coordinateWordBytes tail)))))))) -- Host values are validation/checkpoint shadows, never input to a later -- parameter update. Capture follows the device fence and precedes publication. def cdCaptureMemory = (lambda unrestricted layout : (family CoordinateHostLayout) . (lambda unrestricted data : Nat . (lambda unrestricted tail : (family NativePhysicalCommands) . (phCopyMappedToState (naturalAdd data cdmParameter) (chWeight layout) coordinateWordBytes (phCopyMappedToState (naturalAdd data cdmAccumulator) (chAccumulator layout) coordinateWordBytes tail))))) def cdMemoryPhase = (lambda unrestricted layout : (family CoordinateHostLayout) . (lambda unrestricted gpu : (family NvidiaPlanHostLayout) . (lambda unrestricted component : (pi unrestricted kind : (family NvidiaWholeProgramComponent) . Bytes) . (lambda unrestricted data : Nat . (lambda unrestricted phase : (family CoordinateDevicePhase) . (lambda unrestricted tail : (family NativePhysicalCommands) . (let unrestricted issue = (nvidiaProgramIssue gpu component (cdCounter (chExtent layout)) (cdmOrdinal phase)) in (eliminate CoordinateDevicePhase (lambda unrestricted current : (family CoordinateDevicePhase) . (family NativePhysicalCommands)) phase (branch CoordinateDeviceForward . (issue (phCopyMappedToState (naturalAdd data cdmPartial) (chPartial layout) coordinateWordBytes tail))) (branch CoordinateDeviceBackward . (phCopyStateToMapped (chError layout) (naturalAdd data cdmResidual) coordinateWordBytes (issue (phCopyMappedToState (naturalAdd data cdmGradient) (chGradient layout) coordinateWordBytes tail)))) (branch CoordinateDeviceUpdate . (issue (cdCaptureMemory layout data tail))))))))))) -- argv is the derived manifest/shard layout followed by an exclusive device -- readback path. The session uploads state once, then the three submissions -- share video memory. Final device capture/readback and release precede the -- global checkpoint publication protocol; no success bypasses either lifecycle. def coordinateDeviceArtifact = (lambda unrestricted resources : (family NvidiaResourceContext) . (lambda unrestricted profile : (family RunPodSM86TargetProfile) . (lambda unrestricted device : (family NvidiaDeviceBinding) . (lambda unrestricted worker : Nat . (lambda unrestricted address : (family IPv4Address) . (lambda unrestricted port : Nat . (lambda unrestricted resume : Nat . (lambda unrestricted epoch : Nat . (lambda unrestricted steps : Nat . (eliminate NvidiaProgramPlan (lambda unrestricted plan : (family NvidiaProgramPlan) . Bytes) (cdmPlan profile) (branch NvidiaProgramPlanValue gpu component recipes data dataGPU dataBytes admitted . (nat-eliminate (lambda unrestricted valid : Nat . Bytes) b"" (lambda unrestricted previous : Nat . (lambda unrestricted ignored : Bytes . (chWithLayout cdNumerics epoch phStateExtent (lambda unrestricted layout : (family CoordinateHostLayout) . (eliminate StdResult (lambda unrestricted result : (family StdResult (family SocketAddressFailure) Bytes) . Bytes) (linuxIPv4SocketAddress address port) (branch StdFailure error . b"") (branch StdSuccess socketAddress . (nat-eliminate (lambda unrestricted valid : Nat . Bytes) b"" (lambda unrestricted previous : Nat . (lambda unrestricted ignored : Bytes . (nlsWithProgramRecipes component (lambda unrestricted program : Bytes . (lambda unrestricted qmd : Bytes . (lambda unrestricted push : Bytes . (lambda unrestricted ring : Bytes . (nlsWithPrivateSlots nativePeerEndpointSlots (nlsSessionPrelude cdIdentity b"coordinate-peer" device gpu recipes program qmd push ring (cdSeedMemory layout data chEnd)) (lambda unrestricted prepare : (family NativePhysicalCommands) . (nlsWithPrivateSlots nativePeerEndpointSlots (cdMemoryPhase layout gpu component data (constructor CoordinateDevicePhase CoordinateDeviceForward) chEnd) (lambda unrestricted forward : (family NativePhysicalCommands) . (nlsWithPrivateSlots nativePeerEndpointSlots (cdMemoryPhase layout gpu component data (constructor CoordinateDevicePhase CoordinateDeviceBackward) chEnd) (lambda unrestricted backward : (family NativePhysicalCommands) . (nlsWithPrivateSlots nativePeerEndpointSlots (cdMemoryPhase layout gpu component data (constructor CoordinateDevicePhase CoordinateDeviceUpdate) chEnd) (lambda unrestricted update : (family NativePhysicalCommands) . (nlsWithPrivateSlots nativePeerEndpointSlots (cdCaptureMemory layout data (nvidiaProgramReadback (chArgument (succ chDirectoryIndex)) (prSemaphoreHost gpu) data dataBytes cdmIssueCount (cdCounter (chExtent layout)) (phTelemetry 3 b"coordinate-peer:device-updates-retired" (phRelease chEnd)))) (lambda unrestricted finish : (family NativePhysicalCommands) . (nlsEmbedHost cdIdentity (nlsGenerateResourceELF resources gpu (chProgramFor (naturalAdd nativePeerEndpointSlots phSlots) (cdTrainingExtent (chExtent layout)) (chCommandsWithSession layout worker socketAddress resume steps (prSplice forward) (prSplice backward) (prSplice update) (lambda unrestricted tail : (family NativePhysicalCommands) . (chStore (cdCounter (chExtent layout)) (phImm 0) (prSplice prepare tail))) (prSplice finish)))) program qmd push ring)))))))))))))))))) (arenaCertificate (cdTrainingExtent (chExtent layout)) (cdTrainingResidents (chExtent layout) (chResidents layout)))))))))) (naturalAnd admitted (naturalAnd (cdBindingAdmitted device) (cdTrainingAdmitted worker resume epoch steps))))))))))))))) -- Prediction uses the same numerical policy on a fresh bound device session. -- It restores both saved parameters, executes both products on that GPU, then -- combines the widened results. The extra argv follows the five input paths. def coordinateDevicePredictionArtifact = (lambda unrestricted resources : (family NvidiaResourceContext) . (lambda unrestricted profile : (family RunPodSM86TargetProfile) . (lambda unrestricted device : (family NvidiaDeviceBinding) . (lambda unrestricted epoch : Nat . (eliminate NvidiaLinearStepPlan (lambda unrestricted plan : (family NvidiaLinearStepPlan) . Bytes) (regressionScalarPlan profile) (branch NvidiaLinearStepPlanValue gpu component recipes data inputOffset targetOffset outputOffset dataBytes admitted . (nat-eliminate (lambda unrestricted valid : Nat . Bytes) b"" (lambda unrestricted previous : Nat . (lambda unrestricted ignored : Bytes . (cpWithLayout cdNumerics epoch phStateExtent (lambda unrestricted layout : (family CoordinatePredictionLayout) . (nat-eliminate (lambda unrestricted valid : Nat . Bytes) b"" (lambda unrestricted previous : Nat . (lambda unrestricted ignored : Bytes . (nlsWithProgramRecipes component (lambda unrestricted program : Bytes . (lambda unrestricted qmd : Bytes . (lambda unrestricted push : Bytes . (lambda unrestricted ring : Bytes . (let unrestricted entry = (nvidiaPlanHostTableWord (component (constructor NvidiaWholeProgramComponent NvidiaWholeProgramGPFIFO)) 0) in (nlsWithPrivateSlots cpSlots (nlsSessionPrelude cdIdentity b"coordinate-prediction" device gpu recipes program qmd push ring chEnd) (lambda unrestricted prepare : (family NativePhysicalCommands) . (nlsWithPrivateSlots cpSlots (cdProduct gpu entry data inputOffset targetOffset outputOffset (cpExtent layout) (cpLeft layout) (constructor NvidiaLinearScalarInput NvidiaLinearScalarNatural 1) (cpPrediction layout) (cdProduct gpu entry data inputOffset targetOffset outputOffset (cpExtent layout) (cpRight layout) (constructor NvidiaLinearScalarInput NvidiaLinearScalarNatural 2) (cpFeature layout) chEnd)) (lambda unrestricted products : (family NativePhysicalCommands) . (nlsWithPrivateSlots cpSlots (nlsScalarReadback (chArgument chOutputShardBase) gpu data outputOffset (cdCounter (cpExtent layout)) (phTelemetry 3 b"coordinate-prediction:device-products-retired" (phRelease chEnd))) (lambda unrestricted finish : (family NativePhysicalCommands) . (cpWithRestore layout (chStore (cdCounter (cpExtent layout)) (phImm 0) (prSplice prepare (prSplice products (chBinary chAdd (cpPrediction layout) (cpPrediction layout) (cpFeature layout) (prSplice finish (cpOutput layout chHalt)))))) (lambda unrestricted commands : (family NativePhysicalCommands) . (nlsEmbedHost cdIdentity (nlsGenerateResourceELF resources gpu (chProgramFor (naturalAdd cpSlots phSlots) (cdExtent (cpExtent layout)) commands)) program qmd push ring))))))))))))))))) (arenaCertificate (cdExtent (cpExtent layout)) (cdResidents (cpExtent layout) (cdPredictionResidents layout)))))))) (naturalAnd admitted (cdBindingAdmitted device)))))))))