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

Mã MLIRGen cho XLA

XLA hoạt động trên HloInstruction và thực hiện nhiều tối ưu hóa trên đại diện này, chia sẻ rất nhiều trong số này giữa các thiết bị được nhắm mục tiêu. Vì một số điểm, lịch biểu tuyến tính được tính toán và bộ nhớ đệm được gán cho từng giá trị tĩnh. Bộ mã cụ thể của thiết bị hoạt động bằng cách duyệt qua chuỗi này và gọi "bộ phát" để tạo ra một đại diện phù hợp với thiết bị (ví dụ: một hàm LLVM duy nhất cho mỗi tính toán XLA trên CPU hoặc một chuỗi "thunks" đóng gói các hoạt động GPU và có thể tạo ra PTX khi GPU nhắm mục tiêu).

Là một bước dàn dựng, chúng tôi hiện đang trong quá trình chặn quy trình ngay sau khi XLA hoàn thành giai đoạn gán bộ đệm và thay vào đó phát ra một mô-đun lhlo phương ngữ lhlo . Từ đó, chúng tôi thực hiện codegen bằng cách sử dụng các thành phần MLIR (phương ngữ Linalg, affine và GPU) tùy thuộc vào thiết bị.

Dưới đây là kế hoạch ghi để di chuyển dần dần XLA / GPU bằng cách sử dụng lhlo làm đầu vào codegen.

Nhiệm vụ

Tổ chức Thiết bị
định dạng đầu vào HloInemony * (Nhiệm vụ 1) HloInemony * (Nhiệm vụ 1)
Định dạng đầu ra xla :: Thunk (Nhiệm vụ 2) LLVM IR (Nhiệm vụ 3)
  • Nhiệm vụ 1 thay đổi cả định dạng đầu vào của máy chủ và thiết bị từ HloIn cản * thành LHLO.
  • Nhiệm vụ 2 thay đổi định dạng đầu ra của máy chủ từ thunks thành "một số bãi đáp cho máy chủ" (xem bên dưới).
  • Nhiệm vụ 3 di chuyển đầu ra thiết bị từ LLVM IR sang một số dạng MLIR. Đây là tùy chọn cho dự án này và xem phần "Di chuyển thiết bị LLVM IR" để biết chi tiết.

Dự án này ưu tiên có các mô hình có thể chạy được từ đầu đến cuối với các trình phát LHLO được kích hoạt càng nhiều càng tốt. Điều này ngụ ý rằng danh sách các mục tiêu theo thứ tự ưu tiên sau:

  • Làm cho XLA / GPU có thể chạy được với các trình phát LHLO, với các Thunks và trình phát hiện tại không được sửa đổi.
  • Loại bỏ các tham chiếu đến HloIn cản * trong LHLO, theo từng trường hợp:
    • Chuyển đổi một bộ phát kế thừa sang một bộ phát dựa trên MLIR (ví dụ: Linalg) hoặc
    • Dịch một cách cơ học bộ phát hiện có để lấy đại diện MLIR (di chuyển sang Tiêu chuẩn với Phương ngữ GPU).

Di chuyển Thunks (Nhiệm vụ 2)

xla :: gpu :: Thunk là cấu trúc dữ liệu:

  • Có thể được gọi vào từ máy chủ (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • Mang dữ liệu khác nhau trong các lớp con của nó.
  • Tương tác với Buffer Allocation :: Slice và StreamExecutor.
  • Ra mắt hạt nhân
  • Gọi vào tất cả các thư viện thời gian chạy.

Chi phí đó bao gồm:

  • Đại diện cho dữ liệu cấu hình op cụ thể (ví dụ: cấu hình chập).
  • Di chuyển hình dạng op và hình dạng toán hạng.
  • Đại diện cho một cây thunks (trong khi, điều kiện, vv).

Công việc di chuyển độc lập với di chuyển LHLO / emitter. Trong các nguồn lực hạn chế, nó được ưu tiên sau di chuyển LHLO / emitter.

Chúng tôi có một số lựa chọn về cách hạ thấp phần phía máy chủ từ LHLO:

  • TFRT
    • (Pro) trình bao bọc CUDA và HIP tuyệt vời để sử dụng.
    • (Pro) dễ thực hiện các cuộc gọi thư viện (cuDNN, cuBLAS, cuFFT, v.v.), vì các op TFRT được diễn giải bằng mã C ++.
    • (Con) phía máy chủ đang được phát triển và không được thử nghiệm.
  • Mã CPU bị loại
    • (Pro) khả năng thấp hơn tuyệt vời. Tạo một vài vòng lặp và điều kiện và thế là xong.
    • (Con) GPUDialect chưa mô hình chuỗi / luồng / không đồng bộ / phân bổ thiết bị.
    • (Con) Hỗ trợ thời gian chạy CUDA / HIP là tối thiểu (đường dẫn bộ công cụ, phiên bản, tải động, v.v.).
  • Hiện tại (phiên dịch) XLA

Quyết định: áp dụng TFRT, nhưng cũng hỗ trợ trích xuất mã CPU trong TFRT.

Thiết bị di chuyển LLVM IR (Nhiệm vụ 3)

Một bộ phát nguyên tố tạo ra op mục tiêu bằng cách điền nó vào từng phần tử. Mỗi phần tử đầu ra phụ thuộc vào một tập hợp các phần tử từ toán hạng. Tất cả các yếu tố được mô tả bằng cách kết hợp bộ đệm với các chỉ số động. Nó đủ để mô tả gần như tất cả các op "math", nhưng vì lý do hiệu năng, chỉ một tập hợp lớn các op "math" được triển khai trực tiếp trong (Cpu | Gpu) ElementalIrEuctor.

ElementalIrEuctor là duy nhất ở chỗ:

  • Một phần lớn mã được chia sẻ giữa XLA / GPU và CPU.
  • Nó đại diện cho một phần lớn các op được nhìn thấy trong các mô hình, bao gồm tất cả các op-khôn ngoan.
  • Hầu hết các thân máy bay chỉ phụ thuộc vào ElementalIrEuctor.
  • Nó có cấu trúc đơn giản, vì nó mô tả DAG phụ thuộc dữ liệu giữa các phần tử op và các phần tử toán hạng.
  • Nó chủ yếu là di động và cấp cao (ví dụ, không giống như GPU kReduce và GPU kCopy).
  • Hỗ trợ hình dạng động dễ dàng cho các ops yếu tố khôn ngoan.

Bây giờ, đối với tất cả các op, được phát ra theo nguyên tố hay không, có một số hương vị của trạng thái kết thúc của mỗi op XLA:

  1. Mã thiết bị vẫn là LLVM IR.
  2. Tái cấu trúc bộ phát cũ để giống như LHLO -> MLIR LLVM:
    • (Chi phí) Sẽ là công việc vứt đi nếu chúng ta cuối cùng muốn chuyển sang Tiêu chuẩn.
    • (Lợi ích) Nó dễ dàng và cơ học. Có thể được thực hiện trong một thời gian ngắn.
    • (Lợi ích) Nó không có lợi hơn so với (1).
  3. Tái cấu trúc các trình phát cũ để giống như LHLO -> MLIR GPU + Standard + Loops:
    • (Chi phí) Nâng các trình phát hiện có lên Standard đưa ra một số thách thức. Con trỏ và GEP cần được chuyển đổi thành MemRefs và SubViews. Đảm bảo tính đầy đủ của amdgpu là một số khác.
    • (Chi phí) XLA / GPU phụ thuộc rất nhiều vào siêu dữ liệu LLVM:
      • range cho các chỉ số khối / chủ đề.
      • align , dereferenceable , invariant.load , alias.scope , noalias cho tải / cửa hàng.
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable cho các vòng lặp liên tiếp.
    • (Lợi ích) Có thể lâu dài. Di động hơn.
  4. Tái cấu trúc các trình phát cũ thành LHLO -> Linalg và viết các trình phát Linalg mới
    • (Chi phí) Đây là từng trường hợp. So với các tùy chọn trước đây, việc triển khai mới phù hợp với hiệu suất của XLA cần phải trải qua điểm chuẩn <-> tối ưu hóa quy trình làm việc, có thể là một chi phí đáng kể cho một số ops.
    • (Lợi ích) chồng thống nhất; sự đóng góp cho cộng đồng; tính di động; tiềm năng tối ưu hóa nhiều hơn.

Kết luận:

  • Đừng đi cho (2). (1) hoặc (3) chỉ tốt hơn (2). (2) chi phí nhiều hơn (1), vì nó đòi hỏi rất nhiều tái cấu trúc cơ học. Với (1) chúng ta vẫn có thể đạt được mục tiêu cho phép XLA chọn các trình phát MLIR. Điều này là bằng cách thực hiện LHLO -> LLVM IR -> chạy các trình phát thiết bị cũ.
  • ElementalIrEuctor ops đi cho (4), nhưng không tăng dần. Không có cách nào để làm điều đó bởi op, bởi vì tất cả các op được phát ra từ nguyên tố được kết nối vào cùng một biểu đồ. Công việc này cũng có thể đóng vai trò là điểm thống nhất của một số lực lượng đang hoạt động (xla / service / mlir_gpu, trình tạo nhân, Linalg).
  • Tất cả các op khác đi cho (1). Như một mục tiêu kéo dài, chúng có thể được di chuyển đến (3) hoặc (4).

Ưu tiên

Mặc dù cả ba nhiệm vụ được đề cập ở trên là song song, trong các nguồn lực hạn chế, chúng phải được tuần tự hóa. Ưu tiên tập trung vào kết quả có thể nhìn thấy để hoàn thành mỗi nhiệm vụ.

Ưu tiên là: Nhiệm vụ 1 (LHLO cho các trình phát kế thừa)> Nhiệm vụ 2 (Thunks)> Nhiệm vụ 3 (Trình phát MLIR).

Khi kết thúc Nhiệm vụ 1, người dùng XLA có thể tạo LHLO (ví dụ: trình tạo nhân) và thực thi chúng. Định dạng biên dịch sẽ không được tuần tự MLIR.

Khi kết thúc Nhiệm vụ 2, LHLO hạ xuống MLIR thích hợp, tuần tự hóa. Điều này cho phép biên dịch ngoại tuyến.

Đến cuối Nhiệm vụ 3, tất cả các trình phát XLA đều dựa trên MLIR.

Thiết kế chi tiết

Bước 1: (Nhiệm vụ 1) Hoàn thành LHLO và tạo các trình phát kế thừa

Bước này làm cho tất cả các trình phát XLA / GPU hiện có tương tác với các MLIR op. Bước này là tái cấu trúc thuần túy và NFC.

Bước này chủ yếu là cơ học, nhưng đáng chú ý là sự khác biệt sau đây giữa HloComputing không được kiểm tra và LHLO:

  • Mỗi HloIn cản có quyền truy cập trực tiếp vào toán hạng của nó (DAG luồng dữ liệu). Ngược lại, mỗi op LHLO chỉ có quyền truy cập vào bộ đệm toán hạng của nó (một bipartite giữa op và bộ đệm). Các op của LHLO phải thông qua các chuỗi use-def để truy cập các op toán hạng của chúng.
  • Các trình phát di sản chưa được kiểm chứng theo kinh nghiệm hầu như không bao giờ truy cập vào toán hạng của chúng. Ngoại lệ duy nhất là kReduce.
  • Các trình phát kế thừa không được kiểm tra chỉ truy cập vào BufferAssocation để nhận các lát cắt, không phải để truy cập các cấu trúc dữ liệu bổ trợ như dataflow_analysis () hoặc alias_analysis (). llvm_ir xây dựng alias_analysis () dựa trên thông tin lát.

Kết luận là LHLO phải phù hợp ngay trong mà không gặp rắc rối lớn.

Bước 2: (Tùy chọn) Hỗ trợ hồ sơ

Bước này chỉ cần thiết nếu chúng ta bắt đầu loại bỏ một số logic Thunk XLA (xem bước tiếp theo).

Trước khi thực sự bật bất kỳ trình phát dựa trên MLIR nào, chúng ta cần định hình cho các trình phát dựa trên MLIR.

Hiện tại XLA thực hiện hồ sơ riêng bằng cách gọi vào bộ đếm thời gian của StreamExecutor. Đồng hồ bấm giờ dưới mui xe chèn hai sự kiện trước và sau khi khởi chạy kernel và đo thời gian đồng bộ giữa hai sự kiện này.

Có khoảng ba cách tiếp cận để hỗ trợ hồ sơ trong MLIR:

  • Chạy một profiler end-to-end
  • Thêm một hồ sơ op cho mỗi op trong LHLO, sử dụng một profiler được tiêm.

Cách tiếp cận "từ đầu đến cuối" trong suốt đối với MLIR, nhưng gặp phải vấn đề tương tự khiến XLA không sử dụng nó ở nơi đầu tiên: các cuộc gọi thư viện được thu thập bởi một trình hồ sơ (nvprof / ...) không thể dễ dàng liên quan đến HLO ops Ví dụ, cuDNN khởi chạy nhiều hạt nhân cho mỗi HLO và thật khó để biết hạt nhân nào tương ứng với HLO nào.

Phương pháp "tiêm hồ sơ" yêu cầu:

  • LHLO để lấy một hồ sơ làm tham số.
  • chèn profile.start / profile.end trước và sau mỗi op.
  • một vượt qua từ hồ sơ thấp đó. {bắt đầu, kết thúc} để thực hiện C ++.

Cấu hình chính xác không thể dễ dàng thực hiện đối với các op do MLIR tạo, vì:

  • MLIR không có bộ đếm thời gian, cũng không phụ thuộc vào TFRT / StreamExecutor.
  • MLIR không dễ dàng gọi vào các hàm C với các tham số phức tạp.

Bước 3: (Nhiệm vụ 2) Di chuyển thân cây

Một lưu ý, có khoảng ba loại thunks:

  • KernelThunk, khởi chạy kernel.
  • Điều khiển luồng tks, có logic điều khiển lưu lượng máy chủ (có điều kiện, trong khi, cho, trình tự) và khởi chạy các nhân cơ thể.
  • Thư viện thunks: cuDNN, cuBLAS, cuFFT, NCCL, v.v.

Kế hoạch là:

  • Làm cho Thunks (de) tuần tự hóa.
  • Giúp cải thiện TFRT đến trạng thái có thể hỗ trợ các ngữ nghĩa này.
  • Khi trạng thái được cải thiện, di chuyển tăng dần từng cá nhân.

Các mục hành động này chỉ được đặt hàng một phần. Trình tự thực hiện / song song kỹ thuật thực tế sẽ được đánh giá khi nó đi.

Bước 4: (Nhiệm vụ 3) Đã di chuyển ElementalIrEuctor

Khi hồ sơ đã sẵn sàng, chúng ta có thể hoàn thành và điều chỉnh tất cả các trình phát dựa trên ElementalIrEuctor trong MLIR. Sau đó, chúng tôi bật chúng theo mặc định, giả sử rằng tất cả các trình phát dựa trên MLIR này sử dụng một luồng duy nhất.

Lưu ý rằng cũng có ích khi di chuyển ElementalIrEuctor của XLA / CPU, vì chúng chia sẻ một phần lớn mã.

Với tất cả các hoạt động săn điểm chuẩn và hiệu năng được thực hiện (TODO: xác định tính chẵn lẻ hiệu năng), chúng tôi bật trình phát nguyên tố dựa trên MLIR mới và xóa ElementalIrEuctor kế thừa.

Bước này cũng cung cấp các chuyển tiếp hợp hạch dễ dàng (ops lồng nhau) cho lần di chuyển sau.

Bước 5: Hỗ trợ đa luồng hoặc thả

Chúng tôi không thể xóa một số trình phát cho đến khi chúng tôi hỗ trợ nó trong MLIR hoặc chúng tôi bỏ tính năng này. Đó là một lượng công việc tương đối lớn trong MLIR và một lượng nhỏ lợi nhuận cho XLA. Chúng tôi nên điều tra người dùng hiện tại của người dùng XLA / GPU đa luồng và cố gắng xóa tính năng này nếu hợp lý.

Bước 6: (Nhiệm vụ 3) Ops thiết bị di chuyển

Bước này di chuyển tất cả các op không được kiểm tra, sau đó chúng ta có thể xóa tất cả các trình phát không được kiểm tra.

Điều này kêu gọi viết lại / refactor cho kCopy và kReduce. kReduce đã được làm việc trên rất nhiều, vì vậy số lượng công việc thực tế cần phải được thực hiện vẫn còn được nhìn thấy.