Progettazione dell'interprete

Modello dati

I programmi StabileHLO sono calcoli su tensori (array n-dimensionali), che nel modello attuale vengono implementati utilizzando classe Tensor. La classe di archiviazione sottostante per un oggetto Tensor, detail::Buffer, memorizza il mlir::ShapedType del tensore insieme a un Oggetto mlir::HeapAsmResourceBlob che rappresenta un BLOB mutabile di tensore disposti come array di byte contigui ordine maggiore a minore. detail::Buffer oggetti vengono conteggiati per riferimento per semplificare la gestione della memoria.

I singoli elementi di un tensore sono rappresentati utilizzando la classe Element, che utilizza un sindacato discriminato che detiene uno dei seguenti elementi: APInt, APFloat o pair<APFloat,APFloat> per l'archiviazione. L'ultimo viene utilizzato per memorizzare gli elementi con tipi complessi.

Tensor dispone delle seguenti API per interagire con i singoli elementi:

  • Element Tensor::get(llvm::ArrayRef<int64_t> index): per estrarre un singolo elemento tensore con indice multidimensionale index come Element .
  • void Tensor::set(llvm::ArrayRef<int64_t> index, Element element);: Aggiornare un oggetto Element element in un tensore multidimensionale indice index.

Come funziona l'interprete

La funzione di accesso all'interprete è

SmallVector<Tensor> eval(func::FuncOp func, ArrayRef<Tensor> args);

che comporta le seguenti operazioni:

  1. Tiene traccia degli argomenti SSA di func e del relativo runtime associato Tensor forniti in args, utilizzando una mappa di tabelle di simboli, M.
  2. Per ogni operazione all'interno di func, in ordine SSACFG:
    • Richiama eval nell'operazione. Per ogni operando SSA dell'operazione, estrai la sua valore runtime da M da fornire come argomento alla chiamata eval.
    • Tiene traccia dei risultati dell'SSA dell'operazione e del valore valutato in M.

Il eval a livello di operazione indicato nel punto (2) è responsabile dell'implementazione del la semantica di esecuzione dell'operazione. Di seguito è riportato un esempio per stablehlo::AddOp. Nell'esempio, i singoli elementi dei tensori lhs e rhs sono a coppie estratti come Element oggetti che vengono poi aggiunti. Il risultato dell'aggiunta, un oggetto Element, viene archiviato nel tensore result finale.

Tensor eval(AddOp op, const Tensor &lhs, const Tensor &rhs) {
  Tensor result(op.getType());

  for (auto it = result.index_begin(); it != result.index_end(); ++it)
    result.set(*it, lhs.get(*it) + rhs.get(*it));

  return result;
}

Nel complesso, il design dell'interprete è ottimizzato per la leggibilità delle implementazioni delle funzioni eval per le singole operazioni in quanto lo scopo fungere da implementazione di riferimento per StableHLO. Ad esempio, invece di definendo eval come funzione per il modello e parametrizzandola con i tipi di elementi, incapsuliamo i dettagli di come i diversi tipi di elementi vengono gestiti Element::operator+ e così via, semplificando l'implementazione di eval.

Usare l'interprete per fare piegamento costante

Possiamo usare il meccanismo dell'interprete per comprimere le operazioni con operando costante e i relativi valori. Il seguente snippet di codice illustra un'idea dell'implementazione per piegare stablehlo::AddOp con operandi digitati con rappresentazione in virgola mobile:

OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
  auto attrs = adaptor.getOperands();
  DenseElementsAttr lhsData = dyn_cast<DenseElementsAttr>(attrs[0]);
  DenseElementsAttr rhsData = dyn_cast<DenseElementsAttr>(attrs[1]);
  if (!lhsData || !rhsData) return {};

  auto lhs = Tensor(lhsData);
  auto rhs = Tensor(rhsData);
  auto result = eval(*this, lhs, rhs);

  SmallVector<APFloat> values;
  for (auto i = 0; i < result.getNumElements(); ++i) {
    Element element = result.get(i);
    values.push_back(cast<FloatAttr>(element.getValue()).getValue());
  }

  return DenseElementsAttr::get(result.getType(), values);
}

Al momento, non stiamo lavorando attivamente per integrare l'interprete nei piegamento costante perché non abbiamo intenzione di implementare la cartella per StableHLO. Tuttavia, in futuro, abbiamo intenzione di affidarci all'interprete per folding in MHLO, a quel punto miglioreremo l'ergonomia dello snippet di codice sopra (ad esempio, potremmo avere una funzione helper che raggruppa operandi costanti in Tensor oggetti e decomprime Tensor risultati in OpFoldResult).

Test dell'interprete StableHLO

L'interprete prende come input (A) un programma StableHLO e (B) i valori dei dati da inviare al programma e genera valori di dati di output, che vengono rispetto ai valori dei dati previsti forniti dall'utente. I valori dei dati (B) sono hardcoded nel programma stesso tramite le operazioni stablehlo.constant. La l'interprete valuta il programma di input. Gli output dell'operazione sottoposta a test viene verificata tramite controlli (ad es. check.expect_eq, check.expect_almost_eq), come come mostrato di seguito. check.expect_eq e check.expect_eq_const controllano il bit a bit pari a qualsiasi tipo supportato, nonché check.expect_almost_eq e check.expect_almost_eq_const verifica la quasi uguaglianza all'interno di una tolleranza descritto nella linea guida per i test (G6), per i tipi a virgola mobile e complessi.

// CHECK-LABEL: Evaluated results of function: add_op_test_ui4
func.func @add_op_test_ui4() {
  %0 = stablehlo.constant dense<[0, 2]> : tensor<2xui4>
  %1 = stablehlo.constant dense<[15, 3]> : tensor<2xui4>
  %2 = stablehlo.add %0, %1 : tensor<2xui4>
  check.expect_eq_const %2, [15, 5] : tensor<2xui4>
  func.return
}

Un'utilità di test stablehlo-translate --interpret (codice) responsabile dell'analisi del programma, dell'interpretazione di ogni funzione, incluso operazioni che costituiscono la funzione. Disponiamo di una sala di test dedicata, composta di diversi test che esercitano vari comportamenti di runtime, per ogni operazione StableHLO. I test sono disponibili qui.

Linee guida per l'esecuzione dei test

(G1) Dobbiamo testare tutti i tipi supportati per ogni operazione?

Possiamo utilizzare una combinazione delle seguenti regole per decidere:

  1. Durante l'implementazione di un'operazione, se esiste del codice nella riga eval corrispondente per gestire un particolare tipo, è imperativo avere test per coprire quel tipo. Ad esempio, per l'operazione add è disponibile codice esclusivo per gestire tipi interi, booleani, con rappresentazione in virgola mobile e complessi; di conseguenza, occorre un test per ciascuna categoria di tipi.

  2. Se un insieme di tipi viene gestito in modo uniforme nella funzione eval corrispondente, dovrebbe essere sufficiente un solo test per tutti questi tipi. Ad esempio, per l'operazione add, tutte le varianti dei tipi interi (si4, u4, si8, u8 e così via) vengono gestiti allo stesso modo utilizzando le API llvm::APInt, quindi possiamo saltare aggiungere test per ciascuna di queste varianti e aggiungere invece un singolo test rappresentativo. Per evitare ambiguità nella scelta del rappresentante, devono utilizzare le seguenti linee guida:

    • Se tutti i tipi, gestiti in modo uniforme, hanno lo stesso tipo primitivo (ad es. se sono tutti numeri interi, con rappresentazione in virgola mobile o tipi complessi), scegli quella con la larghezza in bit massima.
    • Se tutti i tipi, gestiti in modo uniforme, presentano una combinazione di tipi primitivi, scegli quello con il seguente tipo primitivo, in ordine decrescente di preferenza: numero intero, virgola mobile, booleano, complesso.

(G2) Come decidiamo il numero di test necessari per coprire la comportamento dell'utente?

L'obiettivo è comprendere in modo esauriente la logica dell'interprete per l'operazione (ovvero tutti i casi d'angolo dell'implementazione) con un numero minimo di test. Ridurre al minimo il numero di test è importante per la manutenibilità. Meno test abbiamo, più facile sarà rivederli e assicurarci che coprire in modo esaustivo l'operazione. Di conseguenza, ci aspettiamo che la maggior parte dei modelli operazioni finiranno per avere un solo test. Se per qualche motivo valido una copertura non attuabile, è possibile fermarsi a valori >= 90%. Verrà deciso caso per caso durante la revisione delle richieste di pull.

(G3) Che ne dici di aggiungere test per l'infrastruttura dell'interprete?

L'infrastruttura degli interpreti è per lo più semplice e può essere aggiunta la nostra base di fiducia. L'unica cosa non banale è il modo in cui i vari tipi vengono raggruppati e scomposto dallo spazio di archiviazione sottostante per l'interprete. Come già detto in (G1), testerà solo i tipi di operazioni gestiti in modo diverso. Con che è possibile che il codice di imballaggio/disimballaggio, corrispondente di tipi interi/con rappresentazione in virgola mobile, potrebbero non essere completamente coperte durante test. Per garantire una copertura completa, possiamo scegliere un'operazione come constant che supporta tutti i tipi di elementi StableHLO e scrive test esaustivi.

(G4) Se l'implementazione di un'operazione dipende da altre operazioni, dobbiamo scrivere test per quest'ultimo?

No. Ad esempio, l'implementazione di batch_norm_grad può essere basata su divide, subtract, multiply e altri. Dovremmo evitare di testare quest'ultima durante il test della prima.

(G5) È necessario scrivere dei test per esercitare uno stile definito dall'implementazione / non definito comportamenti?

Non dobbiamo scrivere test che si avvalgono dei metodi di implementazione definiti o comportamenti indefiniti dell'operazione. Testano l'esercizio di comportamenti definiti dall'implementazione dimostrare un comportamento locale dell'interprete che non dovrebbe essere in generale. I test che praticano comportamenti indefiniti non contribuiscono al la comprensione del comportamento dell'operazione.

(G6) Durante la scrittura dei test per i tipi a virgola mobile, con quale precisione il risultato previsto deve essere specificato nei controlli?

Per le operazioni elementari (addizione, sottrazione, moltiplicazione, divisione e quadrata), un'implementazione che segue le specifiche IEEE dovrebbe fornire una risultato arrotondato entro 0,5 ULP del risultato matematicamente esatto. Detto questo, possiamo tranquillamente immaginare il risultato previsto che usciva da queste operazioni al massimo 1 ULP a parte. Tuttavia, questo metodo potrebbe non funzionare per le funzioni trascendentali (sine, cosine e così via) per cui sono previste garanzie di precisione. dell'implementazione (razionale).

L'implementazione attuale utilizza un modello "universale" di tolleranza pari a 0,0001. L'esempio seguente mostra l'applicazione della tolleranza precedente.

func.func @check_tolerance() {
  %0 = stablehlo.constant dense<0.2> : tensor<f32>

  // The following check succeeds as %0 is almost equal to the provided
  // constant modulo the tolerance, mentioned above.
  check.expect_almost_eq_const %0, dense<0.19999> : tensor<f32>

  // The following check fails as %0 is not bitwise equal to the provided
  // constant.
  check.expect_eq_const %0, dense<0.19999> : tensor<f32>

  func.return
}

Questo è solo il primo passaggio per testare l'accuratezza numerica delle operazioni StableHLO. Al momento, questa è un'area sottospecificata delle specifiche StableHLO e di lavoro per trovare una soluzione #1156 sulla base della nostra esperienza pratica nell'utilizzo di StableHLO e del feedback di le parti interessate. Man mano che questa operazione procede, aggiorneremo l'infrastruttura di conseguenza.

(G7) Qualcosa sullo stile di programmazione dei test?

  1. Assicurati di utilizzare il nome effettivo degli input/output anziché dei valori predefiniti ai valori SSA (ad es. %0, %1 e così via)
  2. Assicurati che i test utilizzino un formato stampato, se esistente.

(G8) Dovremmo includere l'esempio già fornito nella specifica? Sì (per completezza dei test).