OOP-Hw:LuisaCompute Code Analysis

Posted on 2024-09-19 Edited on 2024-10-25

这是本学期按面向对象编程课程的大作业“代码分析”报告。第一次，直面奇迹工程。

~~当然，选这个的原因之一也是我和核心开发者私交甚笃~~

1.

1.1. Main Functions and Procedures

$\text{LuisaCompute}$ （以下简称 $\text{LC}$ ），其中 $\text{LUISA}$ 全称为 $\text{Layered and Unified Interfaces on Stream Architectures}$ ，即“流式架构上的分层统一接口“，是一个面向图形等应用的高性能跨平台计算框架。

~~关于这个取名有一个很有意思的八卦小故事~~

$\text{LC}$ 主要功能可分为如下三部分：

一种用于并行核函数编程的嵌入在现代C++中的领域特定语言，利用即时编译技术进行代码生成和编译。
一个统一的运行时库，提供资源封装器用于跨平台资源管理和命令调度。
多个高度优化的并行执行后端，包括CUDA、DirectX、Metal和CPU。

基本流程：

编写核函数，这些函数的编写是嵌入在C++代码中的，例如：

Callable to_srgb = [](Float3 x) {
    $if (x <= 0.00031308f) {
        x = 12.92f * x;
    } $else {
        x = 1.055f * pow(x, 1.f / 2.4f) - .055f;
    };
    return x;
};
Kernel2D fill = [&](ImageFloat image) {
    auto coord = dispatch_id().xy();
    auto size = make_float2(dispatch_size().xy());
    auto rg = make_float2(coord) / size;
    // invoke the callable
    auto srgb = to_srgb(make_float3(rg, 1.f));
    image.write(coord, make_float4(srgb, 1.f));
};

主机端和设备端的代码逻辑从而可以集成在一起。

将 $\text{C++}$ 代码编译，可以得到一个可执行程序。然而我们编写的DSL在此时并没有被执行，它将被编译器进行宏/模板的展开和替换，转化成普通的 $\text{C++}$ 代码，记录下这些代码的结构信息。
$\text{LC}$ 会在运行时解析并生成IR,并交给不同的后端生成代码执行。不同后端被抽象出平台无关的“资源”的概念，从而可以进行统一的编程。

1.1.1. Record the Operations

从高层次来讲， $\text{C++}$ 中的宏，模板等多种语法特性，为开发人员提供了极高的灵活度，使得“在C++代码内实现内嵌的领域特定语言”成为可能。

个人以为，在 $\text{LC}$ 的“语法糖”中十分重要的设计有两点：

通过高超的 $\text{C++}$ 技巧，实现一套完整且灵活的变量及其运算系统。这套系统的特点是：所有编写的代码经过抽象，用户使用起来与一般的C++代码无异，但是在它们的底层，并没有真正地实现运算，它们实际做的事情是记录下进行的运算，以派发给后端进行代码生成。
而从用户的角度看，用户并不知道底层做了什么。
将复杂代码逻辑转化为可调用对象，并且通过灵活组合，使其自动形成 $\text{AST}$ ，易于访问。

下面我们分别讲述这两部分。

1.1.1.1. Variable and Operator System

在LC中，用户编写的代码里会使用到变量等，但是它们并不是真正的变量，而是类似于“占位符”，对这些占位符，可以使用重载运算符的操作，来让它们“看起来”与正常变量一样，而底层实现逻辑得到高度的自定义。

如果我们观察 $\text{LC}$ 里重载运算符的宏：


#define LUISA_MAKE_GLOBAL_DSL_BINARY_OP(op, op_tag_name)                              \
    template<typename Lhs, typename Rhs>                                              \
        requires luisa::compute::any_dsl_v<Lhs, Rhs> &&                               \
                 luisa::is_basic_v<luisa::compute::expr_value_t<Lhs>> &&              \
                 luisa::is_basic_v<luisa::compute::expr_value_t<Rhs>>                 \
    [[nodiscard]] inline auto operator op(Lhs &&lhs, Rhs &&rhs) noexcept {            \
        using R = luisa::compute::detail::dsl_binary_op_return_type<                  \
            luisa::compute::BinaryOp::op_tag_name, Lhs, Rhs>;                         \
        return luisa::compute::dsl::def<R>(                                           \
            luisa::compute::detail::FunctionBuilder::current()->binary(               \
                luisa::compute::Type::of<R>(),                                        \
                luisa::compute::BinaryOp::op_tag_name,                                \
                luisa::compute::detail::extract_expression(std::forward<Lhs>(lhs)),   \
                luisa::compute::detail::extract_expression(std::forward<Rhs>(rhs)))); \
    }

阅读源码可知，def<R>返回的是一个类似于Var<R>的类型（并不严格是）。知道了这一点，我们就会发现这段代码虽然复杂，但是逻辑很好理解：推断出二元运算操作的返回值类型R，然后在当前正在构造的 $\text{LC}$ 函数对象中创建一个二元表达式，将其作为一个具有R类型的临时变量。与此同时，返回这个临时变量的指针，可以使其继续参与运算，从而自然地形成表达式树。

那么到这里，我们要问一个问题： $\text{C++}$ 特有的“模板”系统，对于LC的设计来说真的很重要吗？

一开始，我的看法是：如果没有模板，那么整个 $\text{LC}$ 的抽象便无法成立了。
但是经过我的阅读和思考，我改变了想法。在这个模块中，“模板”只是一种在编译期实现多态的方式，编译器能够看到用户的代码，然后选取合适的模板进行实例化，在实例化的类型和函数中完成分发和记录。
真正核心的操作在于将“记录”封装为“运算”，而不在于多态分发的过程。即使我们不使用模板，我们一样可以编写一套独立的变量系统，通过动态多态的方式来针对不同的运算完成不同的操作记录。
其缺点在于，由于没有模板，编译器无法拿到更多的信息进行优化，也会引入动态多态的分发开销，会造成极大的性能损失，其优点在于可以带来动态的分析和执行。而 $\text{LC}$ 本身的定位是作为源码库引入项目，
它的使用场景只有静态，并且重视性能，因此选择模板实现性能更高的“静态多态”（或者“编译期多态”）更为合适，实现上虽然技巧复杂，但是配合宏的确减少了大量的重复代码。

对此，我总结为： $\text{LC}$ 本身的设计并不依赖于模板系统，即使没有模板，一样可以通过动态多态完成抽象的任务，同样易于拓展。只是考虑到 $\text{LC}$ 本身为了能够在编译期拿到信息实现大量的优化，出于性能考虑选择了模板来实现“编译期多态”。

到此，我们就知道了 $\text{LC}$ 如何记录用户的基本运算操作。

1.1.1.2. Complex Statments

如果不考虑任何高层语法糖的封装，那么一个基本的实现应该是：

用户编写的基本运算被通过某种方式记录下来，这一点我们已经在前面提过，利用的是 $\text{LC}$ 独特的变量系统和精巧的实现。
更高层的语法树节点对象，例如if/else，for等，它们有着不同的代码执行逻辑，它们能够保存不同的可调用对象，将这些可调用对象按照节点自身的执行特点进行组合。

$\text{实现细节}$
在 $\text{LC}$ 中的 $\text{DSL}$ 中，像$if，$else，$return等关键字都是通过宏定义实现的。这些宏定义在include/luisa/dsl/sugar.h中。

理解 $\text{LC}$ 实现的第一步，是从高层理解这些宏做了什么，它们是如何将用户的代码记录为 $\text{AST}$ 的。在include/luisa/dsl中，我们可以看到有关这一部分的实现。

以前面to_rgb里出现的if/else的定义为例：


#define $return(...) ::luisa::compute::return_(__VA_ARGS__)

#define $if(...)                                                                  \
    ::luisa::compute::detail::IfStmtBuilder::create_with_comment(                 \
        ::luisa::compute::dsl_detail::format_source_location(__FILE__, __LINE__), \
        __VA_ARGS__) %                                                            \
        [&]() noexcept
#define $else \
    / [&]() noexcept
#define $elif(...) \
    *([&] { return __VA_ARGS__; }) % [&]() noexcept

乍一看很复杂，但是如果我们把宏替换进去就能够看出 $\text{LC}$ 的实现意图：


Callable to_srgb = [](Float3 x) {
    ::luisa::compute::detail::IfStmtBuilder::create_with_comment(
        ::luisa::compute::dsl_detail::format_source_location(__FILE__, __LINE__), x <= 0.00031308f) %
        [&]() noexcept {
            x = 12.92f * x;
        } / [&]() noexcept {
            x = 1.055f * pow(x, 1.f / 2.4f) - .055f;
        };
    return x;
};

我们可以看出，每个分支内的代码在宏展开以后都替换为了无返回值的lambda表达式。这两个lambda表达式通过重载运算符，和工厂函数IfStmtBuilder::create_with_comment创建出的对象连接在一起，形成了一个IfStmt的结构。我们暂且可以不管这些重载运算符是什么含义，这里就体现了 $\text{LC}$ 里面非常重要的一个想法，就是将代码转化为可调用对象的形式。

实际上，在 $\text{C++}$ 中，类似的设计可谓常见。降低代码耦合的关键，就是提出不变的部分，然后将“可变化”的部分作为依赖注入其中。对于函数式而言，“可变化”的部分就是函数，对于面向对象而言，“可变化”的部分就是虚接口。这种“可调用对象”的想法，在 $\text{C++ STL}$ 的泛型算法库中得到了大量使用，在 $\text{LC}$ 中也不例外，可以认为是一种函数式编程思想的体现（或者也可以说，函数也是对象，这也是面向对象编程的一种体现）。而在这里，如果使用虚接口，那么用户将不得不自己编写派生类实现虚接口，不仅冗长，性能上也会大打折扣。使用lambda表达式，可以直接在原地编写代码，十分方便。

此外，如果没有宏的封装，那么到此，也不过是用户手写几个lambda函数加上一堆又臭又长的调用，是 $\text{C++}$ 项目里的常见写法，但是这一对于用户来说就暴露了大量的细节。宏的包装真正对用户隐藏了他们不需要关心的部分，把核心逻辑的编写留给了用户。

1.1.2. How to Run the Kernels?

在 $\text{1.1.1}$ 中，我们从高层看到了 $\text{LC}$ 的核函数是如何编写并且被保存记录的，接下来我们需要思考，在程序运行的时候，这些核函数是如何被执行的。

这里摘取tests/test_helloworld.cpp中的部分代码：


Context context{argv[0]};
Device device = context.create_device(argv[1]);
Stream stream = device.create_stream();
Image<float> image{device.create_image<float>(PixelStorage::BYTE4, resolution)};
luisa::vector<std::byte> host_image(image.view().size_bytes());
Kernel2D kernel = [&]() {
    ...
};
Shader2D<> shader = device.compile(kernel);
stream << shader().dispatch(resolution) << image.copy_to(host_image.data()) << synchronize();

可以总结为：传入参数选择后端 -> 创建上下文 -> 创建设备 -> 创建指令流 -> 编译核函数并创建着色器 -> 将着色器分发至后端的执行流 -> 从后端收集数据并同步。

1.1.2.1. Context

“上下文”是一个非常常见的概念，但是在不同的应用中，它的含义也多有所差异，但是总体而言是用来代替全局变量，单例等等，管理一些高层次的状态的。“上下文”的主要作用是避免使用全局的状态，可以认为与“单例“模式截然相反。要理解为什么需要使用上下文，就要理解单例模式有哪些问题。

工程中的直觉告诉我们，在程序中保存一个全局状态是非常危险的。在StackOverflow的指引下，我找到了这篇博文：Singleton Anti-Pattern。

具体而言，保存全局状态其实就是单例模式的一种体现。而“单例”————也就是要求整个进程中只能有一个实例————的要求是非苛刻的。考虑到线程安全等等问题，要在现有的程序中保证这个假设就已经很困难了。而对于未来潜在的拓展需求，这种假设更是自缚手脚。因此，无论是为了编码的方便。还是为了程序的可维护性和可拓展性，我们都应该避免使用全局状态。这也就是为什么“上下文”是一个非常好的设计选择。

在 $\text{LC}$ 中，Context保存了运行时目录，校验层等等，也的确是我们所想的“上下文”的概念。

除此以外，在 $\text{LC}$ 的实现中还有个有趣的小模式： $\text{P-Impl}$ 模式。在include/luisa/runtime/context.h中，我们可以看到Context的定义：


class LC_RUNTIME_API Context {

private:
    luisa::shared_ptr<detail::ContextImpl> _impl;

public:
//  methods...
};

这是 $\text{C++}$ 独有的设计模式，有很多的好处。

这样可以在不改变接口的情况下，修改实现细节。
可以分离成员变量的定义，避免引入过多头文件导致的编译时间过长。
Context实例可以简单快速地拷贝，而无需担忧资源管理的问题。实际上这一点也得益于使用了shared_ptr。在我个人的实践中，如果使用 $\text{P-Impl}$ 模式，我更倾向于使用unique_ptr并禁用移动构造和拷贝构造，明确所有权，保证一份ContextImpl实例必然仅被一个Context实例持有。无论如何， $\text{P-Impl}$ 模式都极大地降低了资源管理的负担。
无论我们如何修改实现，Context编译出的二进制接口都是不变的，这样就可以保证二进制兼容性。如果不使用 $\text{P-Impl}$ 模式，那么我们修改成员变量就很可能导致类内存布局发生变化，进而导致二进制不兼容。

1.1.2.2. Device

Device的结构非常简单，它的成员变量只有一个_impl，这里看起来似乎有点儿 $\text{P-Impl}$ 模式的意思，不过据我观察，实际上Device更像是一个包装类。

class LC_RUNTIME_API Device {

public:
    using Deleter = void(DeviceInterface *);
    using Creator = DeviceInterface *(Context && , const DeviceConfig * );
    using Handle = luisa::shared_ptr<DeviceInterface>;

private:
    Handle _impl;   
}

Device本身的职责基本就是：

作为工厂，创建资源。它有不少create_XXX的方法，当然这些最后还是转发给DeviceInterface来分别实现的。
着色器程序的编译，加载和管理。

实际上Device最后还是会把各种方法转交给DeviceInterface来实现。DeviceInterface非常明显，是一个抽象类。诸如 $\text{Cuda, Vulkan, DirectX}$ 等等与显卡交互的接口其实都有一些共性，例如资源管理之类，这些共性可以被提取出来，也就是DeviceInterface类型。

例如：

// ...

[[nodiscard]] virtual BufferCreationInfo create_buffer(const Type *element, size_t elem_count) noexcept = 0;
  [[nodiscard]] virtual BufferCreationInfo create_buffer(const ir::CArc<ir::Type> *element, size_t elem_count) noexcept = 0;
  virtual void destroy_buffer(uint64_t handle) noexcept = 0;

  [[nodiscard]] virtual ResourceCreationInfo create_texture(
      PixelFormat format, uint dimension,
      uint width, uint height, uint depth,
      uint mipmap_levels, bool simultaneous_access) noexcept = 0;
  virtual void destroy_texture(uint64_t handle) noexcept = 0;

  // bindless array
  [[nodiscard]] virtual ResourceCreationInfo create_bindless_array(size_t size) noexcept = 0;
  virtual void destroy_bindless_array(uint64_t handle) noexcept = 0;

 // ...

我们可以猜想到，不同的后端就刚好对应于不同的DeviceInterface的派生类。这样的设计，使得 $\text{LC}$ 可以很方便地拓展到不同的后端，只需要实现一个新的DeviceInterface即可。

$\text{从字符串到后端}$

这里就不免涉及到一个问题：用户输入的是字符串（argv[1]）， $\text{LC}$ 是如何创建具体的后端的？

Context类作为工厂，调用create_device，接受后端的名字（也就是一个字符串），返回为用户选择的后端的Device对象。

Device Context::create_device(luisa::string_view backend_name_in, const DeviceConfig *settings, bool enable_validation) noexcept {
    luisa::string backend_name{backend_name_in};
    for (auto &c : backend_name) { c = static_cast<char>(std::tolower(c)); }
    auto &&m = _impl->load_backend(backend_name);
    auto interface = m.creator(Context{_impl}, settings);
    // ...
}

所以这里的关键是load_backend返回的东西。


struct BackendModule {
    using BackendDeviceNames = void(luisa::vector<luisa::string> &);
    DynamicModule module;
    Device::Creator *creator;
    Device::Deleter *deleter;
    BackendDeviceNames *backend_device_names;
};

[[nodiscard]] const BackendModule &load_backend(const luisa::string &backend_name) noexcept {
    if (std::find(installed_backends.cbegin(),
                    installed_backends.cend(),
                    backend_name) == installed_backends.cend()) {
        LUISA_ERROR_WITH_LOCATION("Backend '{}' is not installed.", backend_name);
    }
    std::scoped_lock lock{module_mutex};
    if (auto iter = loaded_backends.find(backend_name);
        iter != loaded_backends.cend()) { return *iter->second; }
    BackendModule m{.module = DynamicModule::load(
                        runtime_directory,
                        luisa::format("lc-backend-{}", backend_name))};
    LUISA_ASSERT(m.module, "Failed to load backend '{}'.", backend_name);
    m.creator = m.module.function<Device::Creator>("create");
    m.deleter = m.module.function<Device::Deleter>("destroy");
    m.backend_device_names = m.module.function<BackendModule::BackendDeviceNames>("backend_device_names");
    auto pm = loaded_backends
                    .emplace(backend_name, luisa::make_unique<BackendModule>(std::move(m)))
                    .first->second.get();
    return *pm;
}

BackendModule记录了一个后端的基本信息，包括创建和销毁，它们是通过查找动态链接库中的符号create和destroy来实现的。
这个方法会根据后端的名字，查找相应后端编译出的动态库，加载并记录其创建和销毁函数。

这样就可以获取到用户选择的DeviceInterface的工厂函数create（也就是BackendModule的Creator类型），从而创建一个Device对象。

我之前一直不清楚究竟怎么编写多后端支持的程序， $\text{LC}$ 的实现给了我非常大的启发。

1.1.2.3. Stream

我们关注它重载了哪些<<运算：

Delegate operator<<(luisa::unique_ptr<Command> &&cmd) noexcept;
Delegate operator<<(luisa::move_only_function<void()> &&f) noexcept;
template<typename T>
    requires std::is_rvalue_reference_v<T &&> && is_stream_event_v<T>
Stream &operator<<(T &&t) noexcept {
    luisa::invoke(std::forward<T>(t), device(), handle());
    return *this;
}
Stream &operator<<(CommandList::Commit &&commit) noexcept;
Stream &operator<<(Synchronize &&) noexcept;

先看前两个，Stream能够接受一个命令或者函数（注意这里的右值引用，意味着它转移了所有权），并产生一个委托Delegate。

第三个，如果向Stream推入一个事件，则执行该事件。后两个则是提交和同步，这两个类实际上都是空类，这里只是作 $\text{tag dispatch}$ 的作用。

Delegate的定义如下：

class LC_RUNTIME_API Delegate {
    private:
        Stream *_stream;
        CommandList _command_list;
    // ...
}

我们发现Stream可以一个一个地推入命令或者函数，但是Delegate里只有一个CommandList，因此我们猜测这个CommandList其实是个缓冲区，推入的命令不会立即执行。

Stream支持的<<运算，Delegate全部都支持，我们看其中<<(std::unique_ptr<Command> &&cmd)的实现：

1
2
3

if (!_command_list.callbacks().empty()) { _commit(); }
    _command_list.append(std::move(cmd));
    return std::move(*this);

所以这就很清楚了，Stream推入的命令会被缓存得到Delegate，向Delegate继续推入命令则会继续缓存。

Delegate对于推入事件的处理与Stream小有区别：

template<typename T>
            requires std::is_rvalue_reference_v<T &&> && is_stream_event_v<T>
Stream &operator<<(T &&t) && noexcept {
    _commit();
    luisa::invoke(std::forward<T>(t), _stream->device(), _stream->handle());
    return *_stream;
}

它会立刻提交缓存的所有命令，然后执行推入的事件。包括sycnrhonize也会先进行提交。这和一些并行计算API的设计是类似的。它返回的是Stream的引用，这意味着命令缓存被清空了。

$\text{为什么要使用Delegate?}$

为什么不直接给Stream添加一个CommandList缓冲区，而要使用Delegate呢？关于这个问题我问了核心开发者。

Delegate是一个临时对象，它生存到该语句结束，也就是没有新的东西被推入流中，通过给Delegate设计析构函数，可以实现在一条语句结束时自动提交。

1
2
3

Stream::Delegate::~Delegate() noexcept {
    _commit();
}

Delegate对象生存期只持续到一条语句结束，这样结束的时候就可以自动提交，算是一个小技巧。

1.1.2.4 Shader

着色器原义指在显卡上运行的程序。在这里自然也就拓展为一般意义下交给后端执行的程序。

这一部分反而从设计上暂时没有什么好讲的。Device通过compile方法，由具体后端实现将 $\text{LC}$ 的可调用对象编译为Shader，Shader的dispatch(dispatch_sizes)方法返回一个ShaderInvoke对象，可以直接被推入Stream。

1.2. Modules

$\text{LC}$ 库主要可分为如下几个重要模块：

$\text{Core}$ ：核心模块，主要包括一套专用的STL，并行原语，基础数学函数，日志等基本工具。
$\text{AST}$ ：抽象语法树模块，包括用于表示核函数 $\text{AST}$ 的基本类型。
$\text{Runtime}$ ： $\text{LC}$ 的运行时部分。
$\text{Backends}$ ：后端模块，包括 $\text{CUDA}$ ， $\text{DirectX}$ ， $\text{Metal}$ 和 $\text{CPU}$ 等不同的后端。
$\text{DSL}$ ：领域特定语言模块。
$\text{Tests}$ ：功能测试。

1.3. Summary

这个部分简单分析了一下 $\text{LC}$ 的基本工作流程，目前我们暂不关心 $\text{LC}$ 的后端，只是从 $\text{LC}$ 提供的最上层抽象，也就是它的 $\text{DSL}$ 模块提供的编程模型，来分析它的前端原理。