Tài liệu này trình bày hành trình của một mô-đun Trình tối ưu hoá cấp cao (HLO) XLA từ trạng thái ban đầu đến một tệp thực thi cuối cùng. Đôi khi, chúng tôi sẽ bỏ qua "mô-đun" và chỉ gọi là "HLO".
HLO tối ưu hoá trước
Chúng ta bắt đầu với mô-đun HLO trước khi tối ưu hoá. HLO trước khi tối ưu hoá không chứa các thao tác (ops) được coi là nội bộ đối với XLA, chẳng hạn như fusion hoặc bitcast. Các thao tác không có bố cục ở giai đoạn này hoặc nếu có, bố cục đó sẽ bị bỏ qua. HLO trước khi tối ưu hoá thường được tạo bởi các khung cấp cao hơn như TensorFlow và JAX. Khi sử dụng cờ XLA -xla_dump_to, HLO trước khi tối ưu hoá sẽ được kết xuất vào một tệp có hậu tố tên tệp là "before_optimizations.txt".
Tối ưu hoá mô-đun HLO
Quy trình XLA:GPU chuyển HLO trước khi tối ưu hoá thành HLO được tối ưu hoá bằng cách chạy một chuỗi các lượt truyền. Các lượt truyền có thể được nhóm lại với nhau theo ngữ nghĩa và chạy theo thứ tự sau:
Phân đoạn các thẻ và vé có liên quan
Điều này bao gồm các đường truyền như Shardy Partitioner hoặc các đường truyền cho việc phân mảnh SPMD.
Các lượt tối ưu hoá
Điều này có thể bao gồm cả các lần hợp pháp hoá và đơn giản hoá.
Các lượt tối ưu hoá tập thể
Tương tự như Optimization passes (Các lượt tối ưu hoá), nhưng tập trung vào các hoạt động tập thể.
Lượt gán bố cục
Mỗi thao tác HLO được chỉ định một bố cục là một phần của hình dạng hướng dẫn. Bố cục kiểm soát cách tensor được bố trí thực tế trong bộ nhớ.
Ví dụ về một hình dạng có bố cục:
f32[10,20,30]{2,0,1}
Sau loại phần tử, sẽ có các phương diện logic của hình dạng, theo sau là hoán vị bố cục theo thứ tự từ nhỏ đến lớn. Trong ví dụ này, phương diện nhỏ nhất là 30, phương diện nhỏ thứ hai là 10 và phương diện chính là 20.
Mục tiêu của việc chỉ định bố cục là giảm thiểu số lượng các phép chuyển vị vật lý bắt buộc bằng cách sử dụng chiến lược tham lam. Quá trình này bắt đầu bằng một số ràng buộc về bố cục nhất định (ví dụ: các thư viện cuDNN/cuBLAS mong đợi các phương diện liên tiếp) và truyền bố cục "xuống" rồi "lên" biểu đồ HLO. Vào cuối quá trình truyền bố cục, một số hướng dẫn có thể có bố cục xung đột, một bố cục được truyền từ một toán hạng, một bố cục được truyền từ người dùng. Để giải quyết xung đột này, một chỉ dẫn HLO copy sẽ được chèn vào để thay đổi bố cục từ bố cục toán hạng thành bố cục chỉ dẫn.
Các lượt chuẩn hoá bố cục
Vì khó xác định hình dạng thực tế, nên quá trình chuẩn hoá bố cục sẽ cố gắng viết lại hình dạng để sử dụng bố cục mặc định {rank-1, rank-2, …, 0}. Trong ví dụ trên, hình dạng được chuẩn hoá sẽ là f32[20,10,30]{2,1,0}. Các thao tác sao chép làm thay đổi bố cục sẽ được viết lại dưới dạng tổ hợp của transpose và bitcast. Vì hiện tại chúng ta không thể chuẩn hoá tất cả các hoạt động, nên vẫn có một số hoạt động có thể có bố cục không mặc định, đáng chú ý nhất là gather và dot. Tại ranh giới giữa các hoạt động chuẩn hoá và các hoạt động không chuẩn hoá, sẽ có các hoạt động bitcast đại diện cho một hoạt động chuyển vị, tức là một hoạt động chuyển vị có bố cục được chỉ định khiến hoạt động này không có tác dụng về mặt vật lý.
Việc chuẩn hoá bố cục cũng giúp một số phép chuyển vị ngầm định trở nên rõ ràng. Điều này rất quan trọng vì codegen có thể xử lý các phép chuyển vị rõ ràng bằng một trình phát chuyên dụng. Ví dụ: về mặt kỹ thuật, thao tác định hình lại được phép có bố cục vật lý khác nhau giữa toán hạng và kết quả (ví dụ: do thứ hạng khác nhau). Lượt truyền ReshapeDecomposer chạy trong các lượt truyền chuẩn hoá bố cục sẽ chuyển đổi một thao tác định hình lại thành một chuỗi transpose, định hình lại bitcast và transpose.
Đăng các lượt tối ưu hoá việc chỉ định bố cục
Các lượt truyền quan trọng nhất ở đây là các lượt truyền kết hợp Triton (các lượt truyền kết hợp GEMM + các lượt truyền kết hợp Softmax/Layernorm) hoặc các lượt truyền viết lại thành lệnh gọi thư viện. Tính năng tự động điều chỉnh cũng chạy trong bước này, trong đó XLA chọn giữa các trình phát khác nhau, chọn thuật toán tốt nhất cho các tích chập hoặc dấu chấm, tìm ra cách phân ô tốt nhất cho các hợp nhất do trình phát Triton xử lý, v.v.
Thẻ và vé Fusion
Hai đường chuyền chính là PriorityFusion và Multi-Output hợp nhất.
Trong PriorityFusion, chúng ta tạo ra các hợp nhất được hướng dẫn bởi mô hình chi phí. Khi hợp nhất, chúng tôi sẽ cho phép sao chép các thao tác với nhiều người dùng nếu thao tác đó có thể được hợp nhất vào tất cả người dùng. Chúng tôi cũng sẽ cho phép mở rộng các hợp nhất Triton Softmax hiện có nếu có thể.
Hợp nhất Multi-Output là một lượt riêng biệt cho phép hợp nhất các thao tác/hợp nhất dùng chung một toán hạng. Thao tác này cũng có thể hợp nhất các toán hạng/hợp nhất toán hạng thành người dùng mà không bị trùng lặp bằng cách thêm(các) đầu ra bổ sung, để những người dùng khác của thao tác cần hợp nhất có thể được chuyển hướng đến các đầu ra này. Lượt truyền này cần cẩn thận để không đưa các chu kỳ vào biểu đồ HLO.
Sau khi hợp nhất nhiều đầu ra, quá trình loại bỏ biểu thức con chung (truyền HloCSE) sẽ chạy, có khả năng hợp nhất các thao tác trùng lặp trước đó lại với nhau nếu chúng kết thúc trong cùng một quá trình hợp nhất.
Một số lượt truyền sau khi hợp nhất
Một số lượt truyền liên quan đến các tập hợp (chẳng hạn như chuyển các tập hợp thành không đồng bộ hoặc thực thi một thứ tự tương đối nhất định của các tập hợp).
Cuối cùng, chúng ta chạy CopyInsertion, trong đó các bản sao được thêm vào để đảm bảo rằng các thao tác tại chỗ không ghi đè dữ liệu vẫn cần thiết ở nơi khác.
Vào cuối quá trình tối ưu hoá, HLO được tối ưu hoá sẽ được kết xuất nếu sử dụng cờ -xla_dump_to vào một tệp có hậu tố tên tệp "after_optimizations.txt". Nếu muốn kết xuất HLO sau các lượt truyền trung gian thực sự thay đổi HloModule, bạn có thể sử dụng cờ -xla_dump_hlo_pass_re=.* (hoặc một biểu thức chính quy cụ thể để hạn chế HLO ở một số lượt truyền nhất định).
Lập lịch
Một Mô-đun HLO không có lịch biểu vẫn có một mức độ tự do nhất định trong thứ tự xử lý các hoạt động. Mọi sắp xếp theo cấu trúc liên kết tôn trọng mối quan hệ toán hạng/kết quả và các phần phụ thuộc kiểm soát đều hợp lệ. Lập lịch xác định thứ tự cụ thể cần sử dụng. Mối lo ngại chính ở giai đoạn này là mức tiêu thụ bộ nhớ tối đa, tuỳ thuộc vào thời gian tồn tại của các tensor. Trong bước ban đầu, chúng tôi thử các thuật toán lập lịch khác nhau và chọn lịch biểu giúp giảm thiểu mức tiêu thụ bộ nhớ cao nhất. Xin lưu ý rằng tại thời điểm này, chúng ta chưa làm việc với bộ đệm vật lý (điều này sẽ xảy ra trong "Chỉ định bộ đệm") và mô phỏng mức sử dụng bộ nhớ.
Sau đó, LatencyHidingScheduler sẽ chạy và cố gắng tối đa hoá mức độ trùng lặp giữa hoạt động tính toán và giao tiếp. Nhưng điều đó có thể làm tăng mức sử dụng bộ nhớ trở lại.
Cuối cùng, trong trường hợp mức tiêu thụ bộ nhớ cao hơn lượng bộ nhớ mà chúng ta có, chúng ta sẽ chạy HloRematerialization. Lượt truyền này cố gắng giảm mức sử dụng bộ nhớ bằng cách giảm hiệu suất, chẳng hạn như một số hoạt động hợp nhất có thể được chia tách và một số hoạt động có thể được sao chép để có thời gian tồn tại của vùng đệm ngắn hơn. Nếu quá trình kết xuất lại xảy ra, bạn nên tìm cách giảm yêu cầu về bộ nhớ ở phía mô hình (ví dụ: sử dụng kích thước lô nhỏ hơn).
Chỉ định vùng đệm
Ngay trước khi giảm xuống LLVM IR, chúng ta sẽ chạy các lượt gán vùng đệm để gán các lát vùng đệm cho từng chỉ dẫn trong biểu đồ HLO. Việc chỉ định vùng đệm diễn ra theo một số bước:
HloDataflowAnalysischỉ địnhHloValues(về cơ bản là các vùng đệm logic) cho các chỉ dẫn. Đối với các thao tác tại chỗ, bạn có thể sử dụng lạiHloValuecủa một toán hạng. Một op có thể xác định nhiềuHloValue(ví dụ: với hình dạng kết quả là một bộ).HloAliasAnalysiscố gắng kết hợp các vùng đệm cho các thao tác tạo hiệu ứng răng cưa và tính toán một mối ánh xạ từHloValueđếnHloBuffer.BufferAssignmenttính toán một mối liên kết từHloBuffersđến các phần vùng đệm bên trong một vùng đệm lớn theo cách sao cho cùng một phần vùng đệm không được dùng cho cácHloBufferskhác nhau có thời gian tồn tại trùng nhau. Đối với các thao tác có thể tạo bí danh, bạn có thể thấy có một chút trùng lặp (thời gian kết thúc của mộtHloBuffercó thể trùng với thời gian bắt đầu củaHloBufferkhác). Khi sử dụng cờ-xla_dump_to, một số thông tin về việc chỉ định vùng đệm sẽ được kết xuất vào một tệp có hậu tố tên là "after_optimizations-buffer-assignment.txt".
Thunk
Sau khi được tối ưu hoá và lập lịch, biểu đồ HLO sẽ được giảm xuống thành một chuỗi thunk tuyến tính cho một phần phụ trợ cụ thể (CPU hoặc GPU).
Trong XLA, Thunk là một đơn vị công việc khép kín mà thời gian chạy thực thi. Đó có thể là một lần khởi chạy hạt nhân đã biên dịch, thao tác cụ thể, lệnh gọi thư viện, cấu trúc luồng điều khiển, giao tiếp tập thể, v.v. Trình tự Thunk đại diện cho toàn bộ tệp thực thi cho một phần phụ trợ cụ thể.
Thunk Emission
Quá trình chuyển đổi một phép tính HLO theo lịch biểu thành một chuỗi thunk được gọi là "phát thunk". Việc này do một lớp trình phát chuyên dụng xử lý trong mỗi chương trình phụ trợ.
Đối với phần phụ trợ GPU, việc này do IrEmitterUnnested xử lý.
EmitHloComputation lặp lại danh sách Lệnh HLO theo lịch trong một phép tính và gửi đến một phương thức Emit... chuyên biệt (ví dụ:
EmitFusion, EmitConvolutionThunk, EmitWhile). Mỗi phương thức này sẽ tạo(các) đối tượng Thunk thích hợp và nối chúng vào chuỗi thunk.
Đối với CPU Backend, ThunkEmitter sẽ thực hiện vai trò này và được sắp xếp theo cách tương tự. ThunkSequence cuối cùng được nhúng trong CpuExecutable.
Xin lưu ý rằng mỗi chỉ dẫn trong quá trình tính toán mục nhập của một mô-đun HLO có thể tương ứng với không (kTuple, kConstant, ..), một hoặc nhiều (ví dụ: chỉ dẫn sắp xếp) thunk trong chuỗi thunk cuối cùng.
Bộ đệm lệnh: Tối ưu hoá việc thực thi trên GPU
Phần cứng GPU hiện đại cho phép ghi lại một chuỗi các hoạt động của GPU (khởi chạy kernel, sao chép bộ nhớ, v.v.) một lần rồi phát lại chuỗi đó nhiều lần với mức hao tổn CPU tối thiểu. Đây là một hoạt động tối ưu hoá hiệu suất quan trọng, đặc biệt là đối với những khối lượng công việc có nhiều nhân nhỏ, khởi động nhanh. XLA sử dụng Command Buffer làm một lớp trừu tượng của CUDA Graphs hoặc HIP Graphs. Giao diện cốt lõi được xác định trong GpuCommandBuffer.
Vùng đệm lệnh được biểu thị trong một chuỗi thunk bằng CommandBufferThunk.
Trình phát không tạo thunk này trực tiếp từ các chỉ dẫn HLO. Thay vào đó, việc này được thực hiện bằng CommandBufferConversionPass chạy trên chính ThunkSequence.
Lượt truyền này xác định các chuỗi con liền kề của các thunk tương thích (ví dụ: một chuỗi KernelThunk và GemmThunk). Sau đó, nó sẽ thay thế chuỗi con tìm thấy bằng một CommandBufferThunk duy nhất. Thunk mới đóng gói logic của các thunk ban đầu dưới dạng danh sách các đối tượng CommandBufferCmd đơn giản.
Khi CommandBufferThunk thực thi lần đầu tiên trên một luồng GPU nhất định, nó sẽ "ghi" chuỗi lệnh vào một bộ đệm lệnh phần cứng. Trong tất cả các lần thực thi tiếp theo, lệnh này chỉ đơn giản là đưa ra một lệnh duy nhất cho GPU để "phát lại" chuỗi đã ghi. Điều này giúp tránh hao tổn CPU khi khởi chạy từng nhân riêng lẻ.
Có thể thực thi
Sản phẩm cuối cùng của quy trình biên dịch XLA là một Tệp thực thi độc lập, dành riêng cho nền tảng. Đối tượng này đóng gói tất cả thông tin cần thiết để chạy chương trình đã biên dịch trên một thiết bị mục tiêu, chẳng hạn như CPU hoặc GPU. Đây là cầu nối giữa trình biên dịch và thời gian chạy. Các thời gian chạy hiện đại như PJRT sử dụng các trừu tượng cấp cao hơn một chút (xem PjRtExecutable), nhưng cuối cùng, những trừu tượng này sẽ bao bọc một tệp thực thi dành riêng cho phần phụ trợ.
Executable chứa một số thông tin chính được tạo trong quá trình biên dịch. Mặc dù nội dung chính xác sẽ khác nhau tuỳ theo hệ thống phụ trợ, nhưng thường bao gồm:
Mã đã biên dịch: Đây là mã máy cấp thấp sẽ chạy trên thiết bị. Đối với CPU, đây thường là một hoặc nhiều tệp đối tượng. Đối với GPU, đây là mã thiết bị đã biên dịch ở định dạng PTX hoặc HSACO, được tải vào GPU trong thời gian chạy.
Kế hoạch thực thi (ThunkSequence): Cốt lõi của logic thời gian chạy. Đây là một chuỗi tuyến tính gồm các đối tượng Thunk. Mỗi thunk đại diện cho một đơn vị công việc, chẳng hạn như khởi chạy một hạt nhân, gọi một hàm thư viện (ví dụ: cuBLAS) hoặc xử lý luồng điều khiển. Thời gian chạy sẽ thực thi chương trình bằng cách lặp lại trình tự này.
Bố cục bộ nhớ (BufferAssignment): Phần siêu dữ liệu quan trọng này do BufferAssigner tạo ra, mô tả bố cục bộ nhớ hoàn chỉnh cho quá trình tính toán. Nó chỉ định kích thước của mọi vùng đệm và cách phân bổ cũng như sử dụng lại bộ nhớ cho các tham số, đầu ra và giá trị tạm thời. Thời gian chạy sử dụng thông tin này để phân bổ bộ nhớ thiết bị và truyền các con trỏ chính xác đến từng thunk.
(không bắt buộc) Mô-đun HLO: Để gỡ lỗi và lập hồ sơ, tệp thực thi thường giữ lại một tham chiếu đến HloModule cuối cùng, được tối ưu hoá mà tệp đó được biên dịch.
Trình biên dịch sẽ điều phối việc tạo tệp thực thi cuối cùng cho từng phần phụ trợ cụ thể. Phương thức RunBackend của một quy trình triển khai Trình biên dịch là bước cuối cùng trong quy trình biên dịch, đóng gói tất cả các cấu phần phần mềm đã biên dịch thành một đối tượng Executable.
GpuCompiler và CpuCompiler lần lượt nhắm đến GPU và CPU.
Khi người dùng gọi Execute... trên một tệp thực thi, thời gian chạy sẽ dùng BufferAssignment để phân bổ bộ nhớ, rồi gọi ThunkSequence để khởi chạy các thao tác trên thiết bị bằng mã đã biên dịch.