海狸 🦫

用一些神奇的长生不老药来增强全能的蓝银龙！ 🧙🧙‍♀️🧙‍♂️

动机

在使用MLIR的事实标准方式中，我们需要使用C/C++、TableGen、CMake和Python（在大多数情况下）。这里的每种语言或工具都有一些我们想要利用的功能和便利性。选择最流行和上游支持的解决方案没有任何问题，但是拥有构建基于MLIR项目的替代方法仍然是有价值的，或者至少值得尝试。

Elixir实际上可能是一个很好的MLIR前端选择。Elixir具有SSA、模式匹配、管道操作符。我们可以使用这些语言特性以自然和统一的方式定义MLIR模式和传递管道。Elixir是强类型的，但不是静态类型的，这使它成为快速构建原型以验证和探索新想法的绝佳选择。

在Beaver中构建一段IR：

Func.func some_func(function_type: Type.function([], [Type.i(32)])) do
  region do
    block _() do
      v0 = Arith.constant(value: Attribute.integer(Type.i(32), 0)) >>> Type.i(32)
      cond0 = Arith.constant(true) >>> Type.i(1)
      CF.cond_br(cond0, Beaver.Env.block(bb1), {Beaver.Env.block(bb2), [v0]}) >>> []
    end

    block bb1() do
      v1 = Arith.constant(value: Attribute.integer(Type.i(32), 0)) >>> Type.i(32)
      _add = Arith.addi(v0, v0) >>> Type.i(32)
      CF.br({Beaver.Env.block(bb2), [v1]}) >>> []
    end

    block bb2(arg >>> Type.i(32)) do
      v2 = Arith.constant(value: Attribute.integer(Type.i(32), 0)) >>> Type.i(32)
      add = Arith.addi(arg, v2) >>> Type.i(32)
      Func.return(add) >>> []
    end
  end
end

这是一个小例子，展示了如何在Beaver中定义和运行一个pass（带有一些单子魔法）：

alias Beaver.MLIR.Dialect.Func

defmodule ToyPass do
  use Beaver.MLIR.Pass, on: "func.func"

  defpat replace_add_op() do
    a = value()
    b = value()
    res = type()
    {op, _t} = TOSA.add(a, b) >>> {:op, [res]}

    rewrite op do
      {r, _} = TOSA.sub(a, b) >>> {:op, [res]}
      replace(op, with: r)
    end
  end

  def run(%MLIR.Operation{} = operation) do
    with "func.func" <- Beaver.MLIR.Operation.name(operation),
          attributes <- Beaver.Walker.attributes(operation),
          2 <- Enum.count(attributes),
          {:ok, _} <- MLIR.Pattern.apply_(operation, [replace_add_op(benefit: 2)]) do
      :ok
    end
  end
end

~m"""
module {
  func.func @tosa_add(%arg0: tensor<1x3xf32>, %arg1: tensor<2x1xf32>) -> tensor<2x3xf32> {
    %0 = "tosa.add"(%arg0, %arg1) : (tensor<1x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32>
    return %0 : tensor<2x3xf32>
  }
}
""".(ctx)
|> MLIR.Pass.Composer.nested("func.func", [
  ToyPass.create()
])
|> canonicalize
|> MLIR.Pass.Composer.run!()

目标

• 利用Elixir可组合的模块化和元编程特性，为MLIR提供简单、直观和可扩展的接口。
• 以秒级完成编辑-构建-测试-调试循环。Elixir和Zig中的所有内容都是并行编译的。
• 在MLIR的帮助下，将Elixir编译为本机/WASM/GPU。
• 在硬件加速的世界中重新审视和重生符号AI。Erlang/Elixir有Prolog的根源！
• 为机器学习引入新的技术栈。
  • 高级：Elixir
  • 表示：MLIR
  • 低级：Zig

为什么叫做Beaver（海狸）？

海狸是一种增加生物多样性的伞护种。我们希望这个项目能够像海狸池塘成为许多其他生物栖息地那样，为其他编译器和应用程序提供支持。许多Elixir项目也使用动物名称作为它们的包名，这通常是为了提高人们对濒危物种的认识。要了解更多关于海狸对我们星球重要性的信息，请查看这篇国家地理文章。

快速介绍

Beaver本质上是Erlang/Elixir上的LLVM/MLIR。看到两个成熟社区和四个子社区的交叉很有趣。以下是关于它们每一个的简要信息。

对于Erlang/Elixir分支

• 用一句话向我解释这个MLIR

  MLIR可以被视为编译器的XML，而MLIR方言就像HTTP标准，为通用格式提供了现实世界的语义和功能。

• 查看MLIR的主页。

对于LLVM/MLIR分支

• Elixir这种编程语言有什么好处？

  • 它被编译成Erlang并在BEAM（Erlang的VM）上运行。因此，它具有Erlang的所有容错和并发特性。
  • 作为一种Lisp，Elixir拥有Lisp语言的所有优点，包括卫生宏和基于协议的多态性。
  • Elixir有一个强大的模块系统来保存编译时数据，这允许库用户轻松调整运行时行为。
  • 最小化，很少的关键字。大部分语言都是用自身构建的。

• 查看Elixir的官方指南。

入门

• 教程：使用 Beaver 构建你的第一个编译器！

安装

可以通过在 mix.exs 的依赖列表中添加 beaver 来安装这个包：

def deps do
  [
    {:beaver, "~> 0.3.9"}
  ]
end

在 .formatter.exs 中添加以下内容，可以让格式化工具正确处理 beaver 引入的宏

import_deps: [:beaver],

与 Beaver 相关的 Erlang 应用

LLVM/MLIR 是一个庞大的项目，围绕它构建的 Beaver 包含数千个函数。为了适当地发布 LLVM/MLIR 并简化开发流程，我们需要谨慎地将不同层级的功能拆分到同一个伞形项目下的不同 Erlang 应用中。

• :beaver：Elixir 和 C/C++ 混合。
  • 顶层应用，提供高层功能，包括 IR 生成和模式定义。
  • MLIR CAPI 封装，通过解析 LLVM/MLIR CAPI C 头文件构建，以及一些中层辅助函数以隐藏 C 指针相关操作。这个应用会将加载的 MLIR C 库和管理的 MLIR 上下文添加到 Erlang 监督树中。该应用也使用 Rust，但主要用于 LLVM/MLIR CMake 集成。
  • 所有在标准 MLIR 方言中定义的操作，通过查询注册表构建。这个应用会以符合 Erlang 习惯的方式（如行为遵从）发布 MLIR 操作。
• :kinda：Elixir 和 Zig 混合，从 MLIR C 头文件生成 NIF。仓库：https://github.com/beaver-lodge/kinda
• :manx：纯 Elixir，Nx 的编译器后端。

使用和开发注意事项

• 只有 :beaver 和 :kinda 被设计为可以作为独立应用直接被其他应用使用。
• :manx 只能与 Nx 一起工作。
• 虽然 :kinda 是为 Beaver 构建的，但任何对打包 C API 感兴趣的 Erlang/Elixir 应用也可以利用它。
• 命名空间 Beaver.MLIR 用于任何 MLIR 工具中通常预期的标准功能。
• 命名空间 Beaver 用于仅存在于 Beaver 中的概念和实践，这些主要是作为一组宏提供的 DSL（包括 mlir/0、block/1、defpat/2 等）。实现通常在 Beaver.DSL 命名空间下。
• 在 Beaver 中，Erlang 应用名称和 Elixir 模块名称之间没有严格的一致性要求。两个具有相同命名空间前缀的模块可能位于不同的 Erlang 应用中（这在 Beaver.MLIR 命名空间中经常发生）。当然，应避免重复定义具有相同名称的 Elixir 模块。

工作原理

要实现 MLIR 工具包，我们至少需要以下几组 API：

• IR API，用于创建和更新 IR 中的操作和块
• Pass API，用于创建和运行 Pass
• Pattern API，用于声明特定操作结构的转换

我们借助 MLIR C API 实现 IR API 和 Pass API。既有从 C 头文件生成的低级 API，也有更符合 Elixir 习惯的高级 API。 Pattern API 借助 PDL 方言 实现。我们使用低级 IR API 将 Elixir 代码编译为 PDL。另一种看待这个的方式是，Elixir/Erlang 的模式匹配作为 PDLL 的替代前端。

设计原则

转换优于构建器

使用构建器模式构建 IR 在面向对象的编程语言（如 C++/Python）中很常见。 这种方法的一个问题是，编译器代码看起来与它生成的代码非常不同。 由于 Erlang/Elixir 本质上是 SSA 的，在 Beaver 中，MLIR 操作的创建非常声明式，其容器会用正确的上下文信息对其进行转换。通过这种方式，我们可以：

• 保持编译器代码的结构尽可能接近生成的代码，减少噪音，提高可读性。
• 允许不同目标和语义的方言引入不同的 DSL。例如，CPU、SIMD、GPU 都可以有针对其独特概念定制的专门转换。

一个例子：

module do
  v2 = Arith.constant(1) >>> ~t<i32>
end
# module/1 是一个宏，它会将 SSA `v2 = Arith.constant..` 转换为：
v2 =
 %Beaver.SSA{}
  |> Beaver.SSA.put_arguments(value: ~a{1})
  |> Beaver.SSA.put_block(Beaver.Env.block())
  |> Beaver.SSA.put_ctx(Beaver.Env.context())
  |> Beaver.SSA.put_results(~t<i32>)
  |> Arith.constant()

此外，使用声明式方式构建 IR，可以自然形成正确的支配关系和操作数引用。

SomeDialect.some_op do
  region do
    block entry() do
      x = Arith.constant(1) >>> ~t<i32>
      y = Arith.constant(1) >>> ~t<i32>
    end
  end
  region do
    block entry() do
      z = Arith.addi(x, y) >>> ~t<i32>
    end
  end
end

# 将被转换为：

SomeDialect.some_op(
  regions: fn -> do
    region = Beaver.Env.region() # 创建第一个区域
    block = Beaver.Env.block()
    x = Arith.constant(...)
    y = Arith.constant(...)

    region = Beaver.Env.region() # 创建第二个区域
    block = Beaver.Env.block()
    z = Arith.addi([x, y, ...]) # x 和 y 支配 z
  end
)

Beaver DSL 作为 MLIR 的高级 AST

Beaver SSA DSL 和 MLIR SSA 之间应该有一对一的映射。可以通过解析 MLIR 文本格式并将其转储为 Beaver DSL（本质上是 Elixir AST）来实现往返转换。这使得以更可编程和可读的方式轻松调试 IR 片段成为可能。

在 Beaver 中，无论是生成、转换还是调试，处理 MLIR 都应该使用同一种格式。

符合 Erlang/Elixir 习惯的高级 API

在可能的情况下，低级 C API 应该被封装为支持常见 Elixir 协议的 Elixir 结构。 例如，对 MLIR 操作的操作数、结果、后继、属性、区域的迭代应该在 Elixir 的 Enumerable 协议中实现。 这使得可以使用 Elixir 标准库和 Hex 包中丰富的函数集合。

Beaver 是编译器还是 LLVM/MLIR 的绑定？

Beaver 既是编译器也是 LLVM/MLIR 的绑定。它提供了一套用于构建编译器的工具和 API，同时也作为 LLVM/MLIR 的 Elixir 接口。Beaver 允许用户以 Elixir 风格使用 MLIR 的功能，并提供了额外的抽象和工具来简化编译器开发过程。 Elixir是一种为各种用途而构建的编程语言。在整个Erlang/Elixir生态系统中存在多个子生态系统。每个子生态系统看似互不相关，但在实际生产中它们实际上是相互补充的。举几个例子：

• Phoenix Framework用于Web应用和实时消息
• Nerves Project用于嵌入式设备和物联网
• Nx用于张量和数值计算

这些子生态系统都始于一个种子项目或库。Beaver应该发展成为一个用Elixir和MLIR构建的编译器子生态系统。

MLIR上下文管理

在调用高级API时，最好不要到处传递MLIR上下文。如果没有提供MLIR上下文，属性和类型获取器应该返回一个以MLIR上下文为参数的匿名函数。在Erlang中，所有值都是复制的，所以传递这些匿名函数是非常安全的。在创建操作时，这些函数将在操作状态中使用MLIR上下文被调用。通过这种方法，我们既实现了简洁性又实现了模块化，而不需要全局MLIR上下文。在Beaver中，接受MLIR上下文来创建操作或类型的函数通常被称为"创建器"。

开发

1. 安装Elixir，https://elixir-lang.org/install.html
2. 安装Zig，https://ziglang.org/learn/getting-started/#installing-zig
3. 安装LLVM/MLIR

• 选项1：使用pip安装

  python3 -m pip install -r dev-requirements.txt
  export LLVM_CONFIG_PATH=$(python3 -c 'import mlir;print(mlir.__path__[0])')/bin/llvm-config

• 选项2：从源代码构建 https://mlir.llvm.org/getting_started/ 推荐的安装命令：

  cmake -B build -S llvm -G Ninja -DLLVM_ENABLE_PROJECTS=mlir \
    -DLLVM_TARGETS_TO_BUILD="host" \
    -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
    -DLLVM_ENABLE_ASSERTIONS=ON \
    -DLLVM_ENABLE_OCAMLDOC=OFF \
    -DLLVM_ENABLE_BINDINGS=OFF \
    -DCMAKE_BUILD_TYPE=RelWithDebInfo \
    -DCMAKE_INSTALL_PREFIX=${HOME}/llvm-install
  cmake --build build -t install
  export LLVM_CONFIG_PATH=$HOME/llvm-install/bin/llvm-config

  （可选）使用Vulkan：

  • 安装Vulkan SDK（需要全局安装），参考：https://vulkan.lunarg.com/sdk/home

  • 通过将以下命令添加到你的bash/zsh配置文件来设置环境变量：

    # 你可能需要在这里更改版本
    cd $HOME/VulkanSDK/1.3.216.0/
    source setup-env.sh
    cd -
    

  • 使用vulkaninfo和vkvia来验证Vulkan是否正常工作

  • 在LLVM CMake配置命令中添加-DMLIR_ENABLE_VULKAN_RUNNER=ON

4. 开发和运行测试

• 在同一目录下克隆此仓库和kinda

  git clone https://github.com/beaver-lodge/beaver.git
  git clone https://github.com/beaver-lodge/kinda.git

• 确保LLVM环境变量设置正确，否则可能无法构建

  echo $LLVM_CONFIG_PATH

• 构建并运行Elixir测试

  mix deps.get
  BEAVER_BUILD_CMAKE=1 mix test
  # 使用过滤器运行测试
  mix test --exclude vulkan # 使用此命令跳过vulkan测试
  mix test --only smoke
  mix test --only nx

5. 调试

• 设置环境变量以控制Erlang调度器数量，ERL_AFLAGS="+S 10:5"
• 在LLDB下运行mix test，scripts/lldb-mix-test

6. Livebook

• 请使用Elixir 1.14并从GitHub源代码安装Livebook：

  mix escript.install github livebook-dev/livebook

• 要在Livebook中使用Beaver，在源目录中运行：

  livebook server --name livebook@127.0.0.1 --home .

• 在设置单元格中，将内容替换为：

  beaver_app_root = Path.join(__DIR__, "..")
  
  Mix.install(
    [
      {:beaver, path: beaver_app_root, env: :test}
    ],
    config_path: Path.join(beaver_app_root, "config/config.exs"),
    lockfile: Path.join(beaver_app_root, "mix.lock")
  )

发布新版本

更新Elixir源代码

• 在README.md和mix.exs中更新版本

Linux

• 运行CI，它会生成新的GitHub发布并上传到https://github.com/beaver-lodge/beaver-prebuilt/releases。
• 在mix.exs中更新发布URL

Mac

• 使用以下命令运行macOS构建：

  rm -rf _build/prod
  bash scripts/build-for-publish.sh

• 将beaver-nif-[xxx].tar.gz文件上传到发布

生成checksum.exs

rm checksum.exs
mix clean
mix
mix elixir_make.checksum --all --ignore-unavailable --print

检查输出中的版本是否正确。

发布到Hex

BEAVER_BUILD_CMAKE=1 mix hex.publish

（可选）格式化CMake文件

python3 -m pip install cmake-format
cmake-format -i native/**/CMakeLists.txt native/**/*.cmake