Hai lập trình viên dùng AI viết xong driver GPU OpenGL ES 3.0 cho M4 Mac Mini chỉ trong một tháng

Bài đăng blog trước: https://codyho.dev/blog/hypervisor-macbook-neo/

Chúng tôi đã làm gì

Tóm tắt: Niklas và tôi đã xây dựng một trình điều khiển GPU tương thích hoàn toàn với OpenGL ES 3.0 cho M4 Mac Mini và MacBook Neo trong khoảng một tháng, một quá trình thường mất nhiều năm. Đây là Chrome và Firefox đang chạy WebGL trên M4 Mac Mini với tính năng compositing hoạt động bình thường:

Quan trọng nhất, trình điều khiển này đủ nhanh để chạy Minecraft ở tốc độ 200fps:

Việc xây dựng trình điều khiển này bao gồm kỹ thuật đảo ngược (reverse engineering) firmware ABI cực kỳ phức tạp của AGX (tên Apple đặt cho GPU) và các thành phần user-space. Tất cả đều được thực hiện một cách minh bạch, theo phương pháp "clean room" có thể kiểm chứng bằng các kỹ thuật đã được thiết lập. Mã nguồn hiện chưa sẵn sàng cho người dùng cuối, nhưng chúng tôi đang tìm cách đưa nó đến tay người dùng sớm nhất có thể.

Chúng tôi đã thực hiện như thế nào

Trước đây, tôi đã xây dựng một hypervisor để thực hiện kỹ thuật đảo ngược macOS. Giờ đây, mục tiêu trở thành việc thực sự làm điều gì đó hữu ích với nó, và còn mục tiêu nào tốt hơn là viết một trình điều khiển GPU. GPU thực sự là một yêu cầu đối với bất kỳ hệ thống hiện đại nào, nếu không, mọi thứ sẽ phải được render bằng CPU, vốn chậm hơn và kém hiệu quả năng lượng hơn gấp nhiều lần. Mục tiêu của chúng tôi là triển khai các trình điều khiển OpenGL (và sắp tới là Vulkan) tương thích cho M4 Mac Mini và MacBook Neo.

Thông thường, việc xây dựng một trình điều khiển GPU là một nỗ lực kéo dài nhiều năm; mục tiêu của chúng tôi là thực hiện nó trong vài ngày. Hóa ra "vài ngày" là quá lạc quan, nhưng "vài tuần" vẫn là một sự cải thiện to lớn. Trong những tuần đó, chúng tôi đã:

Trong suốt quá trình này, chúng tôi không xem bất kỳ tệp nhị phân nào của Apple, chỉ xem các dấu vết phần cứng (từ hypervisor của chúng tôi) và các shader do chính chúng tôi xây dựng. Đối với việc RE đồ họa user-space, chúng tôi cẩn thận xử lý mọi blob của Apple như các đối tượng hộp đen (opaque objects). Chúng tôi đã nhờ một người bạn viết tài liệu về các blob này 1 để chúng tôi có thể tự viết một bản triển khai "clean room" (phần lớn được xây dựng bằng cách thử nghiệm mù quáng cho đến khi nó hoạt động). Chúng tôi đã công bố tất cả các thí nghiệm của mình để bất kỳ ai cũng có thể xác minh nguồn gốc công việc của chúng tôi (xem hai kho lưu trữ agx-re trong phần Deliverables).

Bài đăng blog này được chia thành hai phần: user space và kernel space. Điều này phản ánh sự phân chia trong tất cả các trình điều khiển GPU hiện đại: kernel chịu trách nhiệm giao tiếp với firmware, cấp phát bộ đệm và quản lý lập lịch, trong khi nội dung thực tế của các bộ đệm đó và những gì đang được lập lịch là các đối tượng hộp đen. User space chịu trách nhiệm thực sự hiểu cách GPU hoạt động và lấp đầy các bộ đệm đó bằng dữ liệu.

Kernel Space

Trên Apple Silicon, trình điều khiển kernel không giao tiếp trực tiếp với phần cứng. Thay vào đó, nó nói chuyện với firmware GPU đang chạy một RTOS tùy chỉnh có tên là RTKit. Điều đó có nghĩa là bước đầu tiên để tạo trình điều khiển kernel không phải là nói chuyện với phần cứng, mà là tìm ra firmware ABI.

Firmware ABI là phần gây khó chịu nhất trong dự án này, bởi vì thay vì làm điều hợp lý là tạo ra một ABI hợp lý với các giao diện tốt, Apple về cơ bản đã lấy một trình điều khiển kernel thông thường, cắt làm đôi, sau đó đặt một nửa vào AGX và gọi đó là firmware, với nửa còn lại của trình điều khiển kernel giao tiếp bằng cách sử dụng các struct dùng chung trong bộ nhớ. Nhiều struct trong số này có các trường do firmware sở hữu (mà chúng ta không bao giờ được sửa đổi và phải tìm hiểu thông qua kỹ thuật đảo ngược) xen kẽ với các trường do host kiểm soát. Để hình dung ABI phức tạp như thế nào, đây là hình ảnh cây bộ nhớ dùng chung trên M1/M2:

Asahi Lina đã nổi tiếng khi tìm ra tất cả những điều này sau những ngày làm việc vất vả 12 tiếng để xây dựng trình điều khiển kernel cho M1/M2, một thành tựu kỹ thuật đáng kinh ngạc. Thật không may, firmware ABI của A18 Pro (tôi bắt đầu công việc RE của mình trên MacBook Neo và sau đó chuyển sang M4 Mac Mini) phức tạp hơn đáng kể so với firmware ABI của M1 vốn đã rất phức tạp:

Cái quái gì vậy, Apple. Lưu ý cách A18 có:

Có nhiều vấn đề khác gây khó khăn cho quá trình RE 2. Tôi có một số tài liệu về firmware ABI, nhưng nó rất thiếu sót và thành thật mà nói là không hữu ích lắm 3.

Cách tiếp cận của tôi rất đơn giản và dựa trên phương pháp được sử dụng để thực hiện kỹ thuật đảo ngược thành công các máy M1/M2: quan sát những gì macOS đã làm, phát lại nó, sau đó cố gắng tự thực hiện, điều này trở nên khả thi nhờ hypervisor.

Khi tôi mô tả cách tiếp cận này với LLM, nó đã thực hiện việc phát lại theo nghĩa đen: việc đầu tiên nó làm là đợi sự kiện hiển thị firmware đầu tiên (được gọi là “kicks”), sau đó lưu một bản sao của toàn bộ trạng thái bộ nhớ GPU. Sau khi khởi động lại, nó sao chép trạng thái bộ nhớ đã lưu trực tiếp trở lại bộ nhớ host, thực hiện cú "kick" và thấy các trang đầu ra thay đổi. Sau đó, nó cố gắng tái tạo các đối tượng này trong mã, theo dõi tất cả các con trỏ và hiểu nội dung. Qua các thí nghiệm liên tiếp, Codex đã giảm số lượng trang nó sao chép cho đến khi không còn trạng thái phát lại nào và mọi thứ được xây dựng từ mã nguồn. 4 Thật kinh ngạc, tôi nhận thấy Codex có "gu" tốt trong việc khi nào nên tác động thêm vào phần cứng và khi nào nên chỉ chạy hypervisor để tự chụp trạng thái.

Có ba vấn đề lớn, và tất cả đều do chúng tôi không thể có được một bản chụp sạch công việc của host:

Vấn đề đầu tiên là công việc render được gửi sau khi firmware GPU bắt đầu. Chúng tôi có thể chuẩn bị trước công việc trước khi firmware bắt đầu, khởi động GPU và công việc đó sẽ được hoàn thành như mong đợi, nhưng một khi firmware đã bắt đầu, bất kỳ công việc nào được gửi đến sẽ chỉ được ACK và kết thúc mà không thực sự làm gì cả. Khi firmware đã bắt đầu, việc chụp trạng thái khó hơn nhiều vì mọi thứ trở nên động và firmware trở thành một đối tượng có trạng thái mà bạn không thể dễ dàng phát lại.

Tôi phải can thiệp vào thời điểm này và kiểm tra quy trình của Codex. Hóa ra nó đang cố gắng phát lại một bản chụp rất muộn trong vòng đời của AGX, nơi đã có nhiều sự kiện trước đó. Khi tôi bảo nó chọn một bản chụp sớm hơn nhiều trong vòng đời của AGX, bản chụp đầu tiên ngay sau khi firmware bắt đầu, Codex đã có thể phát hiện ra vấn đề gần như ngay lập tức (nó bị thiếu một bộ mô tả một byte). Việc này mất vài ngày.

Vấn đề thứ hai, và là vấn đề chặn chính duy nhất, là tính toán (compute). AGX, nhìn chung, hỗ trợ hai loại công việc: tính toán và render. Trong đường dẫn GUI thông thường, công việc tính toán chỉ được lập lịch sau khi một lượng đáng kể công việc render đã được thực thi. Do đó, mất rất nhiều thời gian để có được một bản chụp sạch của khối lượng công việc tính toán, và khi Codex cuối cùng làm được thì nó nặng 336 MB và không thể phát lại (nó đã cố gắng trong một thời gian dài). Nó cũng cố gắng tự xây dựng các đối tượng bằng cách nhìn vào bản chụp và đã dành hơn một tuần để làm việc này, nhưng cuối cùng không thành công. Có quá nhiều thứ vô nghĩa để sàng lọc. Điều này trở nên trầm trọng hơn do các vấn đề từ phía tôi – sau khi làm cho việc render hoạt động, tôi cho rằng việc gửi công việc tính toán sẽ đơn giản hơn (firmware ABI cho tính toán thực sự đơn giản hơn, vì vậy tôi đã đúng ở đây), nhưng tôi đã mất đi sự khiêm tốn và nghĩ rằng đó sẽ là một việc dễ dàng chỉ mất vài giờ. Do đó, tôi đã không chuẩn bị kỹ lưỡng nhiệm vụ cho LLM.

Giải pháp thực sự đã được đưa ra cho tôi trong một phiên Codex khác. Về cơ bản:

Dấu vết đã được chụp thành công. Trong vòng vài giờ, Codex đã giải mã nó, và trong vòng vài ngày, Codex đã làm cho phần tính toán hoạt động. Còn về lý do tại sao codebase tính toán ban đầu không hoạt động… Codex không biết. Bản hoạt động và bản bị lỗi trông rất giống nhau.

Nhìn lại, đây lẽ ra phải là chiến lược ngay từ đầu – bản chụp nhỏ nhất có thể, chạy ở chế độ single-user để không làm ảnh hưởng đến kết quả. Tôi đã rút kinh nghiệm từ những sai lầm của mình ở đây cho vấn đề cuối cùng:

Render một phần (Partial renders) hóa ra là một trong những thứ khó tìm ra nhất. Chúng xảy ra khi Tiled Vertex Buffer (TVB) không đủ lớn để lưu trữ hình học hiện tại (tức là có quá nhiều hình tam giác để vẽ). Trong những trường hợp này, có hai lựa chọn và trình điều khiển cần hỗ trợ cả hai: hoặc tăng kích thước của TVB, hoặc thực hiện render một phần, tức là render một phần hình học, sau đó tải lại bộ đệm với phần còn lại của các hình tam giác và hoàn thành việc render một phần. Những bản render một phần này hóa ra rất, rất khó chiều, thậm chí còn hơn cả phần còn lại của công việc vì chúng có nghĩa là phải thêm tính năng lưu và tiếp tục vào trình điều khiển GPU.

Quy trình làm việc mà tôi khám phá ra trước đó rất hữu ích ở đây. Codex đã có thể phát lại một giao dịch render một phần, sau đó sửa đổi shader Metal của chúng tôi để thực hiện nhiều lần render một phần (điều này khá dễ dàng bằng cách chỉ cần dồn hàng ngàn hình tam giác vào một ô cho đến khi một bản render một phần được kích hoạt) và sau đó học cách phát lại chúng. Khi Codex đã phát lại thành công, chỉ còn là vấn đề thời gian cho đến khi nó học được cách tự xây dựng nó.

Xây dựng trình điều khiển Kernel

Việc chuyển từ trình điều khiển nguyên mẫu Python sang trình điều khiển Linux đầy đủ tính năng mất ba ngày, và một trong những ngày đó gần như hoàn toàn lãng phí vì Codex, vì lý do nào đó mà tôi vẫn không hiểu, đã chọn giải quyết các bản render một phần trước (nhiệm vụ khó nhất) thay vì thực hiện tính toán trước (nhiệm vụ dễ nhất). Khi tôi bảo nó thực hiện tính toán trước, mọi thứ diễn ra suôn sẻ.

Ở mức độ cao, toàn bộ quá trình về cơ bản là:

Tất cả đây chỉ là công việc kỹ thuật thông thường mà các LLM chắc chắn có khả năng thực hiện.

Điều đáng chú ý duy nhất tôi thấy là Codex đã sử dụng hypervisor một cách tích cực để gỡ lỗi tại sao mã của nó không hoạt động, bao gồm việc chụp toàn bộ không gian địa chỉ và so sánh nó với các mẫu tốt đã biết. Kiểu gỡ lỗi có hệ thống này là lý do tại sao Codex là tác nhân lập trình yêu thích của tôi.

User Space

User space của A18 Pro rất khác so với M1/M2; nó có các định dạng bộ mô tả mới, ISA mới và một loạt các thứ mới khác. Ngoài việc là một trình render trì hoãn dựa trên ô (tile based deferred renderer) được thiết kế để chạy Metal, nó chỉ đơn giản là một GPU khác.

Tin tốt là RE user-space có một quy trình rất rõ ràng. Chỉ cần viết một chương trình Metal nhỏ, biên dịch nó, chạy nó, xem những gì đã thay đổi, sau đó tháo rời nó và bắt đầu mày mò với các bit cho đến khi chúng ta hiểu tất cả chúng làm gì. Nếu bạn đang nghĩ điều này nghe giống như công việc nhàm chán, lặp đi lặp lại mà LLM rất giỏi, thì bạn đã đúng.

Công việc RE diễn ra trong hai giai đoạn. Đối với giai đoạn đầu tiên, tôi đã để Claude xem mọi chương trình Metal mà nó có thể tìm thấy và cố gắng xây dựng một trình tháo rời (disassembler), trình hợp dịch (assembler) và hiểu định dạng của tất cả các bộ mô tả/luồng lệnh/v.v. khác cần thiết cho trình điều khiển GPU. Tôi đã để Claude liệt kê mọi thứ, bao gồm cả những thứ Linux không thể sử dụng (như tessellation) để đảm bảo tính đầy đủ. Điều này đã thành công, nhưng chỉ vì Claude có thể tháo rời và sau đó lắp ráp lại các chương trình không có nghĩa là nó biết cách tự xây dựng một chương trình. Khi cố gắng thu hẹp khoảng cách này, tức là hiểu đủ mọi lệnh để thực sự có thể biên dịch các chương trình tùy ý của riêng chúng tôi, Claude đã làm một công việc khủng khiếp và về cơ bản là không đạt được tiến bộ nào.

Tại thời điểm này, Niklas đã hoàn thành drm-shim của mình cho M4 Mac Mini và tham gia cùng tôi trong giai đoạn thứ hai của RE user-space. Chúng tôi có hai cách tiếp cận khác nhau để thực sự hoàn thành trình điều khiển user-space:

Cách tiếp cận của tôi là ưu tiên RE phần cứng và tập trung vào việc tìm hiểu cách phần cứng và tất cả các lệnh hoạt động. Sau đó, tôi sẽ viết một đặc tả và để LLM triển khai nó, hy vọng sẽ kết thúc với sự tương thích hoàn toàn với OpenGL và Vulkan. Điều này có nghĩa là phần lớn thời gian của LLM của tôi được dành cho việc viết các thí nghiệm trên phần cứng, chứ không phải thực sự triển khai mã Mesa. Ý tưởng là một khi tôi hiểu phần cứng, mọi thứ khác sẽ theo sau.

Niklas đã thực hiện một cách tiếp cận khác, mà tôi mô tả là “Mesa trước”. Về cơ bản, anh ấy đã cố gắng xây dựng Mesa trước và chỉ thực hiện RE để xây dựng một số chức năng. Thời gian của anh ấy được chia cho việc xây dựng và thử nghiệm Mesa, và thực hiện RE.

Hóa ra Niklas đã đạt được tiến bộ nhanh hơn đáng kể so với tôi, bởi vì tác nhân của tôi dành quá nhiều thời gian cho các nhiệm vụ nhỏ, không quan trọng vì mục đích hoàn thiện. Ngược lại, tác nhân của anh ấy được củng cố bởi nhu cầu thực sự xây dựng Mesa, vì vậy nó đã sử dụng thời gian và tài nguyên hiệu quả hơn nhiều. Anh ấy cuối cùng đã tiến nhanh hơn tôi rất nhiều đến mức tôi chỉ cố gắng hỗ trợ công việc của anh ấy bằng cách điều tra bất kỳ hành vi nào mà anh ấy chưa hiểu.