Desain Penerjemah

Model Data

Program StableHLO adalah komputasi pada TensorFlow (array n-dimensi), yang dalam model saat ini, diimplementasikan menggunakan class Tensor. Class penyimpanan yang mendasarinya untuk objek Tensor, detail::Buffer, menyimpan mlir::ShapedType TensorFlow bersama dengan objek mlir::HeapAsmResourceBlob yang mewakili blob data TensorFlow yang dapat diubah yang disusun sebagai array byte yang berdekatan dalam urutan utama ke minor. Objek detail::Buffer dihitung sebagai referensi untuk menyederhanakan pengelolaan memori.

Masing-masing elemen TensorFlow direpresentasikan menggunakan class Element, yang menggunakan union yang didiskriminasi yang menyimpan salah satu dari APInt, APFloat, atau pair<APFloat,APFloat> untuk penyimpanan. Yang terakhir digunakan untuk menyimpan elemen dengan jenis yang kompleks.

Tensor memiliki API berikut untuk berinteraksi dengan setiap elemennya:

  • Element Tensor::get(llvm::ArrayRef<int64_t> index): Untuk mengekstrak elemen TensorFlow individual pada indeks multi-dimensi index sebagai objek Element.
  • void Tensor::set(llvm::ArrayRef<int64_t> index, Element element);: Untuk mengupdate objek Element element menjadi TensorFlow pada indeks multi-dimensi index.

Cara kerja penerjemah

Fungsi entri ke penafsir adalah

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

yang melakukan hal berikut:

  1. Melacak argumen SSA func dan nilai Tensor runtime terkait, yang disediakan di args, menggunakan peta tabel simbol, M.
  2. Untuk setiap op dalam func, dalam urutan SSACFG:
    • Memanggil eval di operasi. Untuk setiap operand SSA operasi, ekstrak nilai runtime-nya dari M untuk disediakan sebagai argumen ke pemanggilan eval.
    • Melacak hasil SSA operasi dan nilai yang dievaluasi di M.

eval level operasi yang disebutkan di (2) bertanggung jawab untuk menerapkan semantik eksekusi operasi. Berikut adalah contoh untuk stablehlo::AddOp. Dalam contoh ini, setiap elemen individual dari TensorFlow lhs dan rhs diekstrak berpasangan sebagai objek Element yang kemudian ditambahkan. Hasil penambahan, objek Element, disimpan dalam TensorFlow result akhir.

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;
}

Secara keseluruhan, desain penafsir dioptimalkan untuk keterbacaan implementasi fungsi eval bagi operasi individual karena dimaksudkan untuk berfungsi sebagai implementasi referensi untuk StableHLO. Misalnya, alih-alih menentukan eval sebagai fungsi template dan membuat parameternya dengan jenis elemen, kami merangkum detail tentang cara berbagai jenis elemen ditangani dalam Element::operator+ dll., yang menyederhanakan implementasi eval.

Menggunakan penafsir untuk folding konstan

Kita dapat menggunakan mekanisme penafsir untuk melipat operasi dengan nilai operand yang konstan. Cuplikan kode berikut menunjukkan ide implementasi untuk melipat stablehlo::AddOp dengan operand berjenis floating-point:

OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
  auto attrs = adaptor.getOperands();
  DenseElementsAttr lhsData = attrs[0].dyn_cast<DenseElementsAttr>();
  DenseElementsAttr rhsData = attrs[1].dyn_cast<DenseElementsAttr>();
  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(element.getValue().cast<FloatAttr>().getValue());
  }

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

Saat ini, kami tidak aktif berupaya mengintegrasikan penafsir ke folding konstan karena kami tidak berencana menerapkan folder untuk StableHLO. Namun, di masa mendatang, kami berencana memanfaatkan penafsir untuk folding konstan di MHLO.Setelah itu, kami akan meningkatkan ergonomi cuplikan kode di atas (mis., kita dapat memiliki fungsi bantuan yang mengemas operand konstan ke dalam objek Tensor dan mengekstrak hasil Tensor ke dalam OpFoldResult).

Menguji penafsir StableHLO

Penafsir mengambil sebagai input (A) program StableHLO, dan (B) nilai data yang akan dimasukkan ke program, dan menghasilkan nilai data output, yang cocok dengan nilai data yang diharapkan yang diberikan pengguna. Nilai data (B) di-hard code dalam program itu sendiri menggunakan operasi stablehlo.constant. Penafsir mengevaluasi program input. Output dari operasi yang sedang diuji diperiksa melalui pemeriksaan (misalnya check.expect_eq, check.expect_almost_eq), seperti yang ditunjukkan di bawah ini. check.expect_eq dan check.expect_eq_const memeriksa kesetaraan bitwise untuk semua jenis yang didukung serta check.expect_almost_eq dan check.expect_almost_eq_const memeriksa kesetaraan yang mendekati dalam toleransi, yang dijelaskan dalam panduan pengujian (G6), untuk floating point dan jenis kompleks.

// 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
}

Utilitas pengujian stablehlo-translate --interpret (kode) bertanggung jawab untuk mengurai program, menafsirkan setiap fungsi termasuk operasi yang merupakan fungsi. Kami memiliki rangkaian pengujian khusus, yang terdiri dari beberapa pengujian yang menjalankan berbagai perilaku runtime, untuk setiap Operasi StableHLO. Pengujian dapat ditemukan di sini (misalnya, interpret_*.mlir).

Pedoman pengujian

(G1) Apakah kami perlu menguji semua jenis yang didukung untuk setiap operasi?

Kita dapat menggunakan kombinasi aturan berikut untuk memutuskan:

  1. Saat menerapkan operasi, jika ada kode dalam fungsi eval yang sesuai untuk menangani jenis tertentu, penting untuk memiliki pengujian yang mencakup jenis tersebut. Misalnya, untuk operasi add, ada kode eksklusif untuk menangani jenis integer, boolean, floating point, dan kompleks, sehingga kita memerlukan satu pengujian untuk setiap kategori jenis.

  2. Jika kumpulan jenis ditangani secara seragam dalam fungsi eval yang sesuai, satu pengujian untuk semua jenis tersebut semestinya sudah memadai. Misalnya, untuk operasi add, semua varian jenis bilangan bulat (si4, u4, si8, u8, dan seterusnya) ditangani dengan cara yang sama menggunakan llvm::APInt API, sehingga kita dapat melewati penambahan pengujian untuk setiap varian tersebut, dan sebagai gantinya menambahkan satu pengujian yang representatif. Untuk menghindari ambiguitas dalam memilih perwakilan, kita harus menggunakan panduan berikut:

    • Jika semua jenis, yang ditangani secara seragam, memiliki jenis primitif yang sama (yaitu jika semuanya adalah jenis integer, floating point, atau kompleks), maka pilih jenis dengan lebar bit maksimum.
    • Jika semua jenis, yang ditangani secara seragam, memiliki campuran jenis primitif, maka pilih jenis yang berjenis primitif berikut, dalam urutan preferensi yang menurun: bilangan bulat, floating point, boolean, kompleks.

(G2) Bagaimana cara menentukan jumlah pengujian yang diperlukan untuk mencakup perilaku operasi?

Tujuannya adalah untuk mencakup logika penafsir untuk operasi (yaitu semua kasus sudut implementasi) secara komprehensif dengan jumlah pengujian yang minimal. Meminimalkan jumlah pengujian sangat penting untuk kemudahan pemeliharaan. Semakin sedikit pengujian yang kami miliki, semakin mudah untuk meninjaunya dan memastikan bahwa pengujian tersebut mencakup operasi secara komprehensif. Dengan demikian, kami berharap sebagian besar operasi yang lebih sederhana hanya akan memiliki satu pengujian. Jika karena alasan yang baik cakupan komprehensif tidak praktis, maka tidak masalah untuk berhenti di >= 90%. Hal ini akan diputuskan berdasarkan kasus per kasus selama peninjauan permintaan pull.

(G3) Bagaimana jika menambahkan pengujian untuk infrastruktur penafsir?

Infrastruktur penerjemah sebagian besar mudah dan dapat ditambahkan ke basis kepercayaan kami. Satu-satunya bagian yang tidak umum adalah bagaimana berbagai jenis dikemas ke dalam dan diekstrak dari penyimpanan penafsir yang mendasarinya. Seperti dibahas di (G1), kami hanya akan menguji jenis operasi yang ditangani secara berbeda. Dengan demikian, ada kemungkinan bahwa kode pengemasan/un-packing, yang sesuai dengan berbagai varian jenis integer/floating point, mungkin tidak sepenuhnya tercakup selama pengujian. Untuk memastikan cakupan penuh, kita dapat memilih op seperti constant yang mendukung semua jenis elemen StableHLO dan menulis pengujian lengkap.

(G4) Jika implementasi operasi bergantung pada operasi lain, apakah kita harus menulis pengujian untuk operasi tersebut?

Tidak. Misalnya, implementasi batch_norm_grad dapat didasarkan pada divide, subtract, multiply, dan lainnya. Kita harus menghindari pengujian operasi yang terakhir saat menguji yang pertama.

(G5) Haruskah kita menulis pengujian untuk menjalankan perilaku yang ditentukan implementasi / tidak ditentukan?

Kita tidak boleh menulis pengujian yang menjalankan perilaku operasi yang ditentukan implementasi atau tidak ditentukan. Pengujian yang menjalankan perilaku yang ditentukan implementasi menunjukkan perilaku lokal penafsir yang tidak boleh digeneralisasi. Pengujian yang melakukan perilaku yang tidak ditentukan tidak berkontribusi terhadap pemahaman tentang perilaku pengoperasian.

(G6) Saat menulis pengujian untuk jenis floating point, seberapa presisi hasil yang diharapkan harus ditentukan dalam pemeriksaan?

Untuk operasi dasar (penambahan, pengurangan, perkalian, pembagian, dan persegi), implementasi yang mengikuti spesifikasi IEEE diharapkan memberikan hasil bulat dalam 0,5 ULP dari hasil yang sama persis secara matematis. Meskipun demikian, kita dapat dengan aman membayangkan hasil yang diharapkan dari operasi ini berjarak maksimal 1 ULP. Namun, cara ini mungkin tidak berfungsi untuk fungsi transendental (sine, cosine, dll.) yang jaminan presisinya ditentukan oleh implementasi (alasan).

Implementasi saat ini menggunakan nilai toleransi "one-size-fits-all" sebesar 0,0001. Contoh berikut menunjukkan toleransi di atas dalam penerapannya.

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
}

Ini hanyalah langkah pertama dalam menguji akurasi numerik operasi StableHLO. Saat ini, ini adalah area spesifikasi StableHLO yang belum ditetapkan, dan masih ada upaya untuk mengetahuinya #1156 berdasarkan pengalaman kami menggunakan StableHLO dalam praktiknya dan berdasarkan masukan dari pemangku kepentingan. Selagi pekerjaan ini berlangsung, kami akan mengupdate infrastruktur sesuai dengan proses ini.

(G7) Ada yang menjelaskan gaya coding pengujian?

  1. Pastikan untuk menggunakan nama input/output yang sebenarnya, bukan menetapkan default ke nilai SSA (misalnya %0, %1, dll.)
  2. Pastikan pengujian menggunakan format yang dicetak, jika ada.

(G8) Haruskah kami menyertakan contoh yang telah disediakan dalam spesifikasi? Ya (untuk kelengkapan pengujian).