Move ERI computation to the IPU, keeping sparse symmetric einsum functionality #111

mihaipgc · 2023-09-27T17:27:21Z

This PR aims to move ERI computation to the IPU. The code started from nanoDFT/sparse_symmetric_ERI.py and aims to reach a similar form once it is complete.

Overview:

sparse_symmetric_intor_ERI.py contains compute_eri which is a reduced version of the relevant integral computing code in electron_repulsion/direct.py
compute_eri_utils.py contains minimal modifications to code from electron_repulsion/direct.py to add prescreening

Related issues:

Summary of changes:

correctly computes integral values
skips all-zero integral subtensors (experimental prescreening)
supports batching
maps integral computation to tiles
supports multi-ipu integral computation / pmap
does not store indices (generates them)

…eparately; work in progress

…mber of values/zeros

…ntly leads to segfault due to elementwise checks (work in progress)

…sors directly; still computing indices inefficiently (work in progress)

…n profiles being created

…shold

…cessary code slowing program down; added integral value tolerance as program argument

…ogress

…ing to estimated final shape; implemented batch mode for ortho; work in progress

…e code; restored complete integral ipu code

…hells

…ing (shell local index recomputation); still using int32, work in progress

…or 6-311G.

…k in progress

Merge updates from main

AlexanderMath

lgtm

mihaipgc added 14 commits September 18, 2023 16:02

added initial code with test against dense eri passing

b39a6e6

fixed index transform bug and tile mapping, now passes for natm>1

5318e1c

able to compute distinct_ERI and compute diff_JK using distinct_ERI s…

75bbbbd

…eparately; work in progress

index and eri value repetition computations moved to ipu

04a54bf

moved distinct_ERI computation to ipu; slow for now

87699d7

constructed distinct_ERI list by explicitly skipping duplicates

6fbb582

block concatenation tricks to obtain faster compilation; increased nu…

a0cf5fb

…mber of values/zeros

cleanup and removed a transpose; passes for natm>1 since previous commit

4f2f1d8

fixed bug in eri value bookkeeping when natm>2 in batched mode; curre…

72ea97b

…ntly leads to segfault due to elementwise checks (work in progress)

drastically reduced compile time by properly concatenating eri subten…

b21bbeb

…sors directly; still computing indices inefficiently (work in progress)

improved index generation; cleaned up code; avoided multiple popvisio…

29fc367

…n profiles being created

added batching

7d2e52b

added nonzero index generation; work in progress

a0710d0

improved screening procedure and added safety check

0f53dfd

mihaipgc self-assigned this Sep 27, 2023

mihaipgc and others added 15 commits September 27, 2023 18:00

add copyright

0e8c85c

added working version of fast nonzero index computation

05ced9c

cleaned code; improved index generation loops; added flag to use thre…

209eb33

…shold

changed to ERI s8 N^2 value sampling for faster testing; removed unne…

f628757

…cessary code slowing program down; added integral value tolerance as program argument

cleaned up code; removed no longer necessary tests

0c61c66

refactored code ahead of multi-ipu eri computation

7977064

refactored and added padding and bookkeeping for eri pmap; work in pr…

9681f73

…ogress

isolated prescreening and orthogonal zero removal with dynamic gather…

f0102cd

…ing to estimated final shape; implemented batch mode for ortho; work in progress

Simplifying integral code.

16514f6

more progress on cleaning.

6fb7df0

Added comments

441d206

removed do-block bookkeeping, now only depends on ijkl; condensed mor…

2a61d2d

…e code; restored complete integral ipu code

removed redundant copy from vertex; reused same vertex code for all s…

390218f

…hells

improved do-tensor rule, to be ported to vertex later

e30a6ef

moved do-block computation from python to poplar

5ffa784

mihaipgc and others added 14 commits October 10, 2023 16:02

moved index computation closer to matmul

0f017bf

refactoring and cleanup; index recomputation work in progress

9537f4d

working shell-based indexing with lower memory consumption with batch…

3d2cc20

…ing (shell local index recomputation); still using int32, work in progress

code cleanup; includes full shell index version; still int32

2f5be9c

cleaned code

d3d716a

using uint8 for indices to reduce memory footprint

df28081

compresssed codelet mapping code; cleanup

0e3bcf6

further compressed code

e686450

Got smallest possible triple valence test case running, two carbons f…

752ef66

…or 6-311G.

reusing tile memory between matmul symmetries; compressed more code

43cd3cd

preparing tile-mapped index generation and tile mapping matching; wor…

c81b013

…k in progress

reverted to encapsulated symmetry indexing

96cd133

update

939ec97

Merge pull request #131 from graphcore-research/main

35fd5b0

Merge updates from main

mihaipgc marked this pull request as ready for review October 26, 2023 09:10

AlexanderMath approved these changes Oct 26, 2023

View reviewed changes

mihaipgc merged commit cd2d875 into main Oct 26, 2023
4 checks passed

mihaipgc deleted the combine-sparse-symmetric-with-eri-computation branch October 26, 2023 09:18

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Move ERI computation to the IPU, keeping sparse symmetric einsum functionality #111

Move ERI computation to the IPU, keeping sparse symmetric einsum functionality #111

mihaipgc commented Sep 27, 2023 •

edited

Loading

AlexanderMath left a comment

Move ERI computation to the IPU, keeping sparse symmetric einsum functionality #111

Move ERI computation to the IPU, keeping sparse symmetric einsum functionality #111

Conversation

mihaipgc commented Sep 27, 2023 • edited Loading

AlexanderMath left a comment

Choose a reason for hiding this comment

mihaipgc commented Sep 27, 2023 •

edited

Loading