Trang này được dịch bởi Cloud Translation API.
Switch to English

Bố cục lát gạch

Hình 1

Hình 1 cho thấy cách một mảng F32 [3,5] được bố trí trong bộ nhớ với cách xếp chồng 2x2. Hình dạng có bố cục này được viết là F32 [3,5] {1,0: T (2,2)}, trong đó 1,0 liên quan đến thứ tự vật lý của kích thước (trường nhỏ_to_major trong Bố cục) trong khi (2,2) sau dấu hai chấm cho biết lát gạch kích thước vật lý bằng một ô 2x2.

Các ô được đặt một cách trực quan để bao phủ hình dạng và sau đó trong mỗi ô, các phần tử sau đó được đặt mà không cần lát, như trong ví dụ trên, trong đó phần bên phải của ví dụ hiển thị bố cục trong bộ nhớ, bao gồm các phần tử đệm màu trắng được thêm vào để có các ô 2x2 hoàn chỉnh mặc dù các giới hạn của mảng ban đầu không đều.

Các phần tử phụ trong phần đệm không bắt buộc phải chứa bất kỳ giá trị cụ thể nào.

Công thức chỉ số tuyến tính để lát gạch cho một hình dạng và một viên gạch

Không có ô, một phần tử e = (e n , e n-1 , ..., e 1 ) trong một mảng có giới hạn mảng d = (d n , d n-1 , ..., d 1 ) (d1 là kích thước nhỏ nhất) được sắp xếp theo thứ tự lớn đến nhỏ tại vị trí:

linear_index (e, d)
= linear_index ((e n , e n-1 , ..., e 1 ), (d n , d n-1 , ..., d 1 ))
= e n d n-1 ... d 1 + e n-1 d n-2 ... d 1 + ... + e 1

Để đơn giản hóa ký hiệu trong tài liệu này, chúng tôi giả sử một ô có cùng số kích thước với mảng. Trong việc triển khai lát gạch của XLA, điều này được khái quát hóa thành các loại gạch lát có ít kích thước hơn bằng cách giữ nguyên các kích thước chính nhất ban đầu không thay đổi và chỉ áp dụng việc ốp lát cho các kích thước nhỏ nhất, để việc lát gạch được chỉ định đề cập đến hậu tố của các kích thước vật lý của hình dạng được lát gạch.

Khi lát gạch có kích thước (t n , t n-1 , ..., t 1 ) được sử dụng, một phần tử trong mảng có các chỉ số (e n , e n-1 , ..., e 1 ) được ánh xạ tới vị trí trong bố cục cuối cùng:

linear_index_with_tile (e, d, t)
= linear_index ((⌊e / t⌋, e mod t), (⌈d / t⌉, t)) (số học là từng phần tử, (a, b) là nối)
= linear_index ((e n / t n ⌋, ..., ⌊e 1 / t 1 ⌋, e n mod t n , ..., e 1 mod t 1 ), (⌈d n / t n ⌉, ..., ⌈d 1 / t 1 ⌉, t n , t n-1 , ..., t 1 ))
= linear_index ((⌊e n / t n ⌋, ..., ⌊e 1 / t 1 ⌋), (⌈d n / t n ⌉, ..., ⌈d 1 / t 1 ⌉)) ∙ t n t n-1 ... t 1 + linear_index ((e n mod t n , ..., e 1 mod t 1 ), (t n , t n-1 , ..., t 1 ))

Bố cục có thể được coi là có hai phần: (⌊e n / t n ⌋, ..., ⌊e 1 / t 1 ⌋), tương ứng với một chỉ số ô trong một mảng ô có kích thước (⌈d n / t n ⌉, ..., ⌈d 1 / t 1 ⌉), và (e n mod t n , ..., e 1 mod t 1 ), tương ứng với một chỉ mục trong ô. Hàm ceil xuất hiện trong ⌈d i / t i ⌉ bởi vì nếu các ô gạch vượt quá giới hạn của mảng lớn hơn, thì phần đệm sẽ được chèn như trong Hình 1. Cả ô và phần tử trong ô đều được đặt một cách đệ quy mà không cần xếp.

Ví dụ trong Hình 1, phần tử (2,3) có chỉ số ô (1,1) và chỉ số bên trong ô (0,1), cho một vectơ tọa độ kết hợp của (1, 1, 0, 1). Các chỉ số ô có giới hạn (2, 3) và ô chính nó là (2, 2) cho một vectơ kết hợp của (2, 3, 2, 2). Chỉ số tuyến tính với ô cho phần tử có chỉ số (2, 3) trong hình dạng logic khi đó

linear_index_with_tile ((2,3), (3,5), (2,2))
= linear_index ((1,1,0,1), (2,3,2,2))
= linear_index ((1,1), (2,3)) ∙ 2 ∙ 2 + linear_index ((0,1), (2,2))
= (1 ∙ 3 + 1) ∙ 2 ∙ 2 + (0 ∙ 2 + 1)
= 17.

Lát gạch dưới dạng pad-reshape-transpose

Bố cục dựa trên lát gạch hoạt động như sau:
Xét một mảng kích thước (d n , d n-1 , ..., d1) (d1 là kích thước nhỏ nhất). Khi nó được lát với kích thước lát (t n , t n-1 , ..., t 1 ) (t 1 là kích thước nhỏ nhất), việc lát gạch đó có thể được mô tả dưới dạng pad-reshape-transpose như sau đường.

  1. Mảng được đệm thành (⌈d n / t n ⌉ ∙ t n , ..., ⌈d 1 / t 1 ⌉ ∙ t 1 ).
  2. Mỗi thứ nguyên i được chia thành (⌈d i / t i ⌉, t i ), tức là mảng được định hình lại thành
    (⌈D n / t n ⌉, t n , ..., ⌈d 1 / t 1 ⌉, t 1 ).
    Không có sự thay đổi bố cục vật lý nào trong việc định hình lại này, vì vậy định hình lại này là một bitcast. Nếu người ta không nghĩ rõ ràng về cách lát gạch, thì việc định hình lại này có thể thể hiện bất kỳ hình dạng nào có cùng số phần tử với hình dạng được đệm - ví dụ ở đây là cách thể hiện một viên gạch theo cách này.
  3. Chuyển vị xảy ra bằng cách di chuyển t n , ..., t 1 đến các kích thước nhỏ nhất trong khi vẫn giữ thứ tự tương đối của chúng, để thứ tự của các kích thước từ kích thước lớn nhất đến kích thước nhỏ nhất trở thành
    (⌈D n / t n ⌉, ..., ⌈d 1 / t 1 ⌉, t n , ..., t 1 ).

Hình dạng cuối cùng có tiền tố
(⌈D n / t n ⌉, ..., ⌈d 1 / t 1 ⌉), mô tả số lượng ô trong mỗi chiều. Một phần tử trong mảng (e n , ..., e 1 ) được ánh xạ tới phần tử này trong hình dạng cuối cùng:
(E n / t n ⌋, ..., ⌊e 0 / t 0 ⌋, e n mod t n , ..., e 1 mod t 1 ). Dễ dàng nhận thấy rằng chỉ số tuyến tính của phần tử tuân theo công thức trên như mong đợi.

Lát gạch nhiều lần

Việc ốp lát của XLA thậm chí còn trở nên linh hoạt hơn bằng cách áp dụng nó nhiều lần.

Hình 2

Hình 2 cho thấy cách một mảng kích thước 4x8 được lát bằng hai cấp độ (đầu tiên là 2x4 sau đó là 2x1). Chúng tôi đại diện cho việc lát gạch lặp lại này là (2,4) (2,1). Mỗi màu chỉ ra một ô 2x4 và mỗi ô viền đỏ là một ô 2x1. Các con số cho biết chỉ số tuyến tính trong bộ nhớ của phần tử đó ở định dạng lát gạch. Định dạng này khớp với định dạng được sử dụng cho BF16 trên TPU, ngoại trừ ô ban đầu lớn hơn, cụ thể là ô xếp là (8,128) (2,1), trong đó mục đích của ô xếp thứ hai x 2x1 là thu thập hai giá trị 16 bit lại với nhau để tạo thành một giá trị 32 bit theo cách phù hợp với kiến ​​trúc của TPU.

Lưu ý rằng ô thứ hai trở lên có thể tham chiếu đến cả các kích thước nhỏ bên trong ô, chỉ sắp xếp lại dữ liệu trong ô, như trong ví dụ này với (8,128) (2,1), nhưng cũng có thể tham chiếu đến ô chéo chính kích thước từ lần lát gạch trước.

Kết hợp các kích thước bằng cách sử dụng gạch

XLA's cũng hỗ trợ kết hợp các kích thước. Ví dụ: nó có thể kết hợp các thứ nguyên trong F32 [2,7,8,11,10] {4,3,2,1,0} thành F32 [112,110] {1,0} trước khi xếp nó với (2,3 ). Ngói được sử dụng là (∗, ∗, 2, ∗, 3). Ở đây, dấu hoa thị trong ô có nghĩa là lấy thứ nguyên đó và kết hợp nó với thứ nguyên nhỏ hơn tiếp theo. Nhiều thứ nguyên liền kề có thể được gộp lại với nhau thành một thứ nguyên. Kích thước cộng gộp được biểu thị bằng giá trị ô -1 trong kích thước đó của ô, giá trị này không hợp lệ trong ô dưới dạng kích thước thứ nguyên.

Chính xác hơn, nếu kích thước i của hình dạng bị loại bỏ thông qua dấu hoa thị trong ô, thì trước khi áp dụng định nghĩa trước đó về lát gạch, kích thước đó sẽ bị xóa khỏi cả hình dạng được lát và vectơ lát và thứ nguyên i-1 là gì của hình dạng có giới hạn mảng của nó tăng từ d i-1 lên d i d i-1 . Bước này được lặp lại cho mỗi dấu hoa thị trong vector ô xếp.