Slang 自动微分装饰器全解析从[Differentiable]到自定义导数与 PrimalSubstitute【免费下载链接】slangMaking it easier to work with shaders项目地址: https://gitcode.com/GitHub_Trending/sl/slang导读本文围绕 Slang 编译器自动微分auto-diff系统的装饰器体系展开详细讲解[Differentiable]、[TreatAsDifferentiable]、[ForwardDerivative]/[BackwardDerivative]及其Of反向引用形式、[PrimalSubstitute]/[PrimalSubstituteOf]等属性在 AST、语义检查和 IR 降级三个层面的完整行为。读完本文你将掌握如何在 Slang 中声明可微函数、为 interface 需求提供可微实现、为复杂或带副作用的函数手工指定导数以及如何为无函数体的硬件内建指令提供参考实现从而让自动微分管线在真实着色器代码中可靠运行。本文为 Slang 编译器贡献者视角的设计文档对应的完整设计脉络可参考 自动微分设计文档总览、类型系统设计、IR Pass 总览 以及 数学基础。一、为什么自动微分需要装饰器Slang 的自动微分目标是给定一个着色器函数编译器自动生成其关于输入参数的导数函数。整个流程从用户显式调用fwd_diff(fn)或bwd_diff(fn)开始由一系列 IR PassForwardDiffTranscriber、BackwardDiffTranscriberBase、DiffUnzipPass等详见 IR Pass 总览完成导数合成。但在合成导数之前编译器必须回答几个关键问题哪些函数允许被微分如果任意函数都能被微分那么库函数、带副作用的函数都会被盲目地尝试变换产生错误或性能灾难。当一个函数结构过于复杂、或包含全局/共享内存访问时如何手工提供更好的导数当一个函数根本没有函数体如硬件内建指令时导数从何而来装饰器decorator正是用来回答这些问题的一整套标注机制。它们在语义检查阶段被解析为 AST 属性节点在降级lowering到 IR 阶段被翻译为对应的 IR Decoration从而指导自动微分各 Pass 的行为。二、[Differentiable]声明函数的可微性2.1 基本语义与 IR 降级[Differentiable]属性用于将函数标记为可微。这是整个自动微分体系的基石性约束自动微分过程只会处理显式标记为[Differentiable]的函数。其他所有函数都被视为不可微从可微函数调用这类函数时调用会被原样复制copy as-is不做任何变换。只有带[Differentiable]的方法会在导数数据流分析 Passderivative>[Differentiable] float sqr(float x) { return x * x; } // 前向模式求导 DifferentialPairfloat fwd_sqr fwd_diff(sqr); // 反向模式求导 void bwd_sqr(inout DifferentialPairfloat dpx, float dOut) bwd_diff(sqr);2.2 废弃的旧装饰器需要特别注意的是[Differentiable]最初被实现为两个独立的装饰器[ForwardDifferentiable]与[BackwardDifferentiable]分别表示支持前向/反向自动微分变换。如今这两个装饰器已废弃deprecated推荐做法是只使用[Differentiable]。在源码中仍保留ForwardDifferentiableAttributesource/slang/slang-ast-modifier.h#L1898与BackwardDifferentiableAttribute的定义以及一个辅助属性HasTrivialForwardDerivativeAttributesource/slang/slang-ast-modifier.h#L1891但新代码不应再使用它们。2.3 调用约束与no_difffwd_diff与bwd_diff不能直接调用未带[Differentiable]标记的方法会产生编译错误。如果在一个[Differentiable]方法内部必须调用非[Differentiable]方法则该调用必须用no_diff()操作包裹这一约束由导数数据流分析 Pass 强制实施float nonDiffFunc(float x) { /* ... */ } [Differentiable] float differentiableFunc(float x) { /* ... */ } [Differentiable] float main(float x) { // 非可微调用必须显式 no_diff否则数据流分析会报错 return no_diff(nonDiffFunc(x)) * differentiableFunc(x); }no_diff只能直接用于函数调用在 IR 层面它转变为TreatAsDifferentiableDecoration表示该函数调用不会产生导数。三、[Differentiable]用于 interface 需求[Differentiable]还可以装饰 interface 的方法需求requirement。由于 interface 没有具体实现编译器无法直接对需求本体做导数变换因此处理方式有所不同。3.1 处理流程整个流程大致如下语义检查阶段在检查 interface 需求方法时checkCallableDeclCommon位于 source/slang/slang-check-decl.cpp检查该方法是否带[Differentiable]属性。生成派生需求声明如果带该属性编译器会构造一组新的方法声明插入到同一个 interface 中ForwardDerivativeRequirementDecl前向模式导数需求BackwardDerivativeRequirementDecl反向模式导数需求 两者的函数类型由原方法的类型经导数类型变换得到。建立引用在原方法中插入一个新成员DerivativeRequirementReferenceDecl用于引用上述新声明。IR 降级DerivativeRequirementReferenceDecl被转换为自定义导数引用——在原方法的需求键requirement key上添加OpBackwardDerivativeDecoration(deriv-fn-req-key)与OpForwardDerivativeDecoration(deriv-fn-req-key)两个 IR Decoration。3.2 示例检查与降级前后对比原文档给出了如下示例interface IFoo { [Differentiable] float bar(float); }; // After checking lowering interface IFoo_after_checking_and_lowering { [BackwardDerivative(bar_bwd)] [ForwardDerivative(bar_fwd)] float bar(float); void bar_bwd(inout DifferentialPairfloat, float); DifferentialPairfloat bar_fwd(DifferentialPairfloat); };可以看到[Differentiable]需求在检查和降级后等价于在原需求上附加了前向与反向自定义导数引用并生成了两个导数需求的函数签名反向导数bar_bwd接收一个inout DifferentialPairfloat参数同时携带 primal 与微分值以及一个额外的float输出导数参数返回void前向导数bar_fwd接收DifferentialPairfloat参数返回DifferentialPairfloat。3.3 实现约束签名的一部分所有满足该 interface 的具体类型其对应实现也必须声明为[Differentiable]这样编译器才能为这些实现合成与 interface 签名匹配的导数实现。从这个意义上说[Differentiable]是函数签名的一部分一个[Differentiable]的 interface 需求只能由一个[Differentiable]的函数实现来满足。这一点与 types.md 中的类型系统设计 相互印证——例如IDifferentiable接口定义于 source/slang/core.meta.slang通过dadd/dzero/dmul三个操作和关联类型Differential定义了可微类型的代数结构而带[Differentiable]的 interface 方法需求则要求在 witness table 上存在对应导数的查找入口。四、[TreatAsDifferentiable]零导数占位在大型代码库中一个 interface 可能有多个实现逐个把所有实现都标记为[Differentiable]并不现实——尤其是某些实现使用了需要额外考量的 hack 或绕过手段暂时不能标注。此时可以使用[TreatAsDifferentiable]装饰器AST 节点TreatAsDifferentiableAttribute定义于 source/slang/slang-ast-modifier.h#L1883继承自DifferentiableAttributeIROpTreatAsDifferentiableDecoration。它的作用是指示自动微分各 Pass构造一个空函数其导数部分返回 0或 0 等价物。这样一个[TreatAsDifferentiable]函数的签名就能匹配[Differentiable]需求而无需真正产生导数。interface ICompute { [Differentiable] float eval(float x); }; // 该实现暂时无法提供导数但签名上仍需满足 interface 需求 [TreatAsDifferentiable] struct HackImpl : ICompute { float eval(float x) { // 一些特殊处理暂不提供真实导数 return x; } };需要强调的是[TreatAsDifferentiable]是绕过而非解决——它相当于显式声明此处的导数为零因此应谨慎使用避免在导数必须非零的位置静默引入错误。五、自定义导数装饰器在很多情况下手工指定某个方法的导数代码比让自动微分 Pass 从函数体合成更好。典型场景有两类函数体过于复杂但存在更简单、数学上等价的求值方式常见于sin(x)、arccos(x)等内建函数。函数涉及全局/共享内存访问合成导数代码可能引发竞态条件或因过度同步而非常慢。为此Slang默认假定全局内存访问不可微并要求用户或 core 模块定义带有不同导数语义的独立访问器。Slang 前端为此提供两套装饰器方向装饰器作用从 primal 函数引用导数[ForwardDerivative(fn)]/[BackwardDerivative(fn)]为被装饰函数指定自定义前向/反向导数函数从导数函数引用 primal 函数[ForwardDerivativeOf(fn)]/[BackwardDerivativeOf(fn)]声明被装饰函数是某 primal 函数的自定义导数对应 AST 节点与 IR DecorationForwardDerivativeAttribute/BackwardDerivativeAttribute均继承自UserDefinedDerivativeAttribute见 source/slang/slang-ast-modifier.h#L1913 与 source/slang/slang-ast-modifier.h#L1950IR 层为OpForwardDervativeDecoration/OpBackwardDerivativeDecorationForwardDerivativeOfAttribute/BackwardDerivativeOfAttribute继承自DerivativeOfAttribute见 source/slang/slang-ast-modifier.h#L1933 与 source/slang/slang-ast-modifier.h#L1958。5.1 两套装饰器的定位差异[ForwardDerivative(fn)]/[BackwardDerivative(fn)]直接在 primal 函数上声明其自定义导数。[ForwardDerivativeOf(fn)]/[BackwardDerivativeOf(fn)]在导数函数上反向声明我是谁的导数。这在不修改原模块的情况下为既有方法提供导数非常有用。典型实例是 source/slang/diff.meta.slang 中的attribute_syntax声明如 source/slang/diff.meta.slang#L143 的[ForwardDerivativeOf(function)]它通过大量[ForwardDerivativeOf(...)]/[BackwardDerivativeOf(...)]标注为hlsl.meta.slang中定义的 core 模块函数如mul、dot、cross、transpose等见 source/slang/diff.meta.slang#L1445 起的大量条目提供手工导数而无需改动hlsl.meta.slang本身。5.2 IR 降级位置当降级到 IR 时这些引用被放置在**目标函数即 primal 函数**上。也就是说无论使用[ForwardDerivative(fn)]还是[ForwardDerivativeOf(fn)]形式最终两套装饰都会降级在 primal 函数上——区别只在于书写位置和引用建立的方向。5.3 适用范围泛型、结构体方法与重载这些装饰器同样适用于泛型定义的方法generic methods结构体方法struct methods重载方法overloaded methods——与函数调用类似重载解析复用ResolveInvoke基础设施完成。// 为结构体方法提供自定义反向导数 struct MyOp { [Differentiable] [BackwardDerivative(myOp_bwd)] float apply(float x) { /* 复杂实现 */ } void myOp_bwd(inout DifferentialPairfloat dpx, float dOut) { // 手工导数 dpx DifferentialPairfloat(dpx.p, dOut * 3.0 * dpx.p * dpx.p); } };5.4 自定义导数签名的检查为保证用户提供的导数函数与期望签名一致并在存在多个重载时解析出正确的方法编译器会检查自定义导数函数与 primal 函数翻译后版本的签名匹配性。这一过程发生在checkDerivativeAttribute()/checkDerivativeOfAttribute()中实现位于 source/slang/slang-check-decl.cpp相关的checkDerivativeAttributeImpl/checkDerivativeOfAttributeImpl也在同一文件中。检查过程复用了ResolveInvoke基础设施构造一个临时的 invoke 表达式用一组**虚构参数imaginary arguments**调用用户提供的导数函数这些参数的类型依据 primal 方法的翻译后类型translated type构造如果ResolveInvoke解析成功则认为提供的导数签名匹配这一方式同时自动支持重载方法解析、泛型类型适配、类型隐式转换type coercion。六、[PrimalSubstitute]与[PrimalSubstituteOf]为无函数体的内建指令提供参考实现6.1 问题背景自定义导数解决的是有函数体但不想自动合成的问题。而[PrimalSubstitute]系列解决的是相反的问题我们希望编译器从函数体自动合成导数但目标函数根本没有函数体可供翻译。这种情况在硬件内建指令hardware intrinsic中非常常见——例如纹理采样texture sampling与插值interpolation操作它们被降级为映射到硬件单元的专用 op-code没有可翻译的软件实现。但这些操作确实存在参考性的软件实现可以用来产生导数。6.2 装饰器语义[PrimalSubstitute(ref-fn)]AST 节点PrimalSubstituteAttribute见 source/slang/slang-ast-modifier.h#L1967IR 层为OpPrimalSubstituteDecoration[PrimalSubstituteOf(orig-fn)]AST 节点PrimalSubstituteOfAttribute见 source/slang/slang-ast-modifier.h#L1976是反向引用形式声明被装饰函数是某 primal 函数的替换实现。其核心价值在于primal 计算路径使用快速硬件内建指令而导数计算路径使用由参考实现合成的导数。这样用户代码既能享受硬件加速的 primal 计算又能获得正确的导数。6.3 完整示例原文档给出的示例完整展现了这一模式[PrimalSubstitute(sampleTexture_ref)] float sampleTexture(TexHandle2D tex, float2 uv) { // Hardware intrinsics } float sampleTexture_ref(TexHandle2D tex, float2 uv) { // Reference SW implementation. } void sampleTexture_bwd(TexHandle2D tex, inout DifferentialPairfloat2 dp_uv, float dOut) { // Backward derivate code synthesized using the reference implementation. }这里sampleTexture是硬件内建版本函数体为空或仅含内建调用sampleTexture_ref是参考软件实现。当自动微分管线遇到sampleTexture时会用sampleTexture_ref替代它来合成导数如sampleTexture_bwd而 primal 调用本身仍走硬件路径。6.4 底层实现原理[PrimalSubstitute(fn)]的实现相对直接当转录器transcriber即实现AutodiffTranscriberBase的前向/反向导数合成器详见 IR Pass 总览被要求合成某个函数的导数时它会先检查该函数上是否存在OpPrimalSubstituteDecoration若存在则在继续导数合成之前将当前函数替换为 substitute 函数。这一先替换、再合成的顺序保证了导数合成的输入永远是那个有真实函数体的参考实现从而让现有的一整套自动微分 Pass 无需为硬件内建指令做任何特判。七、装饰器在自动微分整体架构中的位置将上述装饰器放到整个自动微分管线中看它们构成了一个清晰的声明层用户代码装饰器标注 │ [Differentiable] / [ForwardDerivative] / [PrimalSubstitute] ... ▼ 语义检查slang-check-decl.cppcheckDerivativeAttribute / checkDerivativeOfAttribute / interface 需求展开生成 DerivativeRequirement 声明 ▼ IR 降级OpBackwardDifferentiableDecoration / OpForwardDervativeDecoration / OpBackwardDerivativeDecoration / OpPrimalSubstituteDecoration / OpTreatAsDifferentiableDecoration ▼ 自动微分 IR PassForwardDiffTranscriber / BackwardDiffTranscriberBase / DiffUnzipPass / 转置 / Checkpointing-Recomputation / Extraction见 ir-overview.md ▼ 导数函数DifferentialPairT 签名见 types.md各装饰器的职责可以总结如下装饰器解决的问题关键 IR Decoration[Differentiable]声明函数可微是自动微分的作用域边界OpBackwardDifferentiableDecoration[TreatAsDifferentiable]无导数实现但需匹配可微签名导数为 0OpTreatAsDifferentiableDecoration[ForwardDerivative]/[BackwardDerivative]手工指定前向/反向导数函数OpForwardDervativeDecoration/OpBackwardDerivativeDecoration[ForwardDerivativeOf]/[BackwardDerivativeOf]在不修改原模块的前提下反向声明导数归属同上降级在 primal 函数上[PrimalSubstitute]/[PrimalSubstituteOf]为无函数体的内建指令提供可合成的参考实现OpPrimalSubstituteDecoration这些装饰器与 类型系统设计文档 中描述的IDifferentiable接口、DifferentialPairT、导数数据流分析、以及 IR Pass 总览 中的转录器机制紧密配合共同构成了 Slang 完整、可扩展的自动微分前端声明体系。八、实战要点总结默认不微分只有显式[Differentiable]的函数才会被变换其他调用原样复制。想让非可微调用进入可微函数必须用no_diff()显式包裹。interface 需求的可微性是签名的一部分[Differentiable]的需求只能由[Differentiable]的实现满足编译器会自动生成对应前向/反向导数需求声明。手工导数优先考虑复杂函数体、全局/共享内存访问的场景优先用[ForwardDerivative(fn)]/[BackwardDerivative(fn)]手工指定无法修改原模块时用[ForwardDerivativeOf(fn)]/[BackwardDerivativeOf(fn)]在外部模块如diff.meta.slang提供导数。硬件内建指令用 PrimalSubstituteprimal 走硬件、导数走参考实现通过[PrimalSubstitute(ref-fn)]声明即可编译器会先替换再合成。谨慎使用 TreatAsDifferentiable它产生零导数仅适用于明确知道导数应为零的占位场景。【免费下载链接】slangMaking it easier to work with shaders项目地址: https://gitcode.com/GitHub_Trending/sl/slang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考