Desain Penerjemah

Model Data

Program SttableHLO adalah komputasi melalui tensor (array n-dimensi), yang, dalam model saat ini, diimplementasikan menggunakan kelas Tensor. Class penyimpanan yang mendasarinya untuk objek Tensor, detail::Buffer, menyimpan mlir::ShapedType tensor bersama dengan Objek mlir::HeapAsmResourceBlob yang mewakili blob tensor yang dapat diubah data yang disusun sebagai array byte yang berurutan urutan besar ke kecil. Objek detail::Buffer dihitung dengan referensi untuk menyederhanakan manajemen memori.

Masing-masing elemen tensor direpresentasikan menggunakan class Element, yang menggunakan serikat buruh yang didiskriminasi yang memiliki salah satu dari APInt, APFloat, atau pair<APFloat,APFloat> untuk penyimpanan. Yang terakhir digunakan untuk menyimpan elemen dengan tipe yang kompleks.

Tensor memiliki API berikut untuk berinteraksi dengan setiap elemennya:

  • Element Tensor::get(llvm::ArrayRef<int64_t> index): Untuk mengekstrak elemen tensor individu pada indeks multi-dimensi index sebagai Element .
  • void Tensor::set(llvm::ArrayRef<int64_t> index, Element element);: Untuk mengupdate element objek Element menjadi tensor pada multi-dimensi indeks 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 runtime yang terkait, Tensor nilai, yang disediakan dalam args, menggunakan peta tabel simbol, M.
  2. Untuk setiap op dalam func, dalam urutan SSACFG:
    • Memanggil eval di op. Untuk setiap operand SSA dari op, ekstrak nilai runtime dari M yang akan diberikan sebagai argumen untuk pemanggilan eval.
    • Melacak hasil SSA dari operasi dan nilai yang dievaluasi di M.

eval level op yang disebutkan di (2) bertanggung jawab untuk mengimplementasikan semantik eksekusi op. Berikut adalah contoh untuk stablehlo::AddOp. Dalam contoh, elemen individual tensor lhs dan rhs berpasangan diekstrak sebagai objek Element yang kemudian ditambahkan. Hasil penambahan tersebut, objek Element, disimpan pada tensor 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 penerjemah dioptimalkan agar mudah dibaca implementasi fungsi eval untuk operasi individual karena dimaksudkan untuk berfungsi sebagai implementasi referensi untuk StableHLO. Misalnya, daripada menentukan eval sebagai fungsi template dan memparameterkannya dengan jenis elemen, kita mengenkapsulasi detail tentang bagaimana berbagai jenis elemen ditangani dalam Element::operator+ dll., sehingga menyederhanakan implementasi eval.

Menggunakan penafsir untuk folding yang konstan

Kita dapat menggunakan mekanisme penafsir untuk melipat operasi dengan operand konstanta masing-masing. Cuplikan kode berikut menunjukkan ide penerapan untuk melipat stablehlo::AddOp dengan operand yang memiliki jenis floating point:

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

Saat ini, kami tidak secara aktif mengintegrasikan penerjemah ke dalam folding konstan karena kita tidak berencana menerapkan folder untuk StableHLO. Namun, di masa mendatang, kami berencana untuk memanfaatkan interpreter untuk melipat di MHLO, yang pada saat itu kita akan meningkatkan ergonomi cuplikan kode di atas (misalnya kita bisa memiliki fungsi bantuan yang mengemas operand konstanta ke Objek Tensor dan mengekstrak hasil Tensor ke OpFoldResult).

Menguji penafsir StableHLO

Penerjemah mengambil sebagai input (A) program StableHLO, dan (B) nilai data untuk dipasok ke program, dan menghasilkan nilai data output, yang dicocokkan dengan nilai data diharapkan yang diberikan pengguna. Nilai-nilai data (B) adalah hard code dalam program itu sendiri menggunakan operasi stablehlo.constant. Tujuan {i>Interpreter<i} mengevaluasi program {i>input.<i} Output dari operasi yang sedang diuji diperiksa melalui pemeriksaan (mis. check.expect_eq, check.expect_almost_eq), sebagai seperti yang ditampilkan di bawah ini. check.expect_eq dan check.expect_eq_const memeriksa bitwise untuk semua jenis yang didukung, dan check.expect_almost_eq serta check.expect_almost_eq_const memeriksa kesetaraan yang hampir sama dalam suatu toleransi, dijelaskan dalam pedoman pengujian (G6), untuk jenis floating point dan 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 membentuk fungsi. Kami memiliki rangkaian pengujian khusus, yang terdiri dari beberapa pengujian yang menjalankan berbagai perilaku runtime, untuk setiap Op StableHLO. Pengujian dapat ditemukan di sini.

Pedoman pengujian

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

Kita dapat menggunakan kombinasi aturan berikut untuk memutuskan:

  1. Saat mengimplementasikan op, jika ada kode dalam eval yang sesuai fungsi khusus untuk menangani jenis tertentu, maka sangat penting untuk melakukan pengujian untuk membahas jenis itu. Sebagai contoh, untuk operasi add, ada kode eksklusif untuk menangani bilangan bulat, boolean, floating point, dan jenis kompleks, sehingga kita memerlukan satu tes untuk setiap kategori jenis.

  2. Jika serangkaian jenis ditangani secara seragam dalam fungsi eval yang sesuai, maka satu pengujian untuk semua jenis itu sudah cukup. Sebagai contoh, untuk op 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 menambahkan pengujian untuk setiap varian, lalu menambahkan satu dan representatif. Untuk menghindari ambiguitas dalam memilih perwakilan, kita sebaiknya menggunakan panduan berikut:

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

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

Tujuannya adalah untuk secara komprehensif membahas logika penafsir untuk operasi (yaitu semua kasus penerapan yang berbeda) dengan jumlah pengujian minimal. Meminimalkan jumlah pengujian penting untuk kemudahan pemeliharaan. Semakin sedikit pengujian kita miliki, akan semakin mudah untuk meninjaunya dan memastikan bahwa mereka membahas operasi secara komprehensif. Hasilnya, kami berharap bahwa sebagian besar cara yang lebih sederhana hanya akan ada satu pengujian. Jika, karena alasan tertentu, cakupan tidak praktis, maka tidak masalah untuk berhenti di >= 90%. Hal ini akan diputuskan secara kasus per kasus selama peninjauan permintaan pull.

(G3) Bagaimana dengan menambahkan pengujian untuk infrastruktur penerjemah?

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

(G4) Jika implementasi sebuah operasi bergantung pada operasi lain, apakah kita harus menulis beberapa pengujian untuk yang kedua?

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

(G5) Haruskah kita menulis pengujian untuk menggunakan implementasi yang ditentukan / belum ditentukan perilaku model?

Kita tidak boleh menulis pengujian yang latihannya ditentukan oleh implementasi atau perilaku operasi yang tidak terdefinisi. Pengujian yang menerapkan perilaku yang ditentukan implementasi menunjukkan perilaku lokal penafsir yang seharusnya tidak digeneralisasi. Pengujian yang menerapkan perilaku yang tidak terdefinisi tidak berkontribusi terhadap memahami perilaku operasi.

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

Untuk operasi dasar (penambahan, pengurangan, perkalian, pembagian, dan persegi), implementasi yang mengikuti spesifikasi IEEE diharapkan memberikan hasil yang dibulatkan dalam 0,5 ULP dari hasil pasti secara matematis. Meskipun begitu, kita dengan aman membayangkan hasil yang diharapkan dari operasi ini berada di terpisah maksimal 1 ULP. Namun, ini mungkin tidak berhasil untuk fungsi transendental (sine, cosine, dll.) yang jaminan presisinya ditentukan implementasi (alasan).

Implementasi saat ini menggunakan metode "one-size-fits-all" (satu ukuran untuk semua) nilai toleransi sebesar 0,0001. Contoh berikut menunjukkan penerapan toleransi di atas.

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, area ini belum ditetapkan dari spesifikasi StableHLO, dan ada upaya berkelanjutan untuk mengetahuinya #1156 berdasarkan pengalaman kami dalam menggunakan StableHLO dalam praktiknya dan berdasarkan masukan dari pemangku kepentingan proyek. Seiring berjalannya proses ini, kami akan memperbarui infrastruktur sebagaimana mestinya.

(G7) Ada hal tentang gaya coding pengujian?

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

(G8) Haruskah kita menyertakan contoh yang sudah diberikan dalam spesifikasi? Ya (untuk kelengkapan pengujian).