Mojo 可变参数(Variadics)完整指南:从参数列表、VariadicList 到 VariadicPack 与关键字参数

原创2026-09-10 11:36:341,638 阅读
文章标签:编程语言编译器标准库语言运行时人工智能大模型模型推理服务深度学习算子库高性能计算

Mojo 可变参数(Variadics)完整指南:从参数列表、VariadicList 到 VariadicPack 与关键字参数

本指南以 Mojo/proposals/variadics-design.md 为骨架,系统讲解 Mojo(Modular 平台)中可变参数特性的设计、语法与底层实现:编译期可变参数列表(TypeList/ParameterList)、运行时同构可变参数(VariadicList)、异构可变参数包(VariadicPack)以及可变关键字参数(StringDict,设计文档中写作 OwnedKwargsDict)。读完本文,你将掌握如何在函数签名中使用 *args、*args: *Ts 与 **kwargs,理解其借用/所有权语义、低层 lowering 方式,并能结合标准库源码与测试用例写出惯用的元编程与转发代码。

概述:Mojo 可变参数的四大形态

Mojo 的可变参数(variadic)能力横跨编译期与运行期两条主线:

形态 语法示例 背后类型 作用域
编译期类型列表 *Ts: AnyType TypeList def/struct/comptime 参数位
编译期值列表 *elts: Int ParameterList 同上
运行时同构参数 *args: Int VariadicList 运行时调用点
运行时异构参数包 *args: *Ts VariadicPack 运行时调用点
运行时关键字参数 **kwargs: Int StringDict[V](文档中的 OwnedKwargsDict[V]) 运行时调用点

这些特性统一采用 Python 风格的语法(*args 与 **kwargs),对熟悉 Python 的用户有天然的亲切感。编译期参数列表把静态已知数量的元素打包进一个绑定;运行时同构参数允许单个参数位接收任意数量的"同形状"实参;异构参数包用 *args: *Ts 语法把一组可变参数与一组可变类型一一对应;关键字参数则收集未被命名参数消费的 key=value 操作数。本文档当前状态为 Draft,是面向语言/标准库贡献者与进阶用户的实现级参考。

本文所有示例均可对照标准库测试 Mojo/stdlib/test/builtin/test_variadic.mojo(该文件是语法最新真相的来源)与核心实现 Mojo/stdlib/std/builtin/variadics.mojo。

Variadic parameter lists:编译期参数列表

可变 参数(parameter)列表出现在编译期参数位置(def、struct、comptime 等声明上),把静态已知数量的元素捆绑进一个绑定。与 Python 相同,Mojo 用前导 * 表示一组值或类型:

def takes_types[*Ts: AnyType](): ...
def takes_values[*elts: Int](): ...

Mojo 严格区分这两种情况,标准库分别以 TypeList 与 ParameterList 呈现(均定义在 Mojo/stdlib/std/builtin/variadics.mojo)。两者建立在同一个 KGEN 概念之上:一个 MLIR !kgen.param_list 值,其元素共享同一种编译期形态(要么全部满足某个 trait 的类型,要么全部是同一种类型的值)。它们是内建类型,正常使用无需显式导入——模块头部注释即声明:"These are Mojo built-ins, so you don't need to import them."

TypeList:类型序列的编译期操作

类型列表绑定一系列类型,声明时使用 trait 约束而非值类型,例如 *Ts: AnyType 或 *Ts: Movable。编译器通过 TypeList 操作将其暴露出来:对上面的示例,type_of(Ts) 就是 TypeList。从源码看,TypeList 底层类型正是 _MLIR.KGENParamListType[Self.Trait],即 !kgen.param_list<Trait>(见 variadics.mojo),并借助 #kgen.param_list.size、#kgen.param_list.get 等 MLIR 属性实现查询与索引。

TypeList 提供一系列实用操作:

  • 编译期 size / length:元素个数,如 tl.length;
  • 定长索引:Ts[i](可用于 comptime 循环,通过 __getitem_param__[idx] 实现);
  • 构造与变换助手:TypeList.of、splat、tabulate、map、reduce、filter_idx、contains、reversed、slice 等。

典型用途包括:对多个类型参数做 trait 谓词判断(参见 Mojo/stdlib/std/traits/movable.mojo),以及遍历/变换类型包的元编程。测试用例 test_variadic.mojo 中的 test_type_list_map_to_values、test_type_list_filter_idx_*、test_type_list_reduce_idx 等覆盖了这些操作的常见组合。

ParameterList:同类型编译期值序列

值列表绑定一系列共享同一类型 T 的编译期 值,例如 *args: Int 或 *names: String。它变成一个元素类型为 T 的 ParameterList。可以编译期迭代或索引;ParameterList.get_span() 在需要指针线性布局时,把元素物化为连续常量数组背后的 Span——源码显示它通过 #pop.variadic_to_array 把 param_list 展平成扁平数组、再映射为运行期常量并取首元素指针(见 variadics.mojo)。ParameterList.of、splat 等构造器与 TypeList 的故事一一对应。

两个类型刻意保持平行设计(文档注释表示未来在 Mojo metatype 故事成熟后有望统一):TypeList.map_to_values 把每个元素类型经编译期生成器映射为同构值的 ParameterList——"类型包驱动值包"是常见模式。测试 test_variadic.mojo 的 test_type_list_map_to_values 正是此用法。

需要注意的是:与参数列表不同,参数列表目前不支持异构列表或关键字参数列表,文档明确表示"若未来有足够需求支撑其复杂度,可能加入"。

Homogeneous variadic arguments:运行时同构可变参数

上一节讨论的是声明上的 参数;本节讨论运行时调用点的 实参:一个参数位接受任意数量的实参,但所有实参必须是同一类型 T。表面语法同样是前导 *,但现在它命名的是一个 VariadicList(而非 ParameterList)。

签名与基本用法

同构可变参数形如 *args: Int 或 *parts: String。在 callee 内部,args 可以像一个小型顺序集合使用:len(args)、args[i]、for x in args;当 T 满足 Writable 时还有 args.write_to(s) 之类助手来渲染元组形态:

def print_ints(*values: Int):
    for i in range(len(values)):
        print(values[i])

def main():
    print_ints(10, 20, 30)  # 三个独立实参,一个可变参数绑定

只关心值不关心索引时,迭代器路径同样自然:

def sum_ints(*values: Int) -> Int:
    var total = 0
    for v in values:
        total += v
    return total

callee 收到的是 VariadicList

*args: T 语法映射到标准库类型 VariadicList(variadics.mojo)。它的内部字段是:

var _value: Span[Self._EltPointerType, ImmUntrackedOrigin]  # 元素指针的 Span

也就是说,它携带的是一个 指向各实参引用的指针的 Span,而不是 T 值的密集数组。这种布局让 callee 能以 args[idx] 暴露正确的借用或移动语义(__getitem__ 通过 self._value.unsafe_ptr()[unsafe_offset=idx][] 解引用),同时让胶水对象小到足以按值传递,还允许非 Movable 的值通过可变参数传递。文档给出的粗略 lowering 图景(仅为示意名,非精确 MLIR):

// Caller:  foo(a, b, c)   with   def foo(*xs: T)

1. 编译器照常为每个实参分配存储(栈槽、寄存器或 ABI 要求的形式)。

2. 构造一个临时数组,其第 i 项是指向第 i 个实参的 pointer-to-ref
   (所有元素的指针/引用 MLIR 类型相同)。

3. 调用 foo,传入一个 VariadicList:其内部 Span 指向该数组,长度为实参个数。

4. VariadicList.__init__(由该数组隐式构造)把 POP 数组转成元素指针的 Span;
   __getitem__ 通过该 span 索引并加载引用。

源码佐证了第 4 步:隐式 __init__[size, container_origin] 接收编译器生成的元素指针数组引用,用 pop.array.gep 取首元素指针后 unsafe_bitcast 为元素指针类型,再包成 Span(unsafe_ptr=..., length=size)(见 variadics.mojo)。

所以可变参数束永远是指针的 span——即使 T 是 trivial 类型也是如此。阅读性能文档时需记住:可变参数在 caller 侧仍是独立对象,列表只是一层用于统一索引的间接层。

借用与拥有实参

默认情况下 *args: T 借用每个元素。当需要取得所有权时(例如 T 是线性类型,或想用 consume_elements 移出),使用 var *args: T。VariadicList 用 is_owned 参数跟踪该状态;置位时 __deinit__ 反向遍历列表、逐个销毁元素,与常规实参析构顺序一致(见 variadics.mojo):

@__parameter
def destroy_elem(_idx: Int, var arg: ExplicitDelOnly):
    arg^.destroy()

def take_owned_linear(var *args: ExplicitDelOnly):
    args^.consume_elements[destroy_elem]()

# Caller 传临时对象;callee 逐个消费。
take_owned_linear(ExplicitDelOnly(5), ExplicitDelOnly(10))

该示例直接取自 test_variadic.mojo 的 test_variadic_list_linear_type。consume_elements 调用处的 ^ 选择可变参数束的 owned 视图。

consume_elements 与想要 var 元素的 API

许多"下沉一串值"的 API 与 List 列表字面量构造器形状一致:var *values: Self.T 加上 values^.consume_elements[...] 把每个实参移入新分配的存储。从实现看,consume_elements(deinit self, elt_handler) 只在 Self.is_owned 时可用(where 约束即报错信息"consume_elements may only be called on owned variadic lists"),它用 __get_address_as_owned_value 逐个转移所有权给处理器闭包(见 variadics.mojo)。值得注意的是源码注释:这里刻意不用 Pointer.unsafe_take_pointee,因为它要求元素是 Movable,而 VariadicList 明确不需要这一前提。

打印与调试

当 element_type 满足 Writable 时,VariadicList 实现 write_to/write_repr_to:内部 _write_elements 循环拼接 (a, b, c) 形态(is_repr=True 时逐元素调用 write_repr_to),write_repr_to 还会套上 VariadicList[Int] 类型名(见 variadics.mojo)。这就是为什么测试期望 write_to 输出 (1, 2, 3)、write_repr_to 输出 VariadicListInt, Int(2), Int(3)))。测试用例见 test_variadic_list_write_to 与 test_variadic_list_write_repr_to(test_variadic.mojo)。

Variadic packs:异构可变参数包

"可变参数包"(variadic pack)是同构 VariadicList 的异构对应物。不再用 *args: T(单一静态元素类型),而是把 类型参数包 与 实参包 配对,实参类型取自该包:

def callee*Ts: Writable raises:
    ...

callee 收到一个 VariadicPack(定义于 variadics.mojo)。与 VariadicList 一样,它是 RegisterPassable 且参与所有权(is_owned、var *args: *Ts、consume_elements、__del__),但内部表示是类似 Tuple 的异构值:其底层 MLIR 类型是 !lit.ref.pack<:param_list<Trait> ... isParamPack> 形态的 !kgen.struct(见 variadics.mojo),每个槽位可能对应不同大小、不同 ABI 的具体类型。

为什么包需要 comptime for

每个实参槽位可能是 不同 的具体类型,大小与 ABI 各异。因此包更接近元组而非数组:不存在单一的 T 可用于 args[runtime_idx]。索引通过 __getitem_param__[index] 暴露,使用 编译期 索引(实现中用 lit.ref.pack.extract 抽取槽位),编译器才能为每个位置生成正确的加载指令。

这正是惯用代码用 comptime for 而非运行期 for i in range(len(args)) 遍历包的原因:

def count_many_things*ArgTypes: Intable -> Int:
    var total = 0

    comptime for i in range(args.__len__()):
        # 每个 args[i] 都有来自 *ArgTypes 的不同具体类型。
        total += Int(args[i])

    return total

def main():
    print(count_many_things(Int8(5), UInt32(11), Int(12)))  # 28

该示例改编自 VariadicPack 的 docstring(variadics.mojo)。关键点:循环变量 i 是编译期参数,因此每个 args[i] 都单独 monomorphize。VariadicPack.__len__ 直接返回 Self.Ts.length(类型包长度,编译期已知)。

Writable 包与转发

当每个元素类型都满足 Writable 时,包实现 write_to 及相关操作:

def helper*Ts: Writable raises:
    var s = String()
    args.write_to(s)
    # 对 (1, "hello", True) -> "(1, hello, True)"

def forwarder*Ts: Writable raises:
    helper(*args)  # splat 原样转发包

# Caller:
forwarder(1, "hello", True)

转发模式 callee(*args) 在 test_variadic_pack_forwarding、单元素变体 test_variadic_pack_forwarding_single_element、空包变体 test_variadic_pack_forwarding_empty 以及多跳版本 test_variadic_pack_forwarding_through_two_levels 中均有覆盖(test_variadic.mojo)。空包与单元素包的转发方式完全相同(对应测试中的 forwarder() 与 forwarder(42))。

SomeTypeList 语法糖

当只需要"任意数量的类型,且每个都满足 Trait"时,可以不引入显式 *Ts 绑定,直接在 def 上命名实参包:

def foo(*args: *SomeTypeList[Writable]) raises:
    var s = String()
    args.write_to(s)

SomeTypeList 是定义于 Mojo/stdlib/std/builtin/anytype.mojo 的编译期别名,把 Some[T: Trait] 的思想扩展到受同一 trait 约束的整个 TypeList,在可变调用点尤其有用。对应测试为 test_variadic_pack_some(test_variadic.mojo)。

与 Tuple 的关系

Tuple 是拥有异构序列的典型结构体。其构造函数接收一个与其元素类型列表对齐的可变参数包:

# 概念示意(见 Mojo/stdlib/std/builtin/tuple.mojo):
struct Tuple*element_types: Movable:
    def __init__(out self, var *args: *Self.element_types):
        ...

也就是说,VariadicPack 本质上就是运行期被降级为结构体的异构元组,编译器在编译期掌握每个元素的精确类型,以生成正确的内存布局与访问代码。

Variadic keyword arguments:可变关键字参数

可变 关键字 参数是 Python **kwargs 的运行时对应物。callee 可以接受任意数量的额外关键字实参,且这些实参的 值 共享同一个类型 V。键始终是运行时 String(调用点写下的关键字名)。与可变参数包不同,这里没有异构值类型列表:**kwargs: Int 意味着每个传入值都必须是 Int,而非混合类型元组。

表面语法是参数名前导 **,且必须位于签名末尾、其他参数之后:

def variadic_kwargs(a: Int, b: Int, *args: Int, c: Int, d: Int, **kwargs: Int):
    pass

def print_nicely(**kwargs: Int):
    for item in kwargs.items():
        print(item.key, "=", item.value)

def main():
    print_nicely(a=7, y=8)

每个函数只允许一个 ** 参数,且必须带类型注解:例如 **kwargs: Int,不允许裸写 **kwargs。解析器侧的支持可见于 Mojo/test/mojo-parser/decls/variadic_kwargs.mojo 与 Mojo/lib/MojoParser/Signatures.cpp(后者负责签名中 kw_vararg 槽位的解析)。

callee 收到的是 StringDict(文档中的 OwnedKwargsDict)

设计文档写作时,**kwargs: V 以 var kwargs: OwnedKwargsDict[V] 传入。当前仓库中该类型已更名为 StringDict[V],定义于 Mojo/stdlib/std/collections/dict.mojo,其 docstring 明确写着"用于向函数传递拥有的可变关键字实参的容器",且"用户不应直接实例化它"——编译器在调用点构建、传入 callee。这一更名在 Mojo/docs/site/releases/v1.0.0.md 的发布说明中有记录(OwnedKwargsDict → StringDict)。StringDict 内部包装了一个 Dict[String, V, default_comp_time_hasher],对外暴露字典风格接口:len、in、__getitem__(支持按 String 或 ImmStringSpan 查找,后者免分配)、__setitem__、keys()、values()、items()、find、pop 等,另有 deinit_with 用于值非 Deinitable 时的显式析构。

关键字可变参数以 拥有(var)方式传递,因为字典通常为单个调用点构建。不能显式写 read 或 mut 约定:

# 暂不支持。
def borrowed_kwargs(mut **kwargs: Int): ...

在 callee 内部,kwargs 拥有该字典及其插入的值,这与调用 lowering 用 _insert 把每个关键字操作数转入字典的方式一致。

调用 lowering

文档给出的粗略 lowering 图景(仅为示意名,非精确 MLIR):

// Caller:  foo(x=9, stuff=8)   with   def foo(**kwargs: Int)

1. 分配一个空的 OwnedKwargsDict[Int](局部临时对象)。

2. 对调用点每个关键字操作数:
   - 把键物化为编译期 String 字面量。
   - 求值值表达式。
   - 调用 OwnedKwargsDict::_insert(dict, key, value),把值的所有权转入字典。

3. 把填好的字典作为拥有的 **kwargs 实参传给 foo。

重载解析收集所有未绑定到前面具名参数的关键字操作数,路由到 kw_vararg 槽位;若 callee 没有 **kwargs 参数,这些操作数即报错。位置参数 *args 与 **kwargs 可出现在同一签名中,各自"吃掉"自己种类的操作数。

用 **kwargs^ 转发

要把整个关键字束转发给另一个可变关键字 callee,使用双星解包形式并配合 ^ 转移所有权:

def pass_kwargs(**kwargs: Int):
    takes_int_variadic_kwargs_multiline(**kwargs^)

不加 ^ 时,转发会尝试拷贝字典并因 StringDict 不可隐式拷贝而失败;加上 ^ 后,caller 的字典被移入内层调用。这是 callee(*args) splat 可变参数包的关键字对应物。

当前尚不支持从普通 Dict[String, V] 解包(例如 print_nicely(**my_dict));今天只有从另一个 **kwargs 绑定转发可用。

泛型推断

当函数在值类型上泛型时,关键字实参可以像位置实参一样驱动类型推断:

def infers_from_kwargsT: SomeTrait:
    pass

# T 由关键字值推断为 MemOnly。
infers_from_kwargs(y=MemOnly(), z=s)

推断出的元素类型必须满足声明的 trait 约束,且兼容转入 StringDict._insert(owned、movable 值)。测试 Mojo/stdlib/test/collections/test_dict.mojo 的 test_owned_kwargs_dict 与第 1944 行的 test_owned_kwargs_dict_linear 覆盖了可变关键字参数在 callee 内部暴露的字典 API 表面,以及线性值(非 Copyable)的转移规则。

当前限制

以下缺口是首个版本的有意取舍,亦可参见 Mojo 手册 functions 章节:

  • 仅支持同构值:所有关键字共享一个值类型 V;
  • 键始终是 String,没有类型化键变体;
  • 始终 owned:不支持 read **kwargs 或 mut **kwargs;
  • 调用点不支持从普通 Dict 解包;
  • 不支持编译期可变关键字 参数(声明上的 **kwparams: ...);
  • 值必须是 Movable(并满足 StringDict 的约束);非拥有的线性值遵循与其他 owned 调用实参相同的转移规则。

实践建议与测试参照

  • 选择形态:只需要同类型的一串值时用 *args: T(VariadicList);需要每个槽位类型不同、类型由 trait 约束时用 *args: *Ts(VariadicPack),并务必用 comptime for 遍历;需要收集任意关键字键值对时用 **kwargs: V(StringDict[V])。
  • 所有权:要消费(移动)实参,一律写成 var *args: T / var *args: *Ts,再经 consume_elements 或 **kwargs^ 转移;借用场景保持默认 read 语义。
  • 元编程:TypeList 的 map/reduce/filter_idx/contains 与 ParameterList.get_span() 是编译期类型/值计算的主力工具;SomeTypeList[Trait] 是可变调用点上约束类型包的最简写法。
  • 验证语法:任何语法细节以测试文件 Mojo/stdlib/test/builtin/test_variadic.mojo 为准,其中 test_variadic_*、test_type_list_*、test_parameter_list_*、test_variadic_pack_*、test_dynamic_variadic_pack 等用例分别覆盖了本文讨论的各条路径;关键字字典相关验证在 Mojo/stdlib/test/collections/test_dict.mojo。

文档仍标注为 Draft,TypeList/ParameterList 的统一、**kwparams 编译期参数等能力属于未来演进方向;实际能力请以当前仓库代码与测试为准。

登录后查看全文
mojo