- 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-start
vàall-reduce-done
). - 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
, kAsyncUpdate
và kAsyncDone
. Ý 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-start
và
async-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".