https://akio-tomiya.github.io/julia_imi_workshop2023/
https://joint1.imi.kyushu-u.ac.jp/research_chooses/view/2023a015
7/12(水)数学と物理におけるJuliaの活用

素粒子物理におけるJuliaの活用〜格子QCDの大規模計算に向けて〜
Akio Tomiya (IPUT Osaka, Assistant Prof.)
akio_at_yukawa.kyoto-u.ac.jp

IMI in person

Outline• 素粒子物理学? 格子QCD?
• Julia言語での取り組み• 何ができれば良いのか2

How to treat gauge eldswith neural networks?

fi

3

Applications of neural network in natural science
Neural network have been good job

Akio Tomiya

Folding of protein (AlphaFold2, John Jumper+, Nature, 2020+)
Score:
Higher is better

Neural network wave function for many body (Carleo Troyer, Science 355, 602 (2017) )Variational energy(lower is better)# of units ∝Neural net + Expert knowledge → Best performancehttps://horomary.hatenablog.com/entry/2021/10/01/194825
https://www.aist.go.jp/aist_j/press_release/pr2020/pr20201111/pr20201111.html

4What is neural networks?

Akio Tomiya

A ne transformation + element-wise transformation
Layers of neural nets l = 2,3,⋯, L, u (1)⃗ = x⃗(l−1)⃗{ u (l) = σ (l)(z (l))
(l)

(l)

z⃗ = W u
i

+ b

(l)
⃗

i

A ne transf.
(b=0 called linear transf.)

element-wise (local)

A fully connected neural net⃗ ) + b (3)⃗ )fθ( x ⃗) = σ (3)(W (3)σ (2)(W (2) x ⃗ + b (2)θ is a set of parameters: wij(l), bi(l), ⋯Neural network = map between vectors and vectors
Physicists terminology: Variational transformations

ffi

ffi

5

What is the neural networks?

Akio Tomiya

Neural network is a universal approximator
Example: Recognition of hand-written numbers
6x6

36 dimension

x ⃗ Input

0.000
0.000
0.84340.756=
0.3456
0.64
0.251
⋮

= x⃗

Neural net

Image is a vector(6x6=36 dim)1
10 dimension

Images of “1”

W

(2)

W

(3)

2

“0” = (1,0,0,…)“1” = (0,1,0,…)“2” = (0,0,1,…)…
“9” = (0,0,…,1)

Images of “2”Fact: Neural network can mimic any function= A systematic variational function.
In this example, NN mimics image (36-dim vector) and label (10-dim vector)
6

ML for LQCD is needed• Neural networks• Data processing techniques mainly for 2d
image (a picture = pixels = a set of real #)• Neural network helps data processing
e.g. AlphaFold2• Lattice QCD requires numerical e ort

but is more complicated than pictures• 4 dimension• Non-abelian gauge d.o.f. and symmetry• Fermions (Fermi-Dirac statistics)• Exactness of algorithm is necessary• Q. How can we deal with neural nets?

http://www.physics.adelaide.edu.au/theory/sta /leinweber/VisualQCD/QCDvacuum/

ff

ff

7

Gauge covariant neural network= trainable smearing

Akio Tomiya

AT Y. Nagai arXiv: 2103.11965

Smearing = Gauge covariant way of transform gauge con gurations
Covariant sum

Uμ(n) → Uμsmr(n) =Stapleα †̂ ν†(n + μ)̂ + ⋯(1 − α)Uμ(n) + Vμ [U](n) Vμ†[U ](n) = ∑ Uν(n)Uμ(n + ν)U
[
]
6
μ≠ν

[M] =

M

M †M

Normalization or projection

Gauge covariant neural network = General smearing with tunable parameters w(l)(l) (l−1)(l)zμ (n) = w1 Uμ (n) + w2

{ (zμ(l)(n))

(l)
[U]
θ̄

Train (tune, tting)point-wise (local)

Gauge covariant NN: UμNN(n)[U] = Uμ(4)(n)[Uμ(3)(n)[Uμ(2)(n)[Uμ(n)]]]NN(n) = UμNN(n)[U]

𝒢

𝒩

𝒩

𝒩

Gauge covariant variational map: Uμ(n) ↦ Uμfi

fi

Alternative realization of gauge symmetric neural net: gauge equivariant neural net -> Tilo Wettig’s talk8

Gauge covariant neural network

Akio Tomiya

Schematic illustrations for neural networks (NN)Neural networks for images

Parametrized function

Cat

Ans
Dog

Dense NN
Parameters θ

Convolution
Parameters θTune by backprop (training)

Neural networks for gauge con gurations
Parametrized function
AT Nagai 2103.11965

Wilson loop

U

U NN[U]

Covariant NN(Tunable smearing)Parameters θSD[UDirac op.
(Functionalof con gurations)

NN

[U]]

“Ans”

SD[U]

Tune by backprop (training)
http://www.physics.adelaide.edu.au/theory/sta /leinweber/VisualQCD/QCDvacuum/

fi

ff

fi

Alternative realization of gauge symmetric neural net: gauge equivariant neural net -> Tilo Wettig’s talk9

Gauge covariant neural network= trainable smearing

Akio Tomiya

AT Y. Nagai arXiv: 2103.11965(convolutional)
Neural network

Gauge Covariant
Neural network

Input

Image
(2d data, structured)gauge con g(4d data, structured)

OutputImage(2d data, structured)gauge con g(4d data, structured)

Dictionary

Symmetry

Translation

with Fixed paramImage lter

Translation, rotation(90°),
Gauge sym.
Gauge sym
(APE/stout …) Smearing

Local operation

Summing up nearest
neighbor with weights

Summing up staples
with weights

Activation function

projection/normalization
Tanh, ReLU, sigmoid, …
in Stout/HYP/HISQ

Formula for chain rule

Backprop“Smeared force
calculations” (Stout)

Training?

Backprop + Delta rule

AT Nagai 2103.11965

Well-known(Index i in the neural net corresponds to n & μ in smearing. Information processing with NN is evolution of scalar eld)

fi

fi

fi

fi

10

SLHMC = Exact algorithm with ML

Akio Tomiya

SLHMC for gauge system with dynamical fermionsGauge covariant neural network can mimics gauge invariant functions-> It can be used in simulation? -> Self learning HMC!

HMC

U

U

U

U

U

U

￼

￼

￼

￼

arXiv: 2103.11965
and reference therein

ϕ
U

ϕ

U

Eom
Eom

Eom

U
π

U

U

U′
Eom

U

Metropolis

π′

π′
ϕ

Metropolis

Eom

Eom

U′

ϕ

Eom

U

Gauss

Self
Learning
HMC

π

Eom

Gauss

U

Eom

Eom

Eom

Metropolis

Both use

1
HHMC =
π 2 + Sg + Sf
2∑

Non-conservation of H cancels since
the molecular dynamics is reversible

Metropolis1H=
π 2 + Sg + Sf[U]
2∑
Eom

1
H=
π 2 + Sg+Sf[U NN[U]]2∑
Neural net approximatedfermion action but exact

SLHMC works as an adaptive reweighting!

11

Application for the staggered in 4d

Akio Tomiya

Problems to solve
arXiv: 2103.11965

Mimic di erent action (Dirac operator):(Final target: Domain-wall vs overlap)
A toy problem: Staggered (heavy) vs Staggered (light){Action in MD

S[U] = Sg[U] + Sf[ϕ, U; m = 0.3],

Target action(Metropolis)

Eom

Eom

U

U

U′
π′
ϕ

Metropolis

ϕ

Eom

U
π

U

Metropolis1H=
π 2 + Sg + Sf[U]
2∑
Eom

1
H=
π 2 + Sg+Sf[U NN[U]]2∑
Neural net approximatedfermion action but exact

SLHMC works as an adaptive reweighting!
ff

￼

￼

U

Eom

U

Eom

U

Gauss

Self
Learning
HMC

Sθ[U] = Sg[U] + Sf[ϕ, UθNN[U]; mh = 0.4],

mimic

12

Application for the staggered in 4d

Akio Tomiya

Results are consistent with each otherarXiv: 2103.11965

Expectation value

Implemented by
13

Julia language on Fugaku,
Lattice code in Julia(LatticeQCD.jl and GomalizingFlow.jl)

14

Benchmark

Akio Tomiya

Speed of Julia ~ Clang
1.Open source scienti c language (Just in time compiler/LLVM backend)
2.Fast as C/Fortran (faster sometimes), Practical as Python
3.Machine learning friendly

Compiler

Benchmark(sec)
single coreType

Parallelism

GPU

Pros👍

Cons👎

Column,row

MPI,othersCUDA

FastPractical
ML feiendly

not major

columnmajor

MPI,
others

CUDA

Fast

Long
codes

row-major

Julia
(1.8)

JIT,
LLVM

0.0014

Dynamic& StaticC

Clang(LLVM)

0.0033

StaticPython+Numba(CPython)
JIT, LLVM

Not fully row-major
Practical
0.0131 Dynamic Available Numba- ML
feiendly supported (Numpy)
CUDA

Note

(Rosetta2
is used inbenchmark)C and Julia have similar speed
Benchmarks are performed on m1 mac mini (similar tendency on Xeon)
Benchmark: Multiplications for 12dim vector and 12x12 complex matrix for 10^4 times (repeated 10 times)
fi

15

15

Benchmark

Akio Tomiya

Code comparison
using Randomfunction main()
T = 10
K = 10^4
N = 12#A = zeros(ComplexF64, (N,N))
V = zeros(ComplexF64, N)
W = zeros(ComplexF64, N)
function myprod(A,V,W)
for k = 1:N
for i = 1:N
W[i] += A[i, k]*V[k]
end
end
end
…(cut)…
Attached in backup#include <stdio.h>#include <complex.h>#include <math.h>#include <time.h>#include <stdlib.h>#define T 10
#define K 10000#define N 12…(cut)…
void myprod(double complex A[N][N], double complex *V,
double complex *W) {for (int k = 0; k < N; k++) {for (int i = 0; i < N; i++) {W[i] += V[k] * A[k][i];}}}Attached in backup…(cut)…

• Complex matrix (12x12) times complex vector (d=12)• One set= 10^4 times, and repeated 10 times and averaged• Code of Julia looks like Python (short, simple) but fast as C

Julia: 0.0014 (sec), C: 0.0033 (sec). Single core performance is similar

I thank Taku Izubuchi

1616Benchmark

Akio Tomiya

Why Julia? (My personal opinion)• Modern scienti c programming language[1] https://akio-tomiya.github.io/julia_in_physics/[2] https://qr.ae/prgSG5• Easy to make codes. Fast as C/C++ (Julia& C use LLVM)• Fewer compiling/dependency issues.
• Many people are potentially interested in. (More than 400 people registered to“Julia in physics 2022 online workshop” [1]). 4,923 public repo on Github• No two Language problem. “The fact that while the users are programming ina high-level language such as R and Python, the performance-critical parts
have to be rewritten in C/C++ for performance”. [2]• Neural network friendly (Flux.jl). Tensor networks also (iTensor.jl).
• Works on/with• Xeon, Radeon/Apple silicon/A64FX• MPI, GPU
fi

17

17

LLVM?

Akio Tomiya

LLVM = common backend for making binaries on multiple architectures
SourceExecution
Frontend

const K = 5
const N = 1000
A = zeros(Float64, (N,N))
B = zeros(Float64, (N,N))
C = zeros(Float64, (N,N))

Julia Compiler(Just in time)
Optimizing codes

LLVM-optimizer&
LLVM-backend
LLVM-IR

Frontend

Binary

LLVM#include <stdio.h>#define K 5#define N 1000
double A[N][N];double B[N][N];double C[N][N];x86“Clang” Compiler(Ahead of Time)
Optimizing codes

https://www.fujitsu.com/jp/about/businesspolicy/tech/fugaku/
https://ja.wikipedia.org/wiki/Apple_M1
https://ja.wikipedia.org/wiki/Xeon
https://ja.wikipedia.org/wiki/Ryzen
https://gigazine.net/news/20200623-japan-fugaku-fastest-supercomputer/

m1, m2arm

used in Fugaku
See: https://en.wikipedia.org/wiki/LLVM and related pages

18

Julia is ready on Fugaku(?)
Parallelization with A64FX/Fugaku
M. Giordano, arXiv:2207.12762v1 [cs.DC] 26 Jul 2022

Tests of MPI +on Fugaku

Send-Recv performanceFig. 2. Comparison of latency (top panel) and throughput (bottom panel)of inter-node point-to-point MPI communication between using MPI.jl in
Julia and IMB benchmarks in C (results provided by R-CCS in [18]). Fugaku
scheduler setup: -L "node=2" -mpi "max-proc-per-node=1".

has similar scaling of MPI with C(no obvious overhead)

MPI ranks across 384 nodes using the torus layout, to match the19Lattice QCD code
Open source LQCD code in Julia Language

Akio Tomiya

AT & Y. Nagai in prep

Open source (Julia O cial package, Now updated to v1.0)

Machines: Laptop/desktop/Jupyter/Supercomputers
Functions: SU(Nc)-heatbath, (R)HMC, Self-learning HMC, SU(Nc) Stout
Dynamical Staggered, Dynamical Wilson, Dynamical Domain-wall
Measurements
1. Download Julia binary
2. Add the package through Julia package manager
3. Execute!

Start LQCD
in 5 minhttps://github.com/akio-tomiya/LatticeQCD.jl

ffi

20

Package structure

Akio Tomiya

Our lattice QCD codes are constructed by following repositories
Dependency (Automatically solved)

QCDMeasurements.jl
LatticeDiracOperators.jl
Gauge elds.jl
Wilsonloop.jl

Wrapper for LatticeDiracOperators.jl &Gauge elds.jl, QCDMeasurements.jl- Wizard for parameter les- HMC/RHMC for SU(Nc)- Stout + Wilson/Staggered/DW- Heatbath for SU(Nc)- Measurements- Jupyter, Colab/PC/Supercomputers
etc
Measurements in LQCD
(Correlator, Flow, Qtop, etc)

CLIME_jll
ILDG I/O

Fermions (+HMC), Wilson, KS, DW, MPI
PC/Supercomputers
Gauge elds (+HMC/Heatbath), MPI
PC/Supercomputers

Symbolic operations of Wilson/Polyakov loops
See https://github.com/akio-tomiya/LatticeQCD.jl in detail
fi

fi

fi

fi

21

Benchmark of Julia + QCD

Akio Tomiya

Wilson inversion / MPI parallel, Strong Scaling

AT & Y. Nagai in prep
AT & Y. Nagai in prep

Tested on Yukawa-21@YITP
Absolute execution time

Relative speed up

LatticeDiracOperators.jl# procs

1

2

4

16

32

LatticeDiracOperators.jl

64# procs

1

2

4

16

32

64

It looks scaling well
We need more contributors!
Please help us
We thank to H. Ohno & Issaku Kanamori

22

22
Akio Tomiya

Flow based sampling algorithm

Akio Tomiya

Trivializing map realized using neural network
Normalizing ow? = Trivializing map, exact MCMC with ML

Change of variable by a neural network (Normalizing ow)

ϕ=F

(NN)

[φ]

Sampling from Gaussian→ Inverse trivializing map (neural net)
→ QFT con gurations → Tractable Jacobian (by even-odd strategy)→ After sampling, Metropolis-Hastings test → exact!
Related talk: Nobuyuki Matsumoto (Feb 16)

arxiv 1904.12072, 2003.06413, 2008.05456

fl

fl

fi

23

Normalizing ow in Julia

Akio Tomiya

A public code in Julia Language

https://arxiv.org/abs/2208.08903

A public Julia code for the ow-based
sampling algorithm for scalar eld.
This supports not only 2d but also 3d.
CPU/GPU with Docker.
https://github.com/AtelierArith/GomalizingFlow.jl

A new type of convolution
improves acceptance rate(~ shorten the autocorrelation)
I reported in NeurIPS 2022 workshop
https://ml4physicalsciences.github.io/2022/

fi

fl

fl

24

Summary

Akio Tomiya

LQCD + Machine learning by Julia language

arXiv: 2103.11965
arXiv: 2208.08903• Machine learning for LQCD
• Neural net (NN) + expert knowledge -> Best performancee.g. AlphaFold2, NN wave functions• NN can deal with 4d non-abelian gauge symmetric data now• Self-learning HMC with NN works for dynamical fermions• Julia language for LQCD/HPC/ML• Julia has similar speed with C (w/ & w/o MPI), and machine learning friendly• Two Julia codes for lattice eld theory• LatticeQCD.jl: A suite lattice QCD code, machine learning• GomalizingFlow.jl: Trivializing map via a neural networkfi

25
KAKENHI: 20K14479, 22H05112,
22H05111, 22K03539

Thanks!

25

26

MPI benchmark

Akio Tomiya

PingPong

M. Giordano, arXiv:2207.12762v1 [cs.DC] 26 Jul 2022

This section describes the performance of PingPong as a oneto-one communication. PingPong sends data between two
ranks by Send communication from one rank, receives the dataat the other rank, and then sends the data back to the original
rank by Recv communication.

Fig. 2. Comparison of latency (top panel) and throughput (bottom panel)of inter-node point-to-point MPI communication between using MPI.jl in
Julia and IMB benchmarks in C (results provided by R-CCS in [18]). Fugaku
scheduler setup: -L "node=2" -mpi "max-proc-per-node=1".

Allreduce collects data from all ranks, performs a set
operation, and transfers the result to and transfers
the result to all ranks.

https://jopss.jaea.go.jp/pdfdata/JAERI-Data-Code-2001-010.pdf
27

MPI ranks across 384 nodes using the torus layout, to match thescheduler configuration of the R-CCS benchmarks. MPI.jl
typically showed very small overhead for messages larger than
1-2 KiB—peak throughput of ping-pong communication with
MPI.jl is within 1% of that reported by R-CCS—, but slightly
larger overhead for messages of smaller sizes. We note that,
contrary to IMB, at the present time MPIBenchmarks.jl
does not implement a cache-avoidance mechanism, which may
explain why MPI.jl appears to show better latency than IMB
for messages with size up to 64 KiB, which corresponds tothe size of the L1 cache of the A64FX CPU. We also observe
that, contrary to [16], we did not find a significant performance27

Introduction

Akio Tomiya

Con guration generation with machine learning is developing
Year

Group

ML

Dim.

TheoryGauge sym

Exact?

Fermion?

Reference

2017

AT, Akinori
Tanaka

RBM
+ HMC

2d

Scalar-No

No

arXiv: 1712.03893

2018 K. Zhou+

GAN

2d

Scalar-No

No

arXiv: 1810.12879

2018

J. Pawlowski +

GAN
+HMC

2d

Scalar

-

Yes?

No

arXiv: 1811.03533

2019

MIT+

2d

Scalar-Yes

No

arXiv: 1904.12072

2020

MIT+

2d

U(1)

Equivariant

Yes

No

arXiv: 2003.06413

2020

MIT+

Flow
Flow
Flow

2d

SU(N)

Equivariant

Yes

No

arXiv: 2008.05456

2020

AT, Akinori
Tanaka +

SLMC

4d

SU(N)

InvariantYes

PartiallyarXiv: 2010.11900

2021

M. Medvidovic´+

A-NICE

2d

Scalar-No

No

arXiv: 2012.01442

2021 S. Foreman L2HMC

2d

U(1)

Yes

Yes

No

2021

AT+

SLHMC

4d

QCD

Covariant

Yes

YES!

2021

L. Del
Debbio+2dScalar, O(N)

-

Yes

No

2021

MIT+

Flow
Flow

2d

Yukawa

-

Yes

Yes

2021

S. Foreman,AT+

Flowed
HMC

2d

U(1)

Equivariant

Yes

No but compatible

2021

XY Jing

2d

U(1)Equivariant

Yes

No

2022

J. Finkenrath

Neural
net
Flow

2d

U(1)

Equivariant

Yes

Yes (diagonalization)

arxiv: 2201.02216

2022

MIT+

Flow

2d, 4dU(1), QCD

Equivariant

Yes

Yes

arXiv:2202.11712 +

2022

AT+

Flow

2d, 3dScalarThis talkarXiv: 2112.01586

Yrs
＋…

fi

28

Benchmark

Akio Tomiya

Code comparison
using Randomfunction main()
T = 10
K = 10^4
N = 12#A = zeros(ComplexF64, (N,N))
V = zeros(ComplexF64, N)
W = zeros(ComplexF64, N)
function myprod(A,V,W)
for k = 1:N
for i = 1:N
W[i] += A[i, k]*V[k]
end
end
end
function test(A,V,W)
for jj=1:T
runtimes=[]
for r=1:K
A .= rand(N,N) + im*rand(N,N)
V .= rand(N) + im*rand(N)W .= 0tmp = @elapsed myprod(A,V,W)
push!(runtimes,tmp)
endprintln("$(jj-1) $(sum(runtimes)) #W[1] = $(W[1])")end
end
test(A,V,W)
end
if abspath(PROGRAM_FILE) == @__FILE__
main()
end#include <stdio.h>#include <complex.h>#include <math.h>#include <time.h>#include <stdlib.h>#define T 10
#define K 10000#define N 12
double urand(){double m, a;
m = RAND_MAX + 1;
a = (rand() + 0.5)/m;
a = (rand() + a)/m;
return (rand() + a)/m;}
void myprod(double complex A[N][N], double complex *V, double complex *W) {for (int k = 0; k < N; k++) {for (int i = 0; i < N; i++) {W[i] += V[k] * A[k][i];}}}void test(double complex A[N][N], double complex *V, double complex *W) {for (int jj = 0; jj < T; jj++) {double runtimes = 0;
for (int r = 0; r < K; r++) {for (int i = 0; i < N; i++) {for (int j = 0; j < N; j++) {A[i][j] = urand() + urand() * I;}
V[i] = urand() + urand() * I;
W[i] = 0.0 + 0.0 * I;}clock_t start = clock();
myprod(A, V, W);
clock_t end = clock();
runtimes += (double)(end - start) / CLOCKS_PER_SEC;}
printf("%d %f # W[0] = %f %f\n", jj, runtimes, creal(W[0]), cimag(W[0]) );}}int main() {double complex A[N][N];double complex V[N];double complex W[N];test(A, V, W);}

return 0;2929