Mojo Academy · 课时

编写 GPU 内核函数

在 Mojo 中表达设备代码

第 3 / 4 课13 个步骤

编写 GPU 内核函数 是 CoddyKit 上的免费 Mojo Academy 课时。 这是第 3 节课,共 4 节。 你可以在下方免费阅读本课时的完整内容 — 然后在浏览器中使用内置代码编辑器和全天候 AI 导师进行实践。 这是 Mojo Academy 学习路径的一部分,你的进度在网页和 CoddyKit 应用中同步。 Mojo Academy 课程共包含 4 节课。

什么是 Kernel

GPU kernel 是每个线程执行的函数。您为一个元素编写工作内容,硬件会将其应用到许多元素上。

从单个线程的角度思考

关键在于从单个线程的视角编写 kernel。每个线程都处理工作中的一小部分。

Kernel 就是一个 fn

在 Mojo 中,kernel 是普通的 fn。它严格且带有类型的特性,正好满足设备代码快速编译的需要。

fn add_kernel():
    pass

传入缓冲区

kernel 接受指向设备内存的指针作为参数。这些缓冲区存放输入,以及 kernel 将要写入的输出。

fn add_kernel(a: UnsafePointer[Float32], out: UnsafePointer[Float32]):
    pass

查找当前线程的索引

在 kernel 内部,每个线程首先计算自己的位置。这个索引决定了它必须处理哪个元素。

var i = block_idx.x * block_dim.x + thread_idx.x

保护边界

访问内存前务必检查索引。边界检查可以防止多余线程读取数组末尾之外的数据。

if i < n:
    out[i] = a[i] + b[i]

处理元素的工作

主体非常简短:读取输入、进行计算、写入一个结果。整个 kernel 往往只需一行数学运算。

out[i] = a[i] * b[i]

启动 Kernel

您可以通过选择网格和块大小来启动它。启动操作会将单线程代码扩展到整个网格。

ctx.enqueue_function[add_kernel](grid_dim=blocks, block_dim=256)

没有返回值

Kernel 不会向调用方返回结果。它们会将结果写入输出缓冲区,之后再由您读回。

尽量减少分支

线程在锁步执行时效果最佳。复杂的分支会使各个通道分歧并相互等待,因此请让 kernel 逻辑保持简单且统一。

相同思路,超大规模

一个简短的 kernel 加上一个大型网格,就能得到数百万个结果。规模来自启动方式,而不是更长的代码。

快速检查

您正在为一个线程编写 GPU kernel 的主体。

回顾

kernel 是供单个线程使用的 fn:找到自己的索引、保护边界、完成一个元素的计算,然后在网格上启动。⚡

免费开始

用 AI 导师学习 Mojo — 免费

在浏览器中编写并运行真实代码,获得全天候 AI 导师的即时帮助,并在网页或应用中继续学习。

课程
30
课程
120

常见问题解答

「编写 GPU 内核函数」课时是免费的吗?

是的 — 「编写 GPU 内核函数」的完整文本可在网页上免费阅读。要进行交互式练习(内置代码编辑器和全天候 AI 导师)并解锁 Mojo Academy 课程的其余内容,请升级到 CoddyKit PRO。 Mojo Academy 课程共包含 4 节课。

「编写 GPU 内核函数」这节课中我会学到什么?

在 Mojo 中表达设备代码 你通过在浏览器中直接运行的动手代码来练习 Mojo Academy,全天候 AI 导师会在你学习这节课的过程中回答你的问题。

学习 Mojo Academy 需要有经验吗?

无需任何先前经验。CoddyKit 上的 Mojo Academy 课程适合初学者到高级学习者,你可以从这里开始或从头开始,按照自己的节奏学习。 这是第 3 节课,共 4 节。

「编写 GPU 内核函数」课时需要多长时间?

大多数 CoddyKit 课程大约需要 5–10 分钟。每节课都很精短且互动,所以你能稳步进步,并在网页和应用中从离开的地方继续。

我能在这节 Mojo Academy 课中编写并运行代码吗?

能。每节 Mojo Academy 课都包含内置代码编辑器,你可以在浏览器中直接编写并运行真实代码,并获得即时 AI 反馈 — 无需本地设置。

此课程中的所有课时

  1. 为什么人工智能工作负载需要图形处理器
  2. 线程、块与网格
  3. 编写 GPU 内核函数
  4. 在设备与主机之间传输数据
← 返回 Mojo Academy