Hình dạng và bố cục

Cấu trúc của một XLA Op

Hãy xem xét một ví dụ về HLO:

add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)}
          add(exponential.183, broadcast.3115)

Thành phần này bao gồm những thành phần sau:

Tên Op: add.936
- Đây là tên duy nhất của thao tác.
Hình dạng: bf16[8,1,1280,16384]
- Đây là hình dạng đầu ra của Op. Ở đây, dtype là bf16 và hình dạng là [8,1,1280,16384].
Bố cục (có tính năng Xếp lát): 3,2,0,1:T(8,128)(2,1)
- Điều này mô tả cách mảng được lưu trữ trong bộ nhớ. 3,2,0,1 biểu thị thứ tự của các trục trong bộ nhớ (tức là cột chính, hàng chính, v.v.) và T(8,128)(2,1) biểu thị việc phân ô và khoảng đệm được dùng.
- Bạn không bắt buộc phải sử dụng bố cục. Nếu không được chỉ định, sẽ không có việc xếp lát và các phương diện được giả định là được sắp xếp từ phương diện chính nhất đến phương diện phụ nhất.
Thao tác: add
- Thao tác đang được thực hiện. Ở đây, đó là Add (Thêm), cũng được đề cập trong tên Op.
Đối số: exponential.183, broadcast.3115
- Thao tác này nhận hai đối số, được chỉ định bằng tên riêng biệt của chúng.

Hãy xem xét một ví dụ khác, một Op hợp nhất:

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)}
            fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32),
            kind=kCustom, calls=%all-reduce-scatter.3

Ngoài các thành phần đã mô tả trước đó, thành phần này còn bao gồm:

Thuộc tính: kind và calls
- Các thông số này cung cấp thêm thông tin về thao tác đang được thực hiện, trong trường hợp này là hợp nhất.
Vị trí bộ nhớ (giá trị nhận dạng không gian bộ nhớ): S(1)
- Điều này biểu thị dung lượng bộ nhớ/vị trí lưu trữ mảng. S(1) ở đây biểu thị mảng này nằm trong VMEM (trên TPU).
Thông tin chi tiết về hình dạng và bố cục cho đối số đầu vào %fusion.32

Các phần sau đây mô tả về Hình dạng, Bố cục và Mã nhận dạng không gian bộ nhớ. Bạn có thể tìm hiểu thêm về tính năng Xếp kề trong phần Bố cục xếp kề.

Hình dạng

Proto XLA ShapeProto (xla_data.proto) mô tả số lượng phương diện, kích thước và loại dữ liệu của một mảng N chiều (gọi tắt là mảng).

Thuật ngữ, ký hiệu và quy ước

Số lượng thứ nguyên thực của một mảng là số lượng thứ nguyên có kích thước lớn hơn 1.
Các phương diện được đánh số từ 0 đến N-1 cho một mảng N phương diện. Kích thước của một phương diện là một số nguyên không âm. Cụ thể, kích thước 0 là hợp lệ. Các số đo là nhãn tuỳ ý để thuận tiện. Thứ tự của các số thứ nguyên này không ngụ ý một thứ tự chính/phụ cụ thể trong bố cục của hình dạng. Bố cục được xác định bằng giao thức LayoutProto.
Theo quy ước, các phương diện được liệt kê theo thứ tự tăng dần của số phương diện. Ví dụ: đối với một mảng 3 chiều có kích thước [A x B x C], chiều 0 có kích thước A, chiều 1 có kích thước B và chiều 2 có kích thước C.

Một số tiện ích trong XLA cũng hỗ trợ lập chỉ mục âm tương tự như Python: Phương diện -1 là phương diện cuối cùng (tương đương với N-1 cho một mảng có N phương diện). Ví dụ: đối với mảng 3 chiều được mô tả ở trên, chiều -1 có kích thước C, chiều -2 có kích thước B, v.v.
Mảng hai, ba và bốn chiều thường có các chữ cái cụ thể liên kết với các chiều. Ví dụ: đối với mảng 2 chiều:
- phương diện 0: y
- phương diện 1: x
Đối với mảng 3D:
- phương diện 0: z
- phương diện 1: y
- phương diện 2: x
Đối với mảng 4 chiều:
- phương diện 0: p
- phương diện 1: z
- phương diện 2: y
- phương diện 3: x
Các hàm trong API XLA lấy các phương diện theo thứ tự tăng dần của số phương diện. Thứ tự này khớp với thứ tự được dùng khi truyền các phương diện dưới dạng initializer_list; ví dụ:

ShapeUtil::MakeShape(F32, {A, B, C, D})

sẽ tạo ra một hình dạng có mảng kích thước phương diện bao gồm chuỗi [A, B, C, D].

Bố cục

Proto LayoutProto mô tả cách một mảng được biểu thị trong bộ nhớ. Tệp này bao gồm các trường sau:

message LayoutProto {
  repeated int64 minor_to_major;
  int64 tail_padding_alignment_in_elements;
  ...
}

Thứ tự phương diện phụ đến phương diện chính

Trường bắt buộc duy nhất là minor_to_major. Trường này mô tả thứ tự từ nhỏ đến lớn của các phương diện trong một hình dạng. Các giá trị trong minor_to_major là thứ tự của các phương diện trong mảng (0 đến N-1 đối với mảng N phương diện), trong đó giá trị đầu tiên là phương diện nhỏ nhất cho đến giá trị cuối cùng là phương diện lớn nhất. Phương diện nhỏ nhất là phương diện thay đổi nhanh nhất khi bạn chuyển qua các phần tử của mảng được bố trí trong bộ nhớ tuyến tính.

Ví dụ: hãy xem xét mảng 2 chiều có kích thước [2 x 3] sau đây:

a b c
d e f

Ở đây, phương diện 0 là kích thước 2 và phương diện 1 là kích thước 3. Nếu trường minor_to_major trong bố cục là [0, 1] thì phương diện 0 là phương diện phụ nhất và phương diện 1 là phương diện chính nhất. Điều này tương ứng với bố cục sau trong bộ nhớ tuyến tính:

a d b e c f

Thứ tự chiều từ nhỏ đến lớn này của 0 lên đến N-1 tương tự như cột chính (đối với 2 chiều). Giả sử có một thứ tự đơn điệu của các phương diện, một cách khác mà chúng ta có thể tham chiếu đến bố cục này trong mã là chỉ cần "phương diện 0 là phương diện phụ".

Mặt khác, nếu trường minor_to_major trong bố cục là [1, 0] thì bố cục trong bộ nhớ tuyến tính sẽ là:

a b c d e f

Thứ tự thứ nguyên phụ đến thứ nguyên chính từ N-1 xuống 0 cho một mảng N chiều tương tự như hàng chính (đối với 2 chiều). Giả sử thứ tự kích thước đơn điệu, một cách khác mà chúng ta có thể tham chiếu đến bố cục này trong mã là chỉ cần "dim 0 là chính".

Thứ tự mặc định từ phiên bản phụ đến phiên bản chính

Bố cục mặc định cho các Hình dạng mới tạo là "thứ tự phương diện là chính đến phụ" (tức là [N-1, ..., 0]).

Khoảng đệm

Trường tail_padding_alignment_in_elements xác định cách căn chỉnh mảng được lát thành ô theo số lượng phần tử. Sau khi áp dụng tính năng xếp lát, các phần tử có khoảng đệm sẽ được thêm vào cuối bố cục cho đến khi tổng số phần tử là bội số của giá trị này.

Lập chỉ mục vào mảng

Lớp IndexUtil trong index_util.h cung cấp các tiện ích để chuyển đổi giữa chỉ mục nhiều chiều và chỉ mục tuyến tính cho trước một hình dạng và bố cục. Chỉ mục đa chiều bao gồm một chỉ mục int64 cho mỗi chiều. Chỉ mục tuyến tính là một giá trị int64 duy nhất, lập chỉ mục vào vùng đệm chứa mảng. Xem shape_util.h và layout_util.h trong cùng một thư mục để biết các tiện ích giúp đơn giản hoá việc tạo và thao tác với các hình dạng và bố cục.

Giá trị nhận dạng không gian bộ nhớ

Trong HLO, mỗi mảng có thể được chú thích bằng một giá trị nhận dạng không gian bộ nhớ, được viết dưới dạng S(n).

S(0) (thường bị bỏ qua) biểu thị bộ nhớ băng thông cao (HBM) của thiết bị.
S(1) biểu thị bộ nhớ ảo (VMEM) trên thiết bị.
S(2), S(3), v.v. tương ứng với các không gian bộ nhớ dành riêng cho thiết bị.
S(5) cho biết bộ nhớ máy chủ.