Hiện tượng chồng phổ trong XLA

Tài liệu này mô tả API chồng phổ XLA, cho phép bạn chỉ định tình trạng chồng phổ giữa vùng đệm đầu vào và đầu ra khi xây dựng chương trình XLA.

Xác định tình trạng chồng phổ tại thời gian biên dịch

Ví dụ: hãy xem xét một mô-đun HLO đơn giản chỉ thêm 1 vào dữ liệu đầu vào:

HloModule increment

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

Mô-đun này sẽ phân bổ 2 vùng đệm 4 byte: một cho đầu vào %p và một cho %out đầu ra.

Tuy nhiên, bạn nên thực hiện cập nhật tại chỗ (ví dụ: nếu trong giao diện người dùng tạo biểu thức, biến đầu vào không còn tồn tại sau khi tính toán, như trong p++ tăng dần).

Để thực hiện việc cập nhật như vậy một cách hiệu quả, bạn có thể chỉ định bí danh đầu vào:

HloModule increment, input_output_alias={ {}: 0 }

ENTRY entry {
  %p = f32[] parameter(0)
  %c = f32[] constant(1)
  ROOT %out = f32[] add(%p, %c)
}

Định dạng chỉ định rằng toàn bộ đầu ra (được đánh dấu bằng {}) được đặt bí danh là tham số đầu vào 0.

Để chỉ định bí danh theo phương thức lập trình, hãy xem API XlaBuilder::SetUpAlias.

Xác định tình trạng chồng phổ trong thời gian chạy

Hiện tượng chồng phổ được xác định ở bước trước được chỉ định trong quá trình biên dịch. Trong quá trình thực thi, bạn có thể dùng API LocalClient::RunAsync để chọn có đóng góp vùng đệm hay không.

Các vùng đệm đầu vào của chương trình được gói trong ExecutionInput, theo đó chứa cây MaybeOwningDeviceMemory. Nếu bộ nhớ được chỉ định là sở hữu (quyền sở hữu vùng đệm được chuyển đến thời gian chạy XLA), thì vùng đệm sẽ thực sự được đóng góp và quá trình cập nhật sẽ được thực thi tại chỗ, theo yêu cầu của API đặt bí danh thời gian biên dịch.

Tuy nhiên, nếu vùng đệm được đặt bí danh tại thời điểm biên dịch không được đóng góp trong thời gian chạy, thì tính năng bảo vệ sao chép sẽ khởi động: một bộ đệm đầu ra bổ sung O được phân bổ và nội dung của vùng đệm đầu vào P vốn được đặt bí danh sẽ được sao chép vào O (để thực thi hiệu quả chương trình như thể vùng đệm O được đóng góp trong thời gian chạy).