google-research / dex-lang

Research language for array processing in the Haskell/ML family
BSD 3-Clause "New" or "Revised" License
1.58k stars 107 forks source link

Unallocated memory read with nested dependent pairs #1212

Closed axch closed 1 year ago

axch commented 1 year ago

Reproduction

data Rectangle a = AsRectangle n:Nat m:Nat elts:(Fin n => Fin m => a)

x = AsRectangle _ _ $ for i:(Fin 1). for j:(Fin 1). ordinal j

x
> (AsRectangle 1 1 [[66]])

Expected (AsRectangle 1 1 [[0]])

Other variants of this produce other numbers in the rectangle, can be made non-deterministic, and otherwise indicate reading unallocated memory.

The problem does not occur with a List, or a List of dependent pairs, making me think that the bug is due to the implementation of Place in ImpToLLVM on nested dependent pairs (such as Rectangle becomes).

axch commented 1 year ago

Based on investigation yesterday, it also seems as though this causes AD to fail on tiled_matmul. The current theory is that tiling causes AD to use nested dependent pairs to store the tape, and then this bug causes the tape to be mis-read. In the particular case of AD of tiled matmul, what's misread from the tape seems to be the indices at which to read and write the matrices being multiplied in the tangent pass.