Thiết kế phiên dịch viên

Mô hình dữ liệu

Chương trình StableHLO là các phép tính toán trên tensor (các mảng n chiều), trong mô hình hiện tại, được triển khai bằng cách sử dụng lớp Tensor. Lớp bộ nhớ cơ bản cho đối tượng Tensor, detail::Buffer, lưu trữ mlir::ShapedType của tensor cùng với một đối tượng mlir::HeapAsmResourceBlob đại diện cho một blob có thể thay đổi của dữ liệu tensor được bố trí dưới dạng mảng byte liền kề theo thứ tự từ lớn đến phụ. Các đối tượng detail::Buffer được tính tham chiếu để đơn giản hoá việc quản lý bộ nhớ.

Các phần tử riêng lẻ của tensor được biểu thị bằng cách sử dụng lớp Element, lớp này sử dụng liên kết phân biệt giữ một trong các APInt, APFloat hoặc pair<APFloat,APFloat> để lưu trữ. Lớp cuối cùng được dùng để lưu trữ các phần tử có kiểu phức tạp.

Tensor có các API sau đây để tương tác với các phần tử riêng lẻ:

• Element Tensor::get(llvm::ArrayRef<int64_t> index): Để trích xuất một phần tử tensor riêng lẻ tại chỉ mục đa chiều index dưới dạng đối tượng Element.
• void Tensor::set(llvm::ArrayRef<int64_t> index, Element element);: Để cập nhật một đối tượng Element element thành một tensor tại chỉ mục đa chiều index.

Cách hoạt động của phiên dịch

Hàm nhập thông tin cho trình thông dịch là

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

thực hiện những việc sau:

1. Theo dõi các đối số SSA của func và các giá trị Tensor thời gian chạy được liên kết với những đối số này, được cung cấp trong args, bằng cách sử dụng bản đồ bảng ký hiệu, M.
2. Đối với mỗi op trong func, theo thứ tự SSACFG:
   • Gọi eval trên op. Đối với mỗi toán hạng SSA của op, hãy trích xuất giá trị thời gian chạy từ M để cung cấp làm đối số cho lệnh gọi eval.
   • Theo dõi(các) kết quả SSA của hoạt động và giá trị được đánh giá trong M.

eval ở cấp hoạt động được đề cập trong (2) chịu trách nhiệm triển khai ngữ nghĩa thực thi của hoạt động. Sau đây là ví dụ về stablehlo::AddOp. Trong ví dụ này, các phần tử riêng lẻ của tensor lhs và rhs được trích xuất theo cặp dưới dạng đối tượng Element, sau đó được thêm vào. Kết quả của việc thêm đối tượng Element sẽ được lưu trữ trong tensor result cuối cùng.

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

Nhìn chung, thiết kế của trình thông dịch được tối ưu hoá để dễ đọc cách triển khai các hàm eval cho từng hoạt động riêng lẻ, vì đây là phương thức triển khai tham chiếu cho StableHLO. Ví dụ: thay vì xác định eval làm hàm mẫu và tham số hoá bằng các loại phần tử, chúng ta đóng gói thông tin chi tiết về cách xử lý các loại phần tử khác nhau trong Element::operator+, v.v., đơn giản hoá quá trình triển khai eval.

Dùng trình thông dịch để liên tục gập

Chúng ta có thể sử dụng cơ chế thông dịch để gấp các toán tử có giá trị toán hạng không đổi. Đoạn mã sau đây minh hoạ ý tưởng triển khai cách gập stablehlo::AddOp bằng các toán hạng được nhập dấu phẩy động:

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

Hiện tại, chúng tôi chưa tích cực tích hợp trình thông dịch vào chế độ gập liên tục vì chúng tôi không có kế hoạch triển khai thư mục cho StableHLO. Tuy nhiên, trong tương lai, chúng tôi dự định tận dụng trình thông dịch để gấp gọn liên tục trong MHLO. Tại thời điểm đó, chúng tôi sẽ cải thiện tính hiệu quả của đoạn mã ở trên (ví dụ: chúng ta có thể có một hàm trợ giúp đóng gói các toán hạng hằng số vào các đối tượng Tensor và giải nén kết quả Tensor thành OpFoldResult).

Kiểm thử trình thông dịch StableHLO

Trình thông dịch lấy dữ liệu đầu vào (A) một chương trình StableHLO và (B) các giá trị dữ liệu được cấp cho chương trình và tạo các giá trị dữ liệu đầu ra. Các giá trị này được so khớp với các giá trị dữ liệu dự kiến do người dùng cung cấp. Các giá trị dữ liệu (B) được mã hoá cứng trong chính chương trình bằng cách sử dụng các thao tác stablehlo.constant. Trình diễn giải sẽ đánh giá chương trình đầu vào. (Các) đầu ra của ứng dụng đang được kiểm thử sẽ được kiểm tra thông qua các lượt kiểm tra (ví dụ: check.expect_eq, check.expect_almost_eq), như minh hoạ dưới đây. check.expect_eq và check.expect_eq_const kiểm tra đẳng thức bitwise cho mọi loại được hỗ trợ, cũng như check.expect_almost_eq và check.expect_almost_eq_const kiểm tra xem có đẳng thức gần bằng trong dung sai hay không, được giải thích trong nguyên tắc kiểm thử (G6) cho các loại dấu phẩy động và phức tạp.

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

Tiện ích kiểm thử stablehlo-translate --interpret (mã) chịu trách nhiệm phân tích cú pháp chương trình, diễn giải từng hàm bao gồm cả các thao tác cấu thành hàm đó. Chúng tôi có một bộ kiểm thử chuyên biệt, bao gồm một số chương trình kiểm thử thực thi nhiều hành vi trong thời gian chạy, cho mỗi Op (hoạt động) StableHLO. Bạn có thể xem các chương trình kiểm thử tại đây (ví dụ: filter_*.mlir).

Nguyên tắc kiểm thử

(G1) Chúng tôi có cần thử nghiệm tất cả các loại được hỗ trợ cho mọi hoạt động không?

Chúng ta có thể sử dụng kết hợp các quy tắc sau để quyết định:

1. Trong khi triển khai một op, nếu tồn tại mã trong hàm eval tương ứng để xử lý một loại cụ thể, thì bạn bắt buộc phải có(các) bài kiểm thử để bao gồm loại đó. Ví dụ: đối với hoạt động add, chúng tôi có mã độc quyền để xử lý các kiểu số nguyên, boolean, dấu phẩy động và phức tạp. Do đó, chúng tôi cần một bài kiểm thử cho mỗi danh mục kiểu.

2. Nếu một tập hợp kiểu được xử lý thống nhất trong hàm eval tương ứng, thì bạn chỉ cần kiểm thử một lần cho tất cả các kiểu đó. Ví dụ: đối với op add, tất cả biến thể của các kiểu số nguyên (si4, u4, si8, u8 v.v.) đều được xử lý tương tự bằng cách sử dụng API llvm::APInt. Do đó, chúng ta có thể bỏ qua việc thêm kiểm thử cho từng biến thể đó và thay vào đó hãy thêm một kiểm thử đại diện duy nhất. Để tránh sự không rõ ràng trong việc chọn đại diện, chúng ta nên sử dụng các nguyên tắc sau:

   • Nếu tất cả các kiểu được xử lý thống nhất đều có cùng kiểu dữ liệu nguyên gốc (tức là nếu tất cả đều là số nguyên, dấu phẩy động hoặc loại phức tạp), thì hãy chọn loại có chiều rộng bit tối đa.
   • Nếu tất cả các loại được xử lý thống nhất và có kết hợp các loại nguyên hàm, hãy chọn loại có loại nguyên gốc sau đây theo thứ tự giảm dần: số nguyên, dấu phẩy động, boolean, phức tạp.

(G2) Làm cách nào để chúng tôi quyết định số lượng thử nghiệm cần thiết để kiểm thử hành vi của một op?

Mục tiêu là bao gồm toàn diện logic của thông dịch viên cho op (cụ thể là tất cả các trường hợp góc của hoạt động triển khai) với số lượng kiểm thử tối thiểu. Việc giảm thiểu số lượng bài kiểm thử là rất quan trọng đối với khả năng bảo trì. Càng có ít thử nghiệm thì chúng tôi càng dễ dàng xem xét chúng và để đảm bảo chúng bao quát toàn bộ hoạt động. Do đó, chúng tôi dự kiến rằng hầu hết các hoạt động đơn giản hơn sẽ chỉ có một thử nghiệm. Nếu vì một lý do chính đáng mà phạm vi bao phủ toàn diện là không hợp lý, thì bạn có thể dừng ở mức >= 90%. Điều này sẽ được quyết định theo từng trường hợp trong quá trình xem xét yêu cầu lấy dữ liệu.

(G3) Vậy còn việc thêm bài kiểm thử cho cơ sở hạ tầng phiên dịch thì sao?

Cơ sở hạ tầng phiên dịch hầu như đơn giản và có thể được thêm vào cơ sở tin cậy của chúng tôi. Chỉ có một phần không quan trọng là cách nhiều loại được đóng gói và giải nén từ bộ nhớ thông dịch cơ bản. Như đã thảo luận trong (G1), chúng tôi sẽ chỉ kiểm thử các loại hoạt động được xử lý theo cách khác. Do đó, mã đóng gói/huỷ đóng gói, tương ứng với các biến thể khác nhau của các kiểu số nguyên/dấu phẩy động, có thể không được bao phủ đầy đủ trong quá trình kiểm thử. Để đảm bảo bao phủ đầy đủ, chúng ta có thể chọn một hoạt động như constant hỗ trợ tất cả các loại phần tử StableHLO và viết chương trình kiểm thử đầy đủ.

(G4) Nếu việc triển khai một op phụ thuộc vào các hoạt động khác, thì chúng tôi có nên viết kiểm thử cho hoạt động sau không?

Không. Ví dụ: việc triển khai batch_norm_grad có thể dựa trên divide, subtract, multiply và các yếu tố khác. Chúng ta nên tránh kiểm thử các hoạt động sau trong khi kiểm thử các hoạt động sau.

(G5) Chúng tôi có nên viết mã kiểm thử để thực hiện các hành vi do quy trình triển khai xác định / không xác định không?

Chúng ta không nên viết các kiểm thử thực thi các hành vi triển khai đã xác định hoặc không xác định của op. Các kiểm thử thực hiện các hành vi do phương thức triển khai xác định sẽ minh hoạ một hành vi cục bộ của trình thông dịch không nên được khái quát hoá. Các kiểm thử thực hiện hành vi không xác định không đóng góp vào việc hiểu được hành vi của ứng dụng.

(G6) Trong khi viết mã kiểm thử cho các loại dấu phẩy động, bạn cần chỉ định kết quả dự kiến trong các lượt kiểm tra chính xác đến mức nào?

Đối với các phép toán cơ bản (cộng, trừ, nhân, chia và bình phương), việc triển khai theo thông số kỹ thuật của IEEE dự kiến sẽ cung cấp kết quả làm tròn trong phạm vi 0,5 ULP so với kết quả chính xác về mặt toán học. Dù vậy, chúng ta có thể yên tâm hình dung kết quả dự kiến từ các hoạt động này sẽ cách nhau ít nhất 1 ULP. Tuy nhiên, cách này có thể không hiệu quả đối với các hàm siêu vượt (sine, cosine, v.v.) mà sự đảm bảo về độ chính xác được xác định theo triển khai (lý do).

Cách triển khai hiện tại sử dụng giá trị dung sai "một kích thước-phù hợp cho tất cả" là 0,0001. Ví dụ sau minh hoạ hành vi khoan dung nêu trên trong thực tế.

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
}

Đây chỉ là bước đầu tiên trong quá trình kiểm thử độ chính xác bằng số của các hoạt động StableHLO. Hiện tại, đây là một khía cạnh chưa được xác định cụ thể trong thông số kỹ thuật của StableHLO và chúng tôi vẫn đang tiếp tục tìm ra #1156 dựa trên kinh nghiệm của chúng tôi khi sử dụng StableHLO trong thực tế và dựa trên ý kiến phản hồi của các bên liên quan. Khi quá trình này tiếp tục, chúng tôi sẽ cập nhật cơ sở hạ tầng tương ứng.

(G7) Có gì về kiểu lập trình của thử nghiệm không?

1. Hãy nhớ sử dụng tên thực tế của các đầu vào/đầu ra thay vì đặt mặc định thành các giá trị SSA (ví dụ: %0, %1, v.v.)
2. Đảm bảo kiểm thử sử dụng định dạng được in đẹp, nếu có.

(G8) Có nên thêm ví dụ đã được cung cấp trong phần quy cách không? Có (để hoàn tất thử nghiệm).