Two more loop shapes for vectorize_loops

Aperta
  • Qualcosa non va?
#1,705 2 commenti 1 reazione 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
5/5
Tempo stimato
Più di una settimana
Idoneità per principianti
45/100
Tipo di issue
Funzionalità
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
ocaml
Ambito
compilers, performance

Direzione di ricerca

Inizia individuando il pass vectorize_loops e rivedendo il lavoro precedente a cui si fa riferimento in #1683 e #1681, in particolare l’elenco dei nomi scritti e l’idea di inline_temps. Traccia il modo in cui il pass gestisce le letture del corpo del ciclo, le dichiarazioni, i gather e l’allargamento aritmetico. Il lavoro è completo quando entrambe le forme di ciclo descritte vengono trasformate correttamente, inclusi i temporanei scalari e di tipo int, con copertura di regressione per i pattern mostrati.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

optimization

Thanks for #1683, it picks up a lot. I ran it over posteriordb and the models it rewrites got a lot faster in CmdStan (dogs 2.2x, radon_pooled 3.7x over --O1 alone). Two shapes it doesn't touch yet show up all over the corpus, so I hand-vectorized a few models to see what they'd get. Timings below are CmdStan gradient evaluations, stock stanc at --O1, medians of five runs on the same machine.

Reading something the loop body just wrote

for (n in 1:N) {
  mu[n] = alpha[county_idx[n]];
  target += normal_lpdf(log_radon[n] | mu[n], sigma_y);
}

After the first line mu is in the do-not-read set, so the second line stays a loop even though it only reads the element that was just written. Eleven of the thirteen radon models look like this (two of them write twice before reading), and sigma[n] = exp(sigma[n]) is the same thing in one line. The rewrite is just

mu = alpha[county_idx];
target += normal_lpdf(log_radon | mu, sigma_y);
model data loop vectorized speedup
radon_partially_pooled_centered radon_mn, N = 919 20.1 us 5.9 us 3.4x
radon_partially_pooled_centered radon_all, N = 12777 273 us 76 us 3.6x

I think the rule is that reading x[n] after writing x[n] earlier in the same body is fine, because the vectorized statements run in the same order. #1681 kept a list of written names for this before it was closed.

Scalar and int temporaries in the body

for (t in (K + 1):T) {
  real mu = alpha;
  for (k in 1:K) mu = mu + beta[k] * y[t - k];
  y[t] ~ normal(mu, sigma);
}

for (n in 1:Npts) {
  int irat = rat[n];
  y[n] ~ normal(alpha[irat] + beta[irat] * (x[n] - xbar), sigma_y);
}

Any declaration in the body stops the pass. A real temporary could become a vector over the loop range, and an int pulled out of a data array could be substituted into its uses as a gather. By hand:

vector[T - K] mu = rep_vector(alpha, T - K);
for (k in 1:K) mu = mu + beta[k] * to_vector(y[(K + 1 - k):(T - k)]);
y[(K + 1):T] ~ normal(mu, sigma);

y ~ normal(to_vector(alpha[rat]) + to_vector(beta[rat]) .* (to_vector(x) - xbar), sigma_y);
model data loop vectorized speedup
arK T = 200, K = 5 7.8 us 2.5 us 3.1x
rats_model Npts = 150 4.5 us 2.0 us 2.3x

arK, rats_model and the mixture models with an array[K] real temp are the corpus cases. #1681 called this inline_temps. The rats one also needs to_vector around the array gathers since arrays don't do elementwise arithmetic; the pass already widens arithmetic so maybe that part is small.

Lingua principale
OCaml
Stelle
162
Fork
60
Merge medio
2g 22h
PR unite (30g)
14

Preparare l'ambiente

  • Nessun Dockerfile né file Docker Compose
  • Ha un modello di pull request
  • Nessuna guida per i contributori

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di stan-dev/stanc3

Tutte le issue di stan-dev/stanc3

Issue simili

Altre issue su Compilers

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.