Hướng dẫn HLO không đồng bộ

  1. Việc thêm các hoạt động không đồng bộ vào HLO là cồng kềnh (tức là all-reduce-startall-reduce-done).
  2. Phần phân tách bắt đầu và hoàn tất có thể không phù hợp cho một số trường hợp sử dụng không đồng bộ trường hợp.

Để giải quyết thiếu sót đầu tiên, chúng tôi đề xuất giới thiệu một loạt mã hoạt động không đồng bộ: kAsyncStart, kAsyncUpdatekAsyncDone. Ý tưởng là tạo một mã hoạt động không đồng bộ chung có thể gói bất kỳ lệnh HLO nào. Thao tác thực tế được thực hiện không đồng bộ sẽ được mã hoá bằng một phép tính được gọi chỉ có lệnh là gốc và bất kỳ tham số cho đầu vào. Xử lý vùng đệm đầu vào/đầu ra đang bay và tình trạng chồng phổ sau đó có thể được chia sẻ cho bất kỳ hoạt động không đồng bộ nào. Lệnh bắt đầu không đồng bộ thì hình dạng đầu ra sẽ là một bộ dữ liệu của toán hạng đầu vào, giá trị đầu ra và bất kỳ trạng thái trung gian cần thiết cho async-update hoặc async-done .

%async_op {
  %param0 = f32[64] parameter(0)
  ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=foo
}

%async-start = (f32[64], f32[32], s32[]) async-start(f32[64] %operand),
                                         calls=%async_op
%async-done = f32[32] async-done((f32[64], f32[32], s32[]) %async-start)

Trong cách biểu diễn ở trên, chỉ async-start có phép tính được gọi kể từ thì việc tìm thấy tác dụng của async-done bằng cách theo dõi toán hạng để tìm async-start tương ứng để tìm phép tính được gọi.

Ngoài ra, hãy lưu ý rằng phần tử đầu tiên trong bộ dữ liệu đầu ra của các bí danh async-start có thuộc tính toán hạng, để vùng đệm vẫn hoạt động cho đến khi có lệnh không đồng bộ thực hiện. Tương tự, bí danh của phần tử thứ hai có kết quả là async-done và Phần tử thứ ba là trạng thái ngữ cảnh được sử dụng để theo dõi hoạt động không đồng bộ. Cách biểu diễn này cũng hỗ trợ nhiều tensor trong đầu vào và/hoặc đầu ra của hoạt động không đồng bộ và hoạt động chồng phổ hoạt động như nhau :

%async_op {
  %param0 = f32[64] parameter(0)
  %param1 = f32[64] parameter(1)
  ROOT %op = (f32[32], f32[32]) op(f32[64] %param0, f32[64] %param1),
                                op_specific_attr=foo
}

%async-start = ((f32[64], f32[64]), (f32[32], f32[32]), s32[])
               async-start(f32[64] %operand0, f32[64] %operand1),
               calls=%async_op
%async-done = (f32[32], f32[32]) async-done(%async-start)

Ngoài ra, op có thể bị phân rã thêm thành 0 hoặc nhiều async-update các bước thực hiện phép tính trung gian. Quá trình chồng phổ đầu vào/đầu ra hoạt động theo cách tương tự với lệnh async-update, trong đó mỗi async-startasync-update hướng dẫn phải có một người dùng là người khác async-update hoặc async-done:

%async_op {
  %param0 = f32[64] parameter(0)
  ROOT %op = f32[32] op(f32[64] %param0), op_specific_attr=foo
}

%async-start = (f32[64], f32[32], s32[]) async-start(f32[64] %operand),
                                         calls=%async_op
%async-update0 = (f32[64], f32[32], s32[]) async-update(
                           (f32[64], f32[32], s32[]) %async-start)
%async-update1 = (f32[64], f32[32], s32[]) async-update(
                           (f32[64], f32[32], s32[]) %async-update0)
%async-done = f32[32] async-done((f32[64], f32[32], s32[]) %async-update1)

Cú pháp sugar

Vì có một phép tính riêng để xác định hoạt động sẽ được thực hiện không đồng bộ là hơi cồng kềnh, nên chúng tôi cũng đề xuất một cú pháp dễ hiểu để tự động in và phân tích cú pháp các thao tác không đồng bộ như thể chúng là hạng nhất mã hoạt động. Ý tưởng là xử lý các hậu tố "-bắt đầu", "-cập nhật" và "-Xong" đặc biệt bằng cách tự động tạo phép tính và chỉ dẫn (mà không có hậu tố) khi phân tích cú pháp. Ví dụ: đoạn mã ở trên có thể được in đẹp vào giá trị sau và hai giá trị này có thể được phân tích cú pháp thành cùng một cách trình bày:

%op-start = (f32[64], f32[32], s32[]) op-start(f32[64] %operand),
                                      op_specific_attr=foo
%op-update0 = (f32[64], f32[32], s32[]) op-update(
                        (f32[64], f32[32], s32[]) %op-start),
                        op_specific_attr=foo
%op-update1 = (f32[64], f32[32], s32[]) op-update(
                        (f32[64], f32[32], s32[]) %op-update0)
%op-done = f32[32] op-done((f32[64], f32[32], s32[]) %op-update1)

Để không tạo ra sự không rõ ràng, trình xác minh sẽ không cho phép thao tác được bao bọc bằng async-start nếu chúng ta đã xác định rõ một mã hoạt động cho điều đó hoạt động có hậu tố "-bắt đầu" và/hoặc "hoàn tất". Đây cũng là một lối thoát nếu chúng tôi có bất kỳ hướng dẫn nào yêu cầu xử lý cấp độ HLO không phù hợp với mô hình được mô tả ở trên (ví dụ: đầu vào/đầu ra răng cưa vùng đệm). Vì vậy, ban đầu, copy-start/copy-done, collective-permute-start/collective-permute-done, v.v. sẽ tiếp tục sử dụng mã hoạt động hạng nhất tương ứng của chúng thay vì mã mới async-start/async-done mã cho đến khi chúng tôi dọn dẹp mã để xóa Mã opcode "-start"/"-done".