diff --git a/pfcc/paddle-code-reading/IR_Dialect/README.md b/pfcc/paddle-code-reading/IR_Dialect/README.md index 33dbe6a1a..60d2e9bd7 100644 --- a/pfcc/paddle-code-reading/IR_Dialect/README.md +++ b/pfcc/paddle-code-reading/IR_Dialect/README.md @@ -9,4 +9,6 @@ + [【方案设计】IR 底层基础类型系统设计文档✨New✨](./basic_concepts.md) + [【方案设计】IR 顶层模型结构表示设计文档✨New✨](./ir_program.md) + [【代码约定】IR 代码相关约定](./code_convention.md) -+ [【方案设计】IR Program Translator设计文档 🚀](./program_translator.md) \ No newline at end of file ++ [【方案设计】IR Program Translator设计文档 🚀](./program_translator.md) ++ [【方案设计】控制流设计文档 🚀](./control_flow.md) ++ [【社区贡献文档】PIR源码阅读指南 🚀](./first_step.md) \ No newline at end of file diff --git a/pfcc/paddle-code-reading/IR_Dialect/control_flow.md b/pfcc/paddle-code-reading/IR_Dialect/control_flow.md new file mode 100644 index 000000000..335a83c8f --- /dev/null +++ b/pfcc/paddle-code-reading/IR_Dialect/control_flow.md @@ -0,0 +1,1084 @@ +> 版本,作者,时间 + +| 版本 | 作者 | 时间 | 主要更新 | +| ---- | ------ | ---------- | ------------------ | +| v1.0 | 王明冬 | 2023.08.10 | 初版 | +| v1.1 | 王明冬 | 2023.08.22 | 添加评审时会谈纪要 | + +# 一、概要 + +## 1、相关背景 + +IR升级的核心功能目标是设计一套能够由多方共用且优于当前的IR基础设施。 + +本文是IR升级《控制流模块》的设计评审文档,是在[类型系统](https://ku.baidu-int.com/knowledge/HFVrC7hq1Q/yKeL8Lljko/5K6Iojo8fU/hklvoOpOf9amJs)、[模型结构](https://ku.baidu-int.com/knowledge/HFVrC7hq1Q/yKeL8Lljko/6UmOO2EkH2/7a4ab3fcfc0e42)的基础上继续设计的。 + +本文的功能目标是完成控制流算子的抽象。 + +控制流能够设定特定的顺序执行计算任务,帮助构建更加灵活和复杂的模型。在模型中引入控制流后可以让计算图中某些节点循环执行任意次数,也可以根据条件判断选择某些节点不执行。 + +许多深度学习模型依赖控制流进行训练和推理,基于递归神经网络和强化学习的模型就依赖于循环递归关系和依据输入数据状态条件执行计算。 + +```cpp +def cond(i, ten): + return i < ten + +def body(i, ten): + i = i + 1 + return [i, ten] + +i = paddle.full(shape=[1], fill_value=0, dtype='int64') # loop counter +ten = paddle.full(shape=[1], fill_value=10, dtype='int64') # loop length +i, ten = paddle.static.nn.while_loop(cond, body, [i, ten]) +``` + +以该while循环为例,目前paddle框架对其在计算图中的描述如下: + +```cpp +{ // block 0 + var fill_constant_1.tmp_0 : LOD_TENSOR.shape(1,).dtype(int64).stop_gradient(True) + var fill_constant_3.tmp_0 : LOD_TENSOR.shape(1,).dtype(int64).stop_gradient(True) + var tmp_0 : LOD_TENSOR.shape(1,).dtype(bool).stop_gradient(True) + var _generated_var_0 : STEP_SCOPES) + + {Out=['fill_constant_1.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 3, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [1], str_value = 0, value = 0.0, with_quant_attr = False) + {Out=['fill_constant_3.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 3, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [1], str_value = 10, value = 10.0, with_quant_attr = False) + {Out=['tmp_0']} = less_than(inputs={X=['fill_constant_1.tmp_0'], Y=['fill_constant_3.tmp_0']}, axis = -1, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) + {Out=['tmp_0', 'fill_constant_1.tmp_0'], StepScopes=['_generated_var_0']} = while(inputs={Condition=['tmp_0'], X=['fill_constant_3.tmp_0', 'fill_constant_1.tmp_0']}, is_test = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], sub_block = block[1], with_quant_attr = False) +} +{ // block 1 + var tmp_1 : LOD_TENSOR.shape(1,).dtype(int64).stop_gradient(True) + var tmp_2 : LOD_TENSOR.shape(1,).dtype(bool).stop_gradient(True) + + {Out=['tmp_1']} = scale(inputs={ScaleTensor=[], X=['fill_constant_1.tmp_0']}, bias = 1.0, bias_after_scale = True, op_device = , op_namescope = /, op_role = 0, op_role_var = [], scale = 1.0, with_quant_attr = False) + {Out=['tmp_2']} = less_than(inputs={X=['tmp_1'], Y=['fill_constant_3.tmp_0']}, axis = -1, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) + {Out=['fill_constant_1.tmp_0']} = assign(inputs={X=['tmp_1']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) + {Out=['tmp_0']} = assign(inputs={X=['tmp_2']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) +} +``` + +这存在两个问题: + +1. while算子的部分输入和也是输出,意味着计算图在主block里面存在有向环。[关于“有向环问题”总结](https://ku.baidu-int.com/knowledge/HFVrC7hq1Q/pKzJfZczuc/hPr_u_N_Lx/LNKjesZGk_qCD5) 这种情况对控制流相关分析优化显然是不太友好的。 +2. cond 函数体在主block和子block中同时存在,相应的输入输出变量也会出现两份。目前是因此cond 函数的函数体比较简单,所以影响不大,但cond函数体变复杂以后,这种实现显然是不合理的。 + +新IR设计借鉴了MLIR的设计思想,坚持SSA原则。显然跟当前paddle的这种实现方式是相悖的。因此,本方案在不违反IR整体设计原则的情况下,综合考虑组合算子、自动微分、动态shape等各模块的需求,对控制流算子进行设计实现。 + +## 2、功能目标 + +- 完成新IR体系中控制流算子(If、While)的IR表示. +- 在不破坏IR设计原则的前提下,描述控制流算子的反向IR实现。 + +# 二、意义 + + + +# 三、竞品对照 + +目前,MLIR、TorchMLIR等都不涉及到控制流的反向描述,而控制流设计的最复杂之处就在于反向设计。 + +和普通算子不同,控制流反向的实现需要用到前向的子block中的局部变量。如果直接引用,显然是跟作用域原则相悖的。 + +因此,竞品的参考意义不大。本章只是对MLIR的前向控制流算子的设计进行简要描述。 + +## 1、MLIR的结构化控制流算子定义 + +MLIR对控制流算子的描述是通过scf方言来描述的。 + +### 1.1 WhileOp + +~~~cpp +//===----------------------------------------------------------------------===// +// WhileOp +//===----------------------------------------------------------------------===// + +def WhileOp : SCF_Op<"while", + [DeclareOpInterfaceMethods, + RecursiveMemoryEffects]> { + let summary = "a generic 'while' loop"; + let description = [{ + This operation represents a generic "while"/"do-while" loop that keeps + iterating as long as a condition is satisfied. There is no restriction on + the complexity of the condition. It consists of two regions (with single + block each): "before" region and "after" region. The names of regions + indicates whether they execute before or after the condition check. + Therefore, if the main loop payload is located in the "before" region, the + operation is a "do-while" loop. Otherwise, it is a "while" loop. + + The "before" region terminates with a special operation, `scf.condition`, + that accepts as its first operand an `i1` value indicating whether to + proceed to the "after" region (value is `true`) or not. The two regions + communicate by means of region arguments. Initially, the "before" region + accepts as arguments the operands of the `scf.while` operation and uses them + to evaluate the condition. It forwards the trailing, non-condition operands + of the `scf.condition` terminator either to the "after" region if the + control flow is transferred there or to results of the `scf.while` operation + otherwise. The "after" region takes as arguments the values produced by the + "before" region and uses `scf.yield` to supply new arguments for the + "before" region, into which it transfers the control flow unconditionally. + + A simple "while" loop can be represented as follows. + + ```mlir + %res = scf.while (%arg1 = %init1) : (f32) -> f32 { + // "Before" region. + // In a "while" loop, this region computes the condition. + %condition = call @evaluate_condition(%arg1) : (f32) -> i1 + + // Forward the argument (as result or "after" region argument). + scf.condition(%condition) %arg1 : f32 + + } do { + ^bb0(%arg2: f32): + // "After" region. + // In a "while" loop, this region is the loop body. + %next = call @payload(%arg2) : (f32) -> f32 + + // Forward the new value to the "before" region. + // The operand types must match the types of the `scf.while` operands. + scf.yield %next : f32 + } + ``` + + A simple "do-while" loop can be represented by reducing the "after" block + to a simple forwarder. + + ```mlir + %res = scf.while (%arg1 = %init1) : (f32) -> f32 { + // "Before" region. + // In a "do-while" loop, this region contains the loop body. + %next = call @payload(%arg1) : (f32) -> f32 + + // And also evaluates the condition. + %condition = call @evaluate_condition(%arg1) : (f32) -> i1 + + // Loop through the "after" region. + scf.condition(%condition) %next : f32 + + } do { + ^bb0(%arg2: f32): + // "After" region. + // Forwards the values back to "before" region unmodified. + scf.yield %arg2 : f32 + } + ``` + + Note that the types of region arguments need not to match with each other. + The op expects the operand types to match with argument types of the + "before" region; the result types to match with the trailing operand types + of the terminator of the "before" region, and with the argument types of the + "after" region. The following scheme can be used to share the results of + some operations executed in the "before" region with the "after" region, + avoiding the need to recompute them. + + ```mlir + %res = scf.while (%arg1 = %init1) : (f32) -> i64 { + // One can perform some computations, e.g., necessary to evaluate the + // condition, in the "before" region and forward their results to the + // "after" region. + %shared = call @shared_compute(%arg1) : (f32) -> i64 + + // Evaluate the condition. + %condition = call @evaluate_condition(%arg1, %shared) : (f32, i64) -> i1 + + // Forward the result of the shared computation to the "after" region. + // The types must match the arguments of the "after" region as well as + // those of the `scf.while` results. + scf.condition(%condition) %shared : i64 + + } do { + ^bb0(%arg2: i64) { + // Use the partial result to compute the rest of the payload in the + // "after" region. + %res = call @payload(%arg2) : (i64) -> f32 + + // Forward the new value to the "before" region. + // The operand types must match the types of the `scf.while` operands. + scf.yield %res : f32 + } + ``` + + The custom syntax for this operation is as follows. + + ``` + op ::= `scf.while` assignments `:` function-type region `do` region + `attributes` attribute-dict + initializer ::= /* empty */ | `(` assignment-list `)` + assignment-list ::= assignment | assignment `,` assignment-list + assignment ::= ssa-value `=` ssa-value + ``` + }]; + + let arguments = (ins Variadic:$inits); + let results = (outs Variadic:$results); + let regions = (region SizedRegion<1>:$before, SizedRegion<1>:$after); + + let builders = [ + OpBuilder<(ins "TypeRange":$resultTypes, "ValueRange":$operands, + "function_ref":$beforeBuilder, + "function_ref":$afterBuilder)> + ]; + + let extraClassDeclaration = [{ + using BodyBuilderFn = + function_ref; + + OperandRange getSuccessorEntryOperands(std::optional index); + ConditionOp getConditionOp(); + YieldOp getYieldOp(); + Block::BlockArgListType getBeforeArguments(); + Block::BlockArgListType getAfterArguments(); + }]; + + let hasCanonicalizer = 1; + let hasCustomAssemblyFormat = 1; + let hasVerifier = 1; +} +~~~ + +如上所述,while op包含两个region, 称为before region和after region。 before region以scf.conditon算子结尾,如果conditon的输入条件为true, 就会把控制流传递到after region。否则,将控制流返回到父op, 表示执行结束。 + +after region以scf.yield算子结尾,表示将控制流传递到before region。 + +显然,当主要循环开销在before region时, while op等价于c++的 "do-while"语句。当主要循环开销在after region时,while op等价于c++的while语句。 + +### 1.2 IfOp + +~~~cpp +//===----------------------------------------------------------------------===// +// IfOp +//===----------------------------------------------------------------------===// + +def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods, + SingleBlockImplicitTerminator<"scf::YieldOp">, RecursiveMemoryEffects, + NoRegionArguments]> { + let summary = "if-then-else operation"; + let description = [{ + The `scf.if` operation represents an if-then-else construct for + conditionally executing two regions of code. The operand to an if operation + is a boolean value. For example: + + ```mlir + scf.if %b { + ... + } else { + ... + } + ``` + + `scf.if` may also produce results. Which values are returned depends on + which execution path is taken. + + Example: + + ```mlir + %x, %y = scf.if %b -> (f32, f32) { + %x_true = ... + %y_true = ... + scf.yield %x_true, %y_true : f32, f32 + } else { + %x_false = ... + %y_false = ... + scf.yield %x_false, %y_false : f32, f32 + } + ``` + + The "then" region has exactly 1 block. The "else" region may have 0 or 1 + block. In case the `scf.if` produces results, the "else" region must also + have exactly 1 block. + + The blocks are always terminated with `scf.yield`. If `scf.if` defines no + values, the `scf.yield` can be left out, and will be inserted implicitly. + Otherwise, it must be explicit. + + Example: + + ```mlir + scf.if %b { + ... + } + ``` + + The types of the yielded values must match the result types of the + `scf.if`. + }]; + let arguments = (ins I1:$condition); + let results = (outs Variadic:$results); + let regions = (region SizedRegion<1>:$thenRegion, + MaxSizedRegion<1>:$elseRegion); + + let skipDefaultBuilders = 1; + let builders = [ + OpBuilder<(ins "TypeRange":$resultTypes, "Value":$cond)>, + OpBuilder<(ins "TypeRange":$resultTypes, "Value":$cond, + "bool":$addThenBlock, "bool":$addElseBlock)>, + OpBuilder<(ins "Value":$cond, "bool":$withElseRegion)>, + OpBuilder<(ins "TypeRange":$resultTypes, "Value":$cond, + "bool":$withElseRegion)>, + OpBuilder<(ins "Value":$cond, + CArg<"function_ref", + "buildTerminatedBody">:$thenBuilder, + CArg<"function_ref", + "nullptr">:$elseBuilder)>, + ]; + + let extraClassDeclaration = [{ + OpBuilder getThenBodyBuilder(OpBuilder::Listener *listener = nullptr) { + Block* body = getBody(0); + return getResults().empty() ? OpBuilder::atBlockTerminator(body, listener) + : OpBuilder::atBlockEnd(body, listener); + } + OpBuilder getElseBodyBuilder(OpBuilder::Listener *listener = nullptr) { + Block* body = getBody(1); + return getResults().empty() ? OpBuilder::atBlockTerminator(body, listener) + : OpBuilder::atBlockEnd(body, listener); + } + Block* thenBlock(); + YieldOp thenYield(); + Block* elseBlock(); + YieldOp elseYield(); + }]; + let hasFolder = 1; + let hasCanonicalizer = 1; + let hasCustomAssemblyFormat = 1; + let hasVerifier = 1; +} +~~~ + +如上所示,if op包含两个region, 称为then region和else region。if op只有一个输入conditon,如果输入为true,执行then region, 否则,执行else region。 + +then region和else region的输出都跟if op的输出匹配。如果if op没有输出,那么else region可以为空。 + +# 四、设计思路与实现方案 + +新IR通过Operation、Region、Block三者的循环嵌套来表示结构化控制流。 + +一个Op会包含0个或多个Region, 一个Region会包含0个或多个Block, 一个Block里面包含了0个或多个Operation。 三者循环嵌套包含,用来描述复杂的模型结构。 + +## 1、 基础组件 + +### 1.1 Block + +- **Block** + +新IR的Block等价于基本块, 里面包含了一个算子列表(std::list), 用来表示该基本块的计算语意。 + +```cpp +^block: + %a = "pd.feed" () ... + %b = "pd.feed" () ... + %c = pd.add(%a, %b) ... + pd.fetch(%c) .... +``` + +样例1就是一个简单的block样例。 + +- **BlockArgument** + +Block可以包含一个形参列表(std::vector), 来表示执行该Block所需要的参数数量和类型。 + +```cpp +^block (%a :tensor<...>, %b:tensor<...>): + %c = pd.add(%a, %b) ... + pd.fetch(%c) .... +``` + +样例2是一个简单的带BlockArgument的block样例。 它将样例1的通过feed算子来获取的两个变量%a和%b通过BlockArgument来描述。这意味着,控制流在进入该block之前,必须给%a和%b绑定变量。 + +- **BlockOperand** + +Block可以被封装为BlockOperand(类似Value和OpOperand的关系)。作为终止符一种特殊的操作数, 称为后继块(successor)。终止符算子是指一类有特殊语意的算子,他们可以作为基本块的最后一个op。比如: return、fetch、branch等等。 + +```cpp +#reigon +{ + ^condition_block (%cond): + %1 = pd.constant(1) + %2 = pd.constant(2) + pd.condition_branch %cond, ^then_block(%1), else_block(%2) + ^then_block(%val_1): + pd.return %val_1 + ^else_block(%val_2): + pd.return %val_2 +} +``` + +样例3是一个Block作为终止符算子的操作数的一个例子。 样例中,pd.condition_branch接受三个操作数:%cond、%1、%2的同时,接受两个blockOperand: then_block和else_block,它的语意时,如果%cond的值为True, 就将控制流传递到then_block, 同时将%1作为参数传递给then_block的BlockArgument。否则,就将控制流传递到else_block, 同时将%2作为参数传递给else_block的BlockArgument。 + +注: 在控制流之前, 一个operation由它的输入、输出、属性以及类型信息构成。 加入控制流以后,一个operation的内容包含:它的输入(OpOperand)、输出(OpResult)、属性(AttributeMap)、后继块(BlockOperand)、region组成。 新增了后继块和region。 + +**当Block的最后一个算子执行结束时,根据块内最后一个算子(终止符算子)的语意,控制流会有两种去处:** + +1. **进入同Region的另外一个Block, 该Block一定是终止符算子的后继块。** +2. **返回该Block的父Region, 表示该Region的一次执行的结束。** + +### 1.2 Region + +Region里面包含了一个Block列表(std::vector), 第一个Block(如果存在的话),称为该Region的入口块。 + +**与基本块不同,Region存在一个最显著的约束是:Region内定义的Value只能在该Region内部使用,Region的外面不允许使用****。**这个约束与我们执行期的scope比较类似。 + +当控制流进入一个region, 相当于创建了一个新的子scope, 当控制流退出该region时,该子scope中定义的所有变量都可以回收。 + +**控制流进入Region, 一定会首先进入该Region的入口块。**因此,Region的参数用入口块参数即可描述,不需要额外处理。 + +**当Region的一次执行结束,控制流由子Block返回到该Region时,控制流会有两种去处:** + +1. **进入同Op的某一个Region(可能是自己)。** +2. **返回该Region的父Op,表示该Op的一次执行的结束。** + +具体去处由该Region的父Op的语意决定。 + +## 2、 控制流算子 + +### 2.1 辅助工具 + +为了实现分支语意以及反向函数,需要预定义一些特殊的类型和算子,用来辅助描述控制流算子。 + +这些辅助类型和算子目前先定义在cf(control flow) dialect中。后续有必要的话,可以将部分类型下沉到builtin dialect中。 + +------ + +考虑到WhileOp的反向需要访问前向的局部变量,而且反向的循环迭代与前向是相反的,特提供以下类型和算子: + +- **cf.StackType** + +StackType表示一个支持先进后出的栈类型。 该类型不需要参数。 + +```cpp +class IR_API StackType : public Type { + DECLARE_TYPE_UTILITY_FUNCTOR(StackType, TypeStorage); + ....... +} +``` + +- **cf.CreateStackOp** + +CreateStackOp算子的语意是创建一个空栈。 该算子没有输入、 没有属性、 输出一个类型为 StackType的value。 + +```cpp +// %0是一个stack类型变量 +%0 = cf.create_stack() {} : ()->cf.stack +``` + +- **cf.PushBackOp** + +PushBackOp算子的语意是将一个变量进栈。该算子接受两个输入,第一个为StackType的Value,表示栈, 第二个输入为要被进栈的变量。没有属性,没有输出。 + +```cpp +// %1对应的变量被压栈到了%0中 +cf.push_back(%0, %1){}: (cf.stack, tensor<...>)->() +``` + +- **cf.PopBackOp** + +PopBackOp算子的语意是将栈末尾的变量弹出来。该算子接受一个类型为StackType的输入。 没有属性,有一个输出,表示栈中被弹出的变量。 + +```cpp +// 从%0对应的栈中pop_back出一个变量,记为%2 +%2 = cf.pop_back(%0) {}: cf.stack -> tensor<...> +``` + +- **cf.IsEmptyOp** + +IsEmptyOp算子的语意是判断栈是否为空。该算子接受一个类型为StackType的输入。没有属性,有一个bool类型的输出,表示栈是否为空。 + +```cpp +// 判断栈是否为空,返回bool变量 +%cond = cf.is_empty(%0) {}:cf.stack -> bool +``` + +------ + +同时,我们新增一些终止符算子作为控制流中block的最后一个算子: + +- **cf.YieldOp** + +Yield算子接受可变输入,没有输出和属性。它的语意是将控制流和输入变量传递给父region。类似于函数的return语句。 + +```cpp +// %1,%2等表示该region执行的返回值。 +cf.yield(%1, %2, ....) +``` + +在pd.CondOp的子block中,会用cf.yield算子来描述返回值。 + +- **cf.CondYieldOp** + +CondYield算子比Yield算子多了一个bool类型的条件变量。它的父region会根据该变量,进行一次分支决策。 + +```cpp +// 将控制流传递到父region. 父region会根据%cond的值,进行分支。 +// 对于while op的body region而言,如果%cond为True, 他会将控制流传递给body_region, %0、%1...会被传递给body_region当参数。否则,将控制流返回while_op, %0、%1...会被当作while_op的输出。 +cf.cond_yield (%cond, %0, %1, ...) +``` + +### 2.2 IfOp(CondOp) + +IfOp应该包含两个或三个Region; + +两个region:then_region和else_region对应无反向场景; + +如果存在反向,会额外增加一个init region,同时会增加一个表示变量栈的stack输出。(这种场景会在2.3.1:IfOp的反向实现中进行描述) + +IfOp只有一个输入condition。 输出是可变的。 (**这是因为子block可以直接访问父block的变量,CondOp的内部block的前驱也是唯一的,因此没必要设置参数,在用的地方直接访问原变量即可**) + +如果IfOp的输出为空,那么else_region可以为空。 + +否则,else_region和then_region一样,都必须包含一个不带参数的Block, 分别表示then和else的分支。这两个block都必须以cf.yield算子结尾(如果输出为空,cf.yield可以省略)。 + +cf.yield算子接收可变输入,没有输出。 语意是将它的输入转发给父Op当作输出,它的输入个数与IfOp的输出匹配。 + +IfOp的Verify函数功能包括: + +1. 如果包含两个region: + 1. then_region只包含一个block, 该block + 1. 参数为空 + 2. 以cf.yield算子结尾,且cf.yield算子的输入跟IfOp的输出的数量和类型相匹配。 + 2. 如果IfOp的输入为空,那么else_region也可以为空。否则else_region一定只包含一个block,该bock + 1. 参数为空 + 2. 以cf.yield算子结尾,且cf.yield算子的输入跟IfOp的输出的数量和类型相匹配。 +2. 如果包含三个region: + 1. init_region只包含一个block, 该block只包含两个算子: cf.create_stack 和 cf.yield. + 2. then_region只包含一个block, 该block + 1. 只有一个stack类型的block_argument。 + 2. 以cf.yield算子结尾,且cf.yield算子的输入跟CondOp的输出的数量和类型相匹配。 + 3. 如果else_region为空,那么CondOp的非stack输出也一定为空。否则else_region也一定包含一个block,该bock + 1. 只有一个stack类型的block_argument。 + 2. 以cf.yield算子结尾,且cf.yield算子的输入跟CondOp的输出的数量和类型相匹配。 + +```cpp +# +# pseudocode: +# if 0.1 < 0.23: +# return 1, True +# else: +# return 3, 2 +# +def true_func(): + a = paddle.full(shape=[1, 2], dtype='int32',fill_value=1) + b = paddle.full(shape=[2, 3], dtype='bool', fill_value=True) + return a, b + + +def false_func(): + a = paddle.full(shape=[3, 4], dtype='float32',fill_value=3) + b = paddle.full(shape=[4, 5], dtype='int64', fill_value=2) + return a, b + +x = paddle.full(shape=[1], dtype='float32', fill_value=0.1) +y = paddle.full(shape=[1], dtype='float32', fill_value=0.23) +pred = paddle.less_than(x=x, y=y, name=None) +ret = paddle.static.nn.cond(pred, true_func, false_func) +``` + +对应的新IR下的算子描述为: + +```cpp +%x = pd.full(....) +%y = pd.full(....) +%cond = pd.less_than(x, y) +%ret1, %ret2 = pd.if(%cond) { + %1 = pd.full(....) + %2 = pd.full(...) + cf.yield(%1, %2) + } else { + %1 = pd.full(....) + %2 = pd.full(...) + cf.yield(%1, %2) + } +``` + +当前paddle框架对应的描述为: + +```cpp +{ // block 0 + var fill_constant_1.tmp_0 : LOD_TENSOR.shape(1,).dtype(float32).stop_gradient(True) + var fill_constant_3.tmp_0 : LOD_TENSOR.shape(1,).dtype(float32).stop_gradient(True) + var less_than_0.tmp_0 : LOD_TENSOR.shape(1,).dtype(bool).stop_gradient(True) + var _generated_var_0 : LOD_TENSOR.shape(1, 2).dtype(int32).stop_gradient(True) + var _generated_var_1 : LOD_TENSOR.shape(2, 3).dtype(bool).stop_gradient(True) + var _generated_var_2 : STEP_SCOPES) + var logical_not_0.tmp_0 : LOD_TENSOR.shape(1,).dtype(bool).stop_gradient(True) + var _generated_var_3 : LOD_TENSOR.shape(3, 4).dtype(float32).stop_gradient(True) + var _generated_var_4 : LOD_TENSOR.shape(4, 5).dtype(int64).stop_gradient(True) + var _generated_var_5 : STEP_SCOPES) + var cast_0.tmp_0 : LOD_TENSOR.shape(1,).dtype(int32).stop_gradient(True) + var _generated_var_6 : LOD_TENSOR.shape(-1, -1).dtype(int32).stop_gradient(True) + var _generated_var_7 : LOD_TENSOR.shape(-1, -1).dtype(bool).stop_gradient(True) + + {Out=['fill_constant_1.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 5, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [1], str_value = 0.1, value = 0.10000000149011612, with_quant_attr = False) + {Out=['fill_constant_3.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 5, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [1], str_value = 0.23, value = 0.23000000417232513, with_quant_attr = False) + {Out=['less_than_0.tmp_0']} = less_than(inputs={X=['fill_constant_1.tmp_0'], Y=['fill_constant_3.tmp_0']}, axis = -1, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) + {Out=['_generated_var_1', '_generated_var_0'], Scope=['_generated_var_2']} = conditional_block(inputs={Cond=['less_than_0.tmp_0'], Input=[]}, is_scalar_condition = True, op_device = , op_namescope = /, op_role = 0, op_role_var = [], sub_block = block[1], with_quant_attr = False) + {Out=['logical_not_0.tmp_0']} = logical_not(inputs={X=['less_than_0.tmp_0']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) + {Out=['_generated_var_4', '_generated_var_3'], Scope=['_generated_var_5']} = conditional_block(inputs={Cond=['logical_not_0.tmp_0'], Input=[]}, is_scalar_condition = True, op_device = , op_namescope = /, op_role = 0, op_role_var = [], sub_block = block[2], with_quant_attr = False) + {Out=['cast_0.tmp_0']} = cast(inputs={X=['less_than_0.tmp_0']}, in_dtype = 0, op_device = , op_namescope = /, op_role = 0, op_role_var = [], out_dtype = 2, use_mkldnn = False, with_quant_attr = False) + {Out=['_generated_var_6']} = select_input(inputs={Mask=['cast_0.tmp_0'], X=['_generated_var_3', '_generated_var_0']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) + {Out=['_generated_var_7']} = select_input(inputs={Mask=['cast_0.tmp_0'], X=['_generated_var_4', '_generated_var_1']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) +} +{ // block 1 + var fill_constant_5.tmp_0 : LOD_TENSOR.shape(1, 2).dtype(int32).stop_gradient(True) + var fill_constant_7.tmp_0 : LOD_TENSOR.shape(2, 3).dtype(bool).stop_gradient(True) + + {Out=['fill_constant_5.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 2, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [1, 2], str_value = 1, value = 1.0, with_quant_attr = False) + {Out=['fill_constant_7.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 0, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [2, 3], str_value = 1.0, value = 1.0, with_quant_attr = False) + {Out=['_generated_var_0']} = assign(inputs={X=['fill_constant_5.tmp_0']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) + {Out=['_generated_var_1']} = assign(inputs={X=['fill_constant_7.tmp_0']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) +} +{ // block 2 + var fill_constant_9.tmp_0 : LOD_TENSOR.shape(3, 4).dtype(float32).stop_gradient(True) + var fill_constant_11.tmp_0 : LOD_TENSOR.shape(4, 5).dtype(int64).stop_gradient(True) + + {Out=['fill_constant_9.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 5, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [3, 4], str_value = 3.0, value = 3.0, with_quant_attr = False) + {Out=['fill_constant_11.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 3, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [4, 5], str_value = 2, value = 2.0, with_quant_attr = False) + {Out=['_generated_var_3']} = assign(inputs={X=['fill_constant_9.tmp_0']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) + {Out=['_generated_var_4']} = assign(inputs={X=['fill_constant_11.tmp_0']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) +} +``` + +显然,paddle框架会将控制流的ture分支和false分支分别插入一个condition_block op。再通过select_input op对两个condition_block op的输出进行选择。这是因为在当前框架,一个Op只能包含一个Block,所以遇见IfOp这种算子,必须拆分两个Op。 + +### 2.3 WhileOp + +WhileOp包含两个或三个region。 + +如果是两个region,那就是cond_region和body_region。 + +如果是三个region, 那就是init_region、condition_region、body_region。 + +init_region只做一件事:创建一个stack,将其和输入参数一起, 转发给cond_region,该stack会在循环中压栈一些局部变量,并作为输出传递到更高层次的作用域,提供给反向算子中使用。如果不考虑反向,那么stack输出可以省略,相应的,init_region也可以省略。 WhileOp直接拿cond_region作为入口执行也是可以的。 + +------ + +以下是一个只考虑前向的WhileOp的样例: + +```cpp +def cond(i, ten): + return i < ten + +def body(i, ten): + i = i + 1 + return [i, ten] + +i = paddle.full(shape=[1], fill_value=0, dtype='int64') # loop counter +ten = paddle.full(shape=[1], fill_value=10, dtype='int64') # loop length +i, ten = paddle.static.nn.while_loop(cond, body, [i, ten]) +``` + +对应的新IR下,只观翻译后的初始版本描述为: + +```cpp +%i = pd.full(...) +%ten = pd.full(...) +%i_2, %ten2 = pd.while(%i, %ten) { + // cond region + ^bb0 (%arg1, %arg2): + %cond = pd.less_than(%arg1, %arg2) + cf.cond_yield (%cond, %arg1, %arg2) + } do { + // body region + ^bb1(%arg1, %arg2): + %1 = pd.const(1) + %i_3 = pd.add(%arg1, %1) + cf.yield (%i_3, %arg2) + } +``` + +可以通过数据流分析发现,cond_region和body_region的第二个块参数始终绑定的是%ten, 因此可以进一步优化为: + +```cpp +%i = pd.full(...) +%ten = pd.full(...) +%i_2 = pd.while(%i) { + // cond region + ^bb0(%arg1): + %cond = pd.less_than(%arg1, %ten) + cf.cond_yield (%cond, %arg1) + } do { + // body region + ^bb1(%arg1): + %1 = pd.const(1) + %i_3 = pd.add(%arg1, %1) + cf.yield (%i_3) + } +``` + +再进一步,可以发现body_region里面的%1其实每一轮都是一样的,因此,可以将其提升到循环外: + +```cpp +%i = pd.full(...) +%1 = pd.const(1) +%ten = pd.full(...) +%i_2 = pd.while(%i) { + // cond_region + ^bb0(%arg1): + %cond = pd.less_than(%arg1, %ten) + cf.cond_yield (%cond, %arg1) + } do { + // body_region + ^bb1(%arg2): + %i_3 = pd.add(%arg2, %1) + cf.yield (%i_3) + } +``` + +当前paddle框架对应的描述为: + +```cpp +{ // block 0 + var fill_constant_1.tmp_0 : LOD_TENSOR.shape(1,).dtype(int64).stop_gradient(True) + var fill_constant_3.tmp_0 : LOD_TENSOR.shape(1,).dtype(int64).stop_gradient(True) + var tmp_0 : LOD_TENSOR.shape(1,).dtype(bool).stop_gradient(True) + var _generated_var_0 : STEP_SCOPES) + + {Out=['fill_constant_1.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 3, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [1], str_value = 0, value = 0.0, with_quant_attr = False) + {Out=['fill_constant_3.tmp_0']} = fill_constant(inputs={ShapeTensor=[], ShapeTensorList=[], ValueTensor=[]}, dtype = 3, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], place_type = -1, shape = [1], str_value = 10, value = 10.0, with_quant_attr = False) + {Out=['tmp_0']} = less_than(inputs={X=['fill_constant_1.tmp_0'], Y=['fill_constant_3.tmp_0']}, axis = -1, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) + {Out=['tmp_0', 'fill_constant_1.tmp_0'], StepScopes=['_generated_var_0']} = while(inputs={Condition=['tmp_0'], X=['fill_constant_3.tmp_0', 'fill_constant_1.tmp_0']}, is_test = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], sub_block = block[1], with_quant_attr = False) +} +{ // block 1 + var tmp_1 : LOD_TENSOR.shape(1,).dtype(int64).stop_gradient(True) + var tmp_2 : LOD_TENSOR.shape(1,).dtype(bool).stop_gradient(True) + + {Out=['tmp_1']} = scale(inputs={ScaleTensor=[], X=['fill_constant_1.tmp_0']}, bias = 1.0, bias_after_scale = True, op_device = , op_namescope = /, op_role = 0, op_role_var = [], scale = 1.0, with_quant_attr = False) + {Out=['tmp_2']} = less_than(inputs={X=['tmp_1'], Y=['fill_constant_3.tmp_0']}, axis = -1, force_cpu = False, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) + {Out=['fill_constant_1.tmp_0']} = assign(inputs={X=['tmp_1']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) + {Out=['tmp_0']} = assign(inputs={X=['tmp_2']}, op_device = , op_namescope = /, op_role = 0, op_role_var = [], with_quant_attr = False) +} +``` + +显然,当前版本由于只支持一个子block, 因此会将cond_block的代码复制一遍,一份放在主block,一份放在子block。 + +------ + +### 2.3 对backward的支持 + +下文对于backward的设计支持,基于以下思路: + +1. 问:前向算子和反向算子是否应该处于同一个block?如何描述他们的嵌套关系? + +答:对于最顶层block中的算子,它的前反向处于同一个block。但是对于子作用域中的算子,前反向应该处于不同的block中,但前反向block的辈分(嵌套层级)一定是相同的。 + +------ + +![img](./img/while_op_detail.jpg) + +如图1所示,假设program包含了while_op_1, while_op_1包含了while_op_2, while_op_2嵌套包含了while_op_3, .......嵌套包含了while_op_n....... + +则在经过了backward pass之后,program会包含while_op_1和while_op_1_grad, while_op_1_grad嵌套包含了while_op_2_grad, while_op_2_grad嵌套包含了while_op_3_grad, .......嵌套包含了while_op_n_grad........。 + +while_op_n和while_op_n_grad位于不同的block, 但二者的辈分(离program的嵌套层数)是相同的。 + +------ + +1. 问:如果前反向算子不在同一个block, 当反向算子需要访问前向的输入输出时,如何在不破坏作用域原则(父作用域不允许直接访问子作用域变量)的前提下,构造计算图的拓扑关系? + +答:新IR严格遵守作用域原则:当反向block需要访问前向block中定义的变量时,必须先以返回值传递的方式,将该变量沿着前向嵌套作用层级,逐层向上传递到他们的公共祖先作用域;再以参数传递的方式,沿着反向block嵌套作用层级,逐层向下传递到反向block。 + +通过我们在2.1节中描述的栈类型变量、push_back算子、pop_back算子,我们只需要传递一个栈变量,在前向block中,按照定义顺序,将反向block中需要用到的局部变量依次压栈,在反向block中,依次出栈即可。因为反向block的执行顺序和前向block的执行顺序是互逆的,因此,这儿使用先进后出的栈来进行值传递。 + +backward pass 或者每个op的反向创建的接口:需要保证,在每个前向block中压栈变量的数量和顺序和反向block中出栈变量的数量和顺序是匹配的。 + +------ + +![img](./img/while_op_detail2.jpg) + +因为子作用域可以访问父作用域中定义的变量, 图2所举的例子中,while_3_op的子block可以访问的变量范围是: program的主block、while_1_op的子block、while_2_op的子block。 + +在构建while_3_op_grad算子的子block时,program的主block本来就是它的顶层作用域,可以直接访问。 + +while_1_op的子block中的变量都通过入栈出栈的方式,对偶到了while_1_op_grad的子block中,而while_1_op_grad的子block也是while_3_op_grad的祖先block, 可以直接访问。 + +类似的,while_2_op的子block中变量也被对偶到了while_2_op_grad的子block中,显然,这是while_3_op_grad父block, 因此可以直接访问。 + +综上,while_3_op_grad所在的位置可以访问到while_3_op所能访问的所有变量(可能是对偶变量),因此,while_3_op_grad可以在不违反作用域原则的情况下,成功构建。 + +------ + +1. 问:反向block需要访问前向block中的局部变量,为了实现该目的,我们设计了压栈出栈的实现方式。当添加完反向,训练结束以后,如何进行推理部署呢?或者说如何移除其中的压栈算子呢? + +答:只需要在裁剪反向的pass的最后,追加一个特殊的类似DCE的Pass。比如while_op, 在裁剪了反向以后,我们就会发现,while_op的代表局部变量栈的输出变量已经没有消费者了。这个输出本来就是optional的,所以可以直接将该输出移除即可。 相对应的,该while_op里面的子block的终止符算子也需要移除相应的输入。推往前推,相应的push_back算子、create_stack算子也可以被移除。经过这个pass, 计算图会被变换得和裁剪前一致。 + +更通用的情况,可以提供backward_pass和逆backward_pass。验证二者的可逆性。 + +#### 2.3.1 IfOp的反向实现 + +本方案将IfOp的反向构建流程从逻辑上分为三步:(后面whileOp也类似) + +第一步:对前向IfOp进行改造,增加一个region,创建变量栈,将所有子block中的局部变量压栈,并将该变量栈新增为IfOp的输出。 + +If包含两个或三个region。 如果包含了三个region,说明已经求了一次反向,这种情况我们在后文2.3.3中进行描述。这儿直接假设遇见的IfOp一定只包含了两个region。 + +1. 在then_region的前方,插入一个init region,该init region只包含一个block, 承担两个功能: + 1. 创建一个stack。(一个create_stack算子) + 2. 将该栈变量转发给其它region。(一个cf.yield算子) +2. 在then region和else region的输入输出中都新增stack变量。并将所有该region中定义的局部变量压栈该stack中。 + +```cpp +%x = pd.full(....) +%y = pd.full(....) +%cond = pd.less_than(x, y) +%ret1, %ret2 = pd.if(%cond) { + %1 = pd.full(....) + %2 = pd.full(...) + cf.yield(%1, %2) + } else { + %1 = pd.full(....) + %2 = pd.full(...) + cf.yield(%1, %2) + } +``` + +改造后的if op为: + +```cpp +%x = pd.full(....) +%y = pd.full(....) +%cond = pd.less_than(x, y) +%ret1, %ret2, %stack = pd.if(%cond) + init { + %stack = cf.create_stack() + cf.yield(%stack) + } + then(%arg_stack) + { + %1 = pd.full(....) + %2 = pd.full(...) + cf.push_back(%arg_stack, %1) + cf.push_back(%arg_stack, %2) + cf.yield(%1, %2, %arg_stack) + } + else (%arg_stack) + { + %1 = pd.full(....) + %2 = pd.full(...) + cf.push_back(%arg_stack, %1) + cf.push_back(%arg_stack, %2) + cf.yield(%1, %2, %arg_stack) + } +``` + +第二步:构造反向if_grad op。(反向if_grad op其实也是一个if_op,只是将其命名为if_grad,实现和cond一致) + +创建一个if_grad op, 它包含then_region和else_region。 它的输入和前向cond_op的输入完全一致,只包含一个%cond变量即可。 + +**(正常来说,反向op的输入应该包含了前向op的输出变量的反向,但是我们之前在CondOp的设计中已经描述过了,子block可以直接访问父block变量,因此,不需要作为输入)** + +if_grad op 也没有输出。这是因为if_op是在子block中,直接对父block中的变量进行引用,那相应的,在子block中,如果涉及到对父block中变量的使用,我们之间在原地进行梯度累加即可。 + +在then_region和else_region中将前向中的压栈的所有局部变量都出栈。然后按照bakcward的正常逻辑,依次给后向region中添加前向region的反向算子。 + +```cpp +%x = pd.full(....) +%y = pd.full(....) +%cond = pd.less_than(x, y) +%ret1, %ret2, %stack = pd.if(%cond) + init { + %stack = cf.create_stack() + cf.yield(%stack) + } + then(%arg_stack) + { + %1 = pd.full(....) + %2 = pd.full(...) + cf.push_back(%arg_stack, %1) + cf.push_back(%arg_stack, %2) + cf.yield(%1, %2, %arg_stack) + } else (%arg_stack) + { + %1 = pd.full(....) + %2 = pd.full(...) + cf.push_back(%arg_stack, %1) + cf.push_back(%arg_stack, %2) + cf.yield(%1, %2, %arg_stack) + } + + pd.if_grad(%cond) + then{ + %1 = cf.pop_back(%stack) + %2 = cf.pop_back(%stack) + } + else { + %1 = cf.pop_back(%stack) + %2 = cf.pop_back(%stack) + } +``` + +第三步:优化剪枝,将反向算子中用不到的栈变量,移除它的压栈和出栈算子。 + +为了保证第二步中给每个子block中的算子添加反向可以正常进行。我们在第一步和第二步中对所有的局部变量都进行了压栈出栈操作。 + +但实际上,并不是所有的局部变量在反向计算中都会被使用。因此。这一步主要就是移除无意义的压栈出栈操作。 + +如果发现栈里的局部变量完全用不到的话,可以将if_op的init_region也直接删除。 + +#### 2.3.2 WhileOp的反向实现 + +本方案将WhileOp的反向构建流程从逻辑上也分为三步: + +第一步:对前向WhileOp进行改造,增加一个region,创建变量栈,将所有子block中的局部变量压栈,并将该变量栈新增为while_op的输出。 + +WhileOp包含两个或三个region。 如果包含了三个region,说明已经求了一次反向,这种情况我们在下一节2.3.3中进行描述。这儿直接假设预计的WhileOp一定只包含了两个region。 + +1. 在cond_region的前方,插入一个init region,该init region只包含一个block, 承担两个功能: + 1. 创建一个stack。(一个create_stack算子) + 2. 将输入参数原样转发给condition region。(一个cf.yield算子, 跳转到condition region) +2. 在condition region和body region的输入输出中都新增stack变量。并将所有该region中定义的局部变量压栈该stack中。 + +```cpp +%i = pd.full(...) +%1 = pd.const(1) +%ten = pd.full(...) + +%i_2 = pd.while(%i) + cond(%arg1) { + // cond_region + %cond = pd.less_than(%arg1, %ten) + cf.cond_yield (%cond, %arg1) + } + body(%arg2){ + %i_3 = pd.add(%arg2, %1) + cf.yield (%i_3) + } +``` + +改造后的WhileOp: + +```cpp +%i = pd.full(...) +%1 = pd.const(1) +%ten = pd.full(...) + +%i_2, %stack = pd.while(%i) + init(%arg) { + %stack = cf.create_stack() + cf.yield(%arg, %stack) + } + cond(%arg1, %stack){ + %cond = pd.less_than(%arg1, %ten) + cf.push_back(%stack, %arg1) + cf.push_back(%stack, %cond) + cf.cond_yield(%cond, %arg1, %stack) + } + body(%arg2, %stack){ + %i_3 = pd.add(%arg2, %1) + cf.push_back(%stack, %arg2) + cf.push_back(%stack, %i_3) + pd.yield(%i_3, %stack) + } +``` + +第二步:构造反向while_grad op。(反向while_grad op其实也是一个while_op,只是将其命名为while_grad,实现和while一致) + +创建一个while_grad op, 它包含condition_region和body_region。 它的输入包含前向while_op的输出的所有变量的梯度以及while_op输出的容器栈。在condition_region和body_region中将前向中的压栈的所有局部变量都出栈。然后按照bakcward的正常逻辑,依次给后向region中添加前向region的反向算子。 + +```cpp +%i = pd.full(...) +%1 = pd.const(1) +%ten = pd.full(...) + +%i_2, %stack = pd.while(%i) + init(%arg) { + %stack = cf.create_stack() + cf.yield(%arg, %stack) + } + cond (%arg1, %stack){ + %cond = pd.less_than(%arg1, %ten) + cf.push_back(%stack, %arg1) + cf.push_back(%stack, %cond) + cf.cond_yield(%cond, %arg1, %stack) + } + body(%arg2, %stack){ + %i_3 = pd.add(%arg2, %1) + cf.push_back(%stack, %arg2) + cf.push_back(%stack, %i_3) + pd.yield(%i_3, %stack) + } +........ +%i_grad = pd.while_grad(%i_2_grad, %stack) + cond(%arg1_grad, %stack) { + %cond = cf.pop_back(%stack) + %arg1 = cf.pop_back(%stack) + // less_than的输入是:%arg1, %ten, 输出是%cond, 正常来说, + // 反向的输出应该是%arg1_grad, %ten_grad. + // 通过反向接口发现,less_than的反向算子是空的,也就是说less_than算子对%arg1_grad和%ten_grad没有贡献。 + // 此处直接跳过less_than的反向。 + %new_cond = cf.is_empty(%stack) + cf.cond_yield(%new_cond, %arg1_grad) + } + body(%arg2_grad)){ + %arg2 = cf.pop_back(%stack) + %i_3 = cf.pop_back(%stack) + + // add的输入是%arg2, %i1, 输出是%i_3 + // 所以反向的输出应该是%arg2_grad, %i1_grad + // 这儿需要注意的一点是,反向变量的定义域一定要和前向变量的定义域对偶。如果%i1_grad没有定义,那我们应该上溯到前向的对偶定义域中去定义一个初始为0的梯度向量,在这儿对其进行累加。 + %tmp_arg2_grad, %tmp_1_grad = pd.add_grad(%arg2, %1, %i_3, %arg_2_grad) + //在构建pd.add_grad的时候,可以发现目前已经存在一个arg2_grad 和 1_grad。 因此我们需要将pd.add_grad的输出累加到之前的变量上。 + pd.inplace_add(%arg2_arg, %tmp_arg2_grad) + pd.inplace_add(%1_grad, %temp_1_grad) + pd.yield(%arg2_grad) + } +``` + +第三步:优化剪枝,将反向算子中用不到的栈变量,移除它的压栈和出栈算子。 + +实际上,在前两步中,我们通过压栈出栈操作,给反向block中的子算子提供了前向算子的所有输入输出。 但其实反向算子不一定会用到前向的所有输出输入,或者有些算子的反向算子直接就是空的。这个时候就会出现很多用不到的局部变量。在这一步中对这些无用的变量和算子进行剪枝。 + +```cpp +%i = pd.full(...) +%1 = pd.const(1) +%ten = pd.full(...) + +%i_2, %stack = pd.while(%i) + init(%arg){ + %stack = cf.create_stack() + cf.yield(%arg, %stack) + } + cond(%arg1, %stack) { + %cond = pd.less_than(%arg1, %ten) + //cf.push_back(%stack, %arg1) + //cf.push_back(%stack, %cond) + cf.cond_yield(%cond, %arg1, %stack) + } + body(%arg2, %stack){ + %i_3 = pd.add(%arg2, %1) + cf.push_back(%stack, %arg2) + cf.push_back(%stack, %i_3) + pd.yield(%i_3, %stack) + } +........ +^%i_grad = pd.while_grad(%i_2_grad, %stack) + cond(%arg1_grad, %stack){ + // 发现%cond, %arg1都没有被使用,所以这两行的pop_back可以删除。 + / 删除的时候,需要同步删除前向中的push_back算子。 + //%cond = cf.pop_back(%stack) + //%arg1 = cf.pop_back(%stack) + %new_cond = cf.is_empty(%stack) + cf.cond_yield(%new_cond, %arg1_grad) + } + body(%arg2_grad){ + %arg2 = cf.pop_back(%stack) + %i_3 = cf.pop_back(%stack) + %tmp_arg2_grad, %tmp_1_grad = pd.add_grad(%arg2, %1, %i_3, %arg_2_grad) + pd.inplace_add(%arg2_arg, %tmp_arg2_grad) + pd.inplace_add(%1_grad, %temp_1_grad) + pd.yield(%arg2_grad) + } +``` + +#### 2.3.3 更近一步的考虑 + +我们在上一节(2.3.1)的最开始假设了backward pass中遇见的while op一定是包含了两个region。但实际上,在高阶微分场景,会对一个算子多次求反向,如果此时模型包含了while op, 在第二次及以后的backward pass中,遇见的while_op就会包含三个region。 + +目前paddle没有支持控制流高阶微分的需求。但在本节中我们可以提供一些简单的思路,方便后续有需求的时候,去进行相关实现。主要是为了论述新IR设计的完备性,证明后续对控制流高阶微分的实现只需要做增量开发。 + +当我们遇见while_op包含了三个region之时,那它一定已经包含了stack输出。(这是因为init region的存在就是为了创建stack输出,如果没有stack输出,那说明init region可以被移除,这就转变为了2.3.1中描述的两个region的情况)。而while_op包含了stack输出,那么一定存在while_grad_op使用了该stack。(如果不存在,说明该stack可以移除,因而init_region也可以被移除,再次转变为了2.3.1中描述的两个region的情况)。 + +一种简单的做法是,直接将根据stack变量获得的while_grad op直接复制一份,修改一下输入的反向变量即可。这样做唯一的问题在于 stack里面的变量被压栈了一次,却出栈了两次,而这个问题,我们只要将stack变量进行一次拷贝即可解决。 + +这种方式会有一个性能问题,那就是模型中存在两个只有输入和输出不同的while_grad op。 这个可以通过定义函数算子来进行解决。 将while_grad可以封装成一个函数。调用调用两次即可。 + diff --git a/pfcc/paddle-code-reading/IR_Dialect/first_step.md b/pfcc/paddle-code-reading/IR_Dialect/first_step.md new file mode 100644 index 000000000..81c52e06f --- /dev/null +++ b/pfcc/paddle-code-reading/IR_Dialect/first_step.md @@ -0,0 +1,1838 @@ +# 一起从代码层面熟悉 PIR —— PIR 源码阅读指南 + +| 版本 | 作者 | 指导/校验 | 时间 | 主要更新 | +| ---- | ----- | --------- |---------- | ------------------ | +| 0.7 | [Ryan](https://github.com/drryanhuang) | [Aurelius84](https://github.com/Aurelius84) |2023.10.01 | 初版| +| 1.0 | [Ryan](https://github.com/drryanhuang) | [Aurelius84](https://github.com/Aurelius84) |2023.10.13 | 初版| + +阅读本指南的收获 + +- 阅读 Paddle 工业级源码, 提升 C++ 编码能力 +- 从代码层面了解 PIR 体系的设计 +- 了解 `Pimpl` 设计模式 +- 熟悉 Paddle 代码风格 +- ~~学习胡乱使用倒叙和插叙的混乱啰嗦的博客风格写作手法~~ + + +由于 PIR 依旧在高频更新, 如果笔者写的有问题, 各位开发者大佬提个PR帮我修改一下, 万分感谢! 本文也穿插了 [pfcc/paddle-code-reading/IR_Dialect/ir_program.md](./ir_program.md) 的相关内容, 在阅读完本文后, 再去阅读 [pfcc/paddle-code-reading/IR_Dialect/ir_program.md](./ir_program.md) 相信各位家人会对 PIR 体系有更深入的了解. + + +PIR的整体设计自底向上分为三层: + +

+ +

+ +乍一看对刚接手新IR相关工作的我们来说有些抽象, 本文档旨在通过源码阅读与抽象概念相对应, 从而达到对PIR体系的理解. + + +## 1. 新 IR 体系下的 Program + +我想各位家人们一定看过《水浒传》, 《水浒传》中经常由一个好汉引出另一个好汉的出场,本小节也是如此,从 `Program` 开始,依次引出 `IrContext`, `ModuleOp`, `Block`, `Region` 和 `Operation` + +主要位置在 `paddle/pir/core/program.h` 和 `paddle/pir/core/program.cc` + +`Program` 是模型结构的抽象,分为计算图 `graphs` 和权重 `weights`。 现阶段,计算图以 `list` 的形式表示。 +Todo:未来将进行控制流算子的详细设计,引入基本块、闭包、函数等概念,不断提高 `Program` 表示计算图的能力。 + +可以在 `Program` 的源码中看到私有变量 计算图`module_` 和 权重`parameters_` + +```c++ +class IR_API Program { + public: + // ...... + + private: + // computation graph + ModuleOp module_; + // weight + ParameterMap parameters_; +}; +``` + +再来看下 `Program` 的构造函数: +```c++ +Program::Program(IrContext* context) { + module_ = ModuleOp::Create(context, this); +} +``` +只是传入了 `context` 参数, `IrContext* context` 目前可以简单理解为一个包含代码上下文参数的类, 后面会介绍 + +`ModuleOp::Create(context, this)` 的第二个参数 `this` 说明在构造 `ModuleOp` 对象时, 需要传入对应的 `Program` 对象指针 + + + + +这里我们稍微深入看一下类型 `ModuleOp` 和 `ParameterMap` + +`ModuleOp` 中有两个显眼包 `Program *program()` 和 `Block *block()`, 前者返回构造它时的`Program` 对象指针 + +```c++ +// paddle/pir/core/builtin_op.h + +class IR_API ModuleOp : public pir::Op { + public: + // ...... + Program *program(); // <-------- 这里, 显眼包在这里 + Block *block(); + + // + // As the top operation, ModuleOp only support create&destroye through + // below interface: "create"&"destroy". + static ModuleOp Create(IrContext *context, Program *pointer); + void Destroy(); +}; +``` +后者 `Block *block()` 用来存放计算图,点开 `Block` 的源码,看到私有变量 `OpListType ops_` 该变量就是用来存计算图内容的list. + +```c++ +class IR_API Block { + using OpListType = std::list; // <----- 计算图以 `list` 的形式表示 + + public: + + Block() = default; // 这里是默认构造函数 + + // ...... + + private: + Region::iterator position_; + BlockOperand first_use_; + OpListType ops_; // <------ `ops_` 变量用来存计算图内容 + BlockArgListType arguments_; + Region *parent_; +}; +``` + +再看一眼 `Block` 内部的几个方法, 基本上都是围绕 `list ops_` 的方法做了一些封装 + +```c++ + bool empty() const { return ops_.empty(); } + size_t size() const { return ops_.size(); } + + ConstIterator begin() const { return ops_.begin(); } + ConstIterator end() const { return ops_.end(); } + Iterator begin() { return ops_.begin(); } + Iterator end() { return ops_.end(); } + ReverseIterator rbegin() { return ops_.rbegin(); } + ReverseIterator rend() { return ops_.rend(); } + + Operation *back() const { return ops_.back(); } + Operation *front() const { return ops_.front(); } + void push_back(Operation *op); + void push_front(Operation *op); + Iterator insert(ConstIterator iterator, Operation *op); + Iterator erase(ConstIterator position); + void clear(); +``` + + +再看一下 `Block` 的私有变量 `Region *parent_`, 也简要看一下 `Region` 的源码, 其实只需要看前几句都 OK 了, 逻辑和 `Block` 基本类似,也有一个私有变量 `std::list blocks_` 用来存一些列 `Block *` 指针 + +```c++ +class IR_API Region { + // ...... + private: + Operation *parent_{nullptr}; // not owned // <------- + std::list blocks_; // owned // <------- 用来存一系列 Block* 指针 +}; +``` + + +位置在 `paddle/pir/core/region.cc` 和 `paddle/pir/core/region.h`, 以下内部函数都是在 `std::list blocks_` 上的封装 + +```c++ + using iterator = std::list::iterator; + using reverse_iterator = std::list::reverse_iterator; + using const_iterator = std::list::const_iterator; + ~Region(); + bool empty() const { return blocks_.empty(); } + size_t size() const { return blocks_.size(); } + + iterator begin() { return blocks_.begin(); } + iterator end() { return blocks_.end(); } + const_iterator begin() const { return blocks_.begin(); } + const_iterator end() const { return blocks_.end(); } + reverse_iterator rbegin() { return blocks_.rbegin(); } + reverse_iterator rend() { return blocks_.rend(); } + + Block *back() const { return blocks_.back(); } + Block *front() const { return blocks_.front(); } + void push_back(Block *block); + void emplace_back(); + void push_front(Block *block); + iterator insert(const_iterator position, Block *block); + iterator erase(const_iterator position); + void clear(); +``` + +另一个私有变量是 `Operation *parent_` 由此, 我们可以猜测, `Block`, `Region` 和 `Operation` 三者的关系是: +``` +Block ~= [Operation*, Operation*, ..., Operation*] +Region ~= [Block*, Block*, ..., Block*] +``` +而 `Region` 也可以被 `Operation` 所拥有, 我们简要查看一下 `Operation` 的源码 + +位置在 `paddle/pir/core/operation.h` 和 `paddle/pir/core/operation.h` + +```c++ +class IR_API alignas(8) Operation final { + private: + // ....... + Region *regions_{nullptr}; // <------ `Operation` 也拥有一个 `Region` + Block *parent_{nullptr}; // <------ `Operation` 被一个 `Block` 拥有 + Block::Iterator position_; +}; +``` + +好的, 到此, 简单梳理一下内容 + +- `Program` 中 `ModuleOp module_`存计算图 , `ParameterMap parameters_` 存权重 +- `ModuleOp` 类中, 用 `Block` 来存计算图中的内容 +- 之后了解了 `Block`, `Region` 和 `Operation` 三者的关系 + + +接下来介绍刚刚跳过的 `Program` 的 `ParameterMap parameters_` + +`Program` 与 `ParameterMap parameters_` 有关的代码如下: + +```c++ +class IR_API Program { + public: + using ParameterMap = + std::unordered_map>; + + ParameterMap& parameters() { return parameters_; } + void set_parameters(ParameterMap&& parameters) { + parameters_ = std::move(parameters); + } + + private: + // computation graph + ModuleOp module_; + // weight + ParameterMap parameters_; +}; +``` + +`ParameterMap parameters_` 是一个 key 为参数名字, value 为 `Parameter *` 的 `unordered_map`, (可以理解为Python的字典) + +- `ParameterMap& parameters()` 用来返回当前 `Program` 的参数map +- `void set_parameters(ParameterMap&& parameters)` 用来设置当前 `Program` 的参数 + +好的, 接下来来看 `Parameter` 类: + +```c++ + +class IR_API Parameter { + public: + Parameter(void* data, size_t size, pir::Type type) { + data_ = malloc(size); + memcpy(data_, data, size); + size_ = size; + type_ = type; + } + + Parameter(const Parameter& param) { + data_ = malloc(param.size_); + memcpy(data_, param.data_, param.size_); + size_ = param.size_; + type_ = param.type_; + } + + Parameter& operator=(const Parameter& param) { + data_ = malloc(param.size_); + memcpy(data_, param.data_, param.size_); + size_ = param.size_; + type_ = param.type_; + return *this; + } + + private: + void* data_; + size_t size_; + Type type_; +``` + +(目前参数内存申请依旧是新IR那部分) ~~之前杰哥考我, Paddle参数在何时进行内存申请(malloc), 我想应该就是在 `Parameter` 类初始化的时候, 有参构造/拷贝构造和 `operator=` 都重新申请了内存, 并修改了数据起始地址 `void* data_`, 数据长度 `size_t size_` 和 数据类型 `Type type_`.~~ + +这里的内存申请过程其实也验证了一句话, `Tensor` 在底层中的存储都是一维的 + +`Type` 类也是新 IR 体系下关键的一环, 我们这里也将其入栈, 之后细说 + +到此为止, 我们对新IR体系的 **模型表示层** 有了一半的了解, 以下这段摘自 [pfcc/paddle-code-reading/IR_Dialect/ir_program.md](./ir_program.md), 以方便大家了解接下来的内容 + + +> 1. `Program` 用来表示一个具体的模型。它包含两部分:`计算图` 和 `权重` 。 +> 2. `Weight` 用来对模型的权重参数进行单独存储,这也是深度学习框架和传统编译器不一样的地方。传统编译器会将数据段内嵌到程序里面。这是因为传统编译器里面,数据和代码是强绑定的,不可分割。但是对神经网络而言,一个计算图的每个 `epoch` 都会存在一份权重参数,多个计算图也有可能共同一份权重参数,二者不是强绑定的。 +> 3. `Value`、`Operation` 用来对计算图进行抽象 +> + `Operation` 表示计算图中的节点。 +> + 一个 `Operation` 表示一个算子,它里面包含了零个或多个 `Region` 。 +> + `Region` 表示一个闭包,它里面包含了零个或多个 `Block`。 +> + `Block` 表示一个符合 `SSA` 的基本块,里面包含了零个或多个 `Operation` 。 +> + 三者循环嵌套,可以实现任意复杂的语法结构。 +> + `Value` 表示计算图中的有向边,他用来将两个 `Operaton` 关联起来,描述了程序中的 `UD链` 。 +> + `OpResult` 表示定义端,定义了一个 `Value` 。 +> + `OpOperand` 表示使用端,描述了对一个 `Value` 的使用。 + + + +## 2. 新 IR 体系下的计算图 + + +我们先看一下主框架的计算图中的计算图: + +> 目前主框架和编译器分别定义了 `Program` & `Graph` 来描述计算图。 +> 主框架相对历史悠久一点,在 `Program` 中,变量的定义和算子是解藕的,算子通过变量名 (`字符串`) 简接关联到变量。一方面,`计算图有环`。另一方面,效率也不高,要想知道一个变量都被哪些算子关联了,就必须遍历 `block` 中所有算子的所有输入输出,进行字符串比对。在 `Graph` 中,一方面,变量和算子被同时抽象为了计算图节点,这增加了图优化的复杂度。另一方面,`Graph` 内嵌了 `Program` ,图优化不仅要处理图节点的UD链,还得处理图节点内嵌的 `OpDesc` & `VarDesc` 的 `UD链`,进一步增加了图优化的复杂度。 + +接下来通过以下代码来看主框架中的计算图,算子节点和变量节点都可以通过 `Graph::CreateEmptyNode` 来创建, 通过第二个参数 `ir::Node::Type type` 来指定是创建算子节点还是变量节点. + +```c++ +// 代码位置 paddle/fluid/framework/ir/pass_test.cc +void BuildCircleGraph(Graph* g) { + ir::Node* o1 = g->CreateEmptyNode("op1", Node::Type::kOperation); + ir::Node* o2 = g->CreateEmptyNode("op2", Node::Type::kOperation); + ir::Node* v1 = g->CreateEmptyNode("var1", Node::Type::kVariable); + ir::Node* v2 = g->CreateEmptyNode("var2", Node::Type::kVariable); + + o1->outputs.push_back(v1); + o2->inputs.push_back(v1); + v1->inputs.push_back(o1); + v1->outputs.push_back(o2); + + o2->outputs.push_back(v2); + o1->inputs.push_back(v2); + v2->inputs.push_back(o2); + v2->outputs.push_back(o1); +} +``` + +每个算子(如 `o1`)节点要记录自己的输入和输出, 每个变量也要记录自己的 `inputs` 和 `outputs`, 以上代码建立了一个有环图, 拓扑结构如下: + + +```mermaid +graph TD; + o1 --> v1; + v1 --> o2; + o2 --> v2; + v2 --> o1; +``` + +新 IR 体系下计算图的有向边是变量(`Value`),节点是算子(`Operation`). + +算子信息分为四部分:输入(`OpOperandImpl`)、输出(`OpResultImpl`)、属性(`Attribute`)、类型信息(`OpInfo`). + +我们从 `Operation *Create` `Operation` 的创建方法参数入手, 先来看属性 `AttributeMap` 和类型信息 `OpInfo`. + + +```c++ + static Operation *Create(const std::vector &inputs, + const AttributeMap &attributes, + const std::vector &output_types, + pir::OpInfo op_info, + size_t num_regions = 0, + const std::vector &successors = {}); // 控制流 + static Operation *Create(OperationArgument &&op_argument); +``` + +`Create` 函数的重载版本输入类型是 `OperationArgument` , 从其第三个构造函数可以看出, 该结构体是对上述输入做的封装 + +位置在 `paddle/pir/core/operation_utils.cc` 和 `paddle/pir/core/operation_utils.h` + +```c++ +struct OperationArgument { + std::vector inputs; + AttributeMap attributes; + std::vector output_types; + OpInfo info; + std::vector successors; + std::vector> regions; + + public: + OperationArgument(IrContext* ir_context, const std::string& name); // <---- 这里也是 ctx 统一管理 ? + explicit OperationArgument(OpInfo info) : info(info) {} + OperationArgument(const std::vector& inputs, + const AttributeMap& attributes, + const std::vector& types, + OpInfo info, + const std::vector successors = {}) + : inputs(inputs), + attributes(attributes), + output_types(types), + info(info), + successors(successors) {} +} +``` + +输入参数 `attributes` 的类型是 +```c++ +using AttributeMap = std::unordered_map; +``` +用 `unordered_map` 来存属性的名字和属性的内容 + +定义 `Attribute` 类描述一个具体的属性对象。里面包含一个 `AttributeStorage` 作为该属性对象的真实存储对象,该存储对象由 `IRContext` 负责管理。 + +位置在 `paddle/pir/core/attribute.h` 和 `paddle/pir/core/attribute.cc` + +```c++ +// Attribute类的统一接口。 所有 Attribute 类的派生仅派生接口,而不派生成员。 +class IR_API Attribute { + public: + using Storage = AttributeStorage; + + Attribute() = default; + + Attribute(const Storage *storage) // NOLINT + : storage_(storage) {} + + protected: + const Storage *storage_{nullptr}; +}; +``` + +`Attribute`类 与 `AttributeStorage`类 是一一对应的, 二者分别派生了对应的类别, 通过宏 `DECLARE_ATTRIBUTE_UTILITY_FUNCTOR` 关联起来 + +```c++ +// 位置在 `paddle/pir/core/attribute_base.h` + +// 在自定义Attribute类中添加一些必要的函数 + +#define DECLARE_ATTRIBUTE_UTILITY_FUNCTOR(concrete_attribute, storage_type) \ + using Storage = storage_type; \ + \ + const Storage *storage() const { \ + return static_cast(this->storage_); \ + } \ + \ + static pir::TypeId type_id() { \ + return pir::TypeId::get(); \ + } \ + \ + template \ + static bool classof(T val) { \ + return val.type_id() == type_id(); \ + } \ + \ + template \ + static concrete_attribute get(pir::IrContext *ctx, Args... args) { \ + return pir::AttributeManager::template get(ctx, \ + args...); \ + } +``` + +接下来只需要在 `Attribute` 的派生类中使用这个宏即可: + +```c++ +class IR_API BoolAttribute : public Attribute { + public: + using Attribute::Attribute; + + DECLARE_ATTRIBUTE_UTILITY_FUNCTOR(BoolAttribute, BoolAttributeStorage); + + bool data() const; +}; + +class IR_API FloatAttribute : public Attribute { + public: + using Attribute::Attribute; + + DECLARE_ATTRIBUTE_UTILITY_FUNCTOR(FloatAttribute, FloatAttributeStorage); + + float data() const; +}; +``` + + +`AttributeStorage` 派生了以下类直接与一个 `BaseType` 对应, 二者通过宏 `DECLARE_BASE_TYPE_ATTRIBUTE_STORAGE` 进行关联, 以减少代码量, 方便维护 + +```c++ +// 位置在 paddle/pir/core/builtin_attribute_storage.h + +#define DECLARE_BASE_TYPE_ATTRIBUTE_STORAGE(ConcreteStorage, BaseType) \ + struct ConcreteStorage : public AttributeStorage { \ + using ParamKey = BaseType; \ + \ + explicit ConcreteStorage(ParamKey key) { data_ = key; } \ + \ + static ConcreteStorage *Construct(ParamKey key) { \ + return new ConcreteStorage(key); \ + } \ + \ + static size_t HashValue(ParamKey key) { \ + return std::hash{}(key); \ + } \ + \ + bool operator==(ParamKey key) const { return data_ == key; } \ + \ + BaseType data() const { return data_; } \ + \ + private: \ + BaseType data_; \ + } +``` + +```c++ +// 位置在 paddle/pir/core/builtin_attribute_storage.h + +DECLARE_BASE_TYPE_ATTRIBUTE_STORAGE(BoolAttributeStorage, bool); +DECLARE_BASE_TYPE_ATTRIBUTE_STORAGE(FloatAttributeStorage, float); +DECLARE_BASE_TYPE_ATTRIBUTE_STORAGE(DoubleAttributeStorage, double); +DECLARE_BASE_TYPE_ATTRIBUTE_STORAGE(Int32AttributeStorage, int32_t); +DECLARE_BASE_TYPE_ATTRIBUTE_STORAGE(Int64AttributeStorage, int64_t); +DECLARE_BASE_TYPE_ATTRIBUTE_STORAGE(PointerAttributeStorage, void *); +DECLARE_BASE_TYPE_ATTRIBUTE_STORAGE(TypeAttributeStorage, Type); +``` + +没有使用宏 `DECLARE_BASE_TYPE_ATTRIBUTE_STORAGE` 的 `StrAttributeStorage` 和 `ArrayAttributeStorage` 分别用来存字符串和多个 `Attribute` 数组. + +在 `paddle/fluid/pir/dialect/operator/ir/op_attribute.h` 和 `paddle/fluid/pir/dialect/operator/ir/attribute_storage.h` 下, 也有其他 `Attribute` 和 `AttributeStorage` 派生类, 各位家人自行搜索查看 + + + +接下来通过一个例子, 来看 `IRContext` 是如何管理 `AttributeStorage` 的 + +```c++ +struct EmbeddingGradOpTranscriber : public OpTranscriber { + void HandleNonexistentAttribute(pir::IrContext* ctx, + pir::AttributeMap* attribute_map, + const OpAttributeInfo& info) override { + if (info.name == "padding_idx") { + (*attribute_map)[info.name] = pir::Int64Attribute::get(ctx, -1); // <------- + } else if (info.name == "sparse") { + (*attribute_map)[info.name] = pir::BoolAttribute::get(ctx, false); + } + } + // ...... +} +``` + +这里 `pir::Int64Attribute::get(ctx, -1)` 返回一个 `Int64Attribute` 对象, 注册到 `attribute_map` 指针指向的 `AttributeMap` + +`pir::Int64Attribute::get` 是一个静态方法, 调用了 `pir::AttributeManager` 的 `get` 静态方法 + +```c++ +#define DECLARE_ATTRIBUTE_UTILITY_FUNCTOR(concrete_attribute, storage_type) \ + // ...... + template \ + static concrete_attribute get(pir::IrContext *ctx, Args... args) { \ + return pir::AttributeManager::template get(ctx, \ + args...); \ + } +``` + + +`AttributeManager` 中重载了 3 个 `get` 和 3 个 `RegisterAttribute`. 这里建议看源码, 注释很全 + +- `get` 从 `IrContext` 获取一份 `Attribute` 类型的实例. 如果四有参属性, 且 `IrContext` 中没有, 则会创建一个新对象, 并在 `IrContext` 中注册, 而对于无参属性, 仅进行搜索 +- `RegisterAttribute` 用于将属性 `Attribute` 注册到 `IrContext`. + +```c++ +// 代码位置 paddle/pir/core/attribute_base.h + +struct IR_API AttributeManager { + /// + /// \brief Get a unique instance of Attribute T from IrContext. Note: For a + /// parametric attribute, if not found in IrContext, it will try to create a + /// new instance and register it to IrContext; for a parameterless attribute, + /// only search. + + template + static T get(IrContext *ctx, Args &&...args) { + return get( + ctx, pir::TypeId::get(), std::forward(args)...); + } + + // ...... + + /// + /// \brief Register a unique instance of Attribute T to IrContext. + /// + /// \param ctx The IrContext instance. + /// + template + static void RegisterAttribute(IrContext *ctx) { + RegisterAttribute(ctx, pir::TypeId::get()); + } + + // ...... +}; +``` + + +所以之后我们在写单测的时候, 可以这样创建 `Attribute` + +```c++ + // 代码位置 paddle/pir/dialect/shape/ir/shape_op.cc + + argument.AddAttribute("knownNonNegative", attr_knownNonNegative); + Attribute attr_knownNegativeOne = + BoolAttribute::get(IrContext::Instance(), knownNegativeOne); + argument.AddAttribute("knownNegativeOne", attr_knownNegativeOne); + Attribute attr_knownNonSizeOne = + BoolAttribute::get(IrContext::Instance(), knownNonSizeOne); + argument.AddAttribute("knownNonSizeOne", attr_knownNonSizeOne); + Attribute attr_knownNonSizeZero = + BoolAttribute::get(IrContext::Instance(), knownNonSizeZero); + argument.AddAttribute("knownNonSizeZero", attr_knownNonSizeZero); +``` + +此处的 `argument` 是上文提到的 `OperationArgument` 结构体, 将外部传入的属性注册到自己的 `AttributeMap attributes` 中 + +```c++ + // 位置在 paddle/pir/core/operation_utils.h + + /// Add an attribute with the specified name. + void AddAttribute(const std::string& name, Attribute attr) { + attributes[name] = attr; + } +``` + +这里浅浅的引出 `Builder`, 除了我们刚提到的 `XXXAttribute::get` 静态方法来创建 `Attribute` 实例, 我们可以使用 `Builder` 来创建: + +```c++ +void Operation1::Build(pir::Builder &builder, // NOLINT + pir::OperationArgument &argument) { // NOLINT + std::unordered_map attributes{ + {"op1_attr1", builder.str_attr("op1_attr2")}, // <------------------- 这里这里, 看这里 + {"op1_attr2", builder.str_attr("op1_attr2")}}; + argument.AddOutput(builder.float32_type()); + argument.AddAttributes(attributes); +} +``` + +`builder.str_attr` 直接创建了一个 `StrAttribute` 实例, 实际上 `Builder` 内部对这些 `Attribute` 的创建做了一系列封装, 依旧是调用 `get` 静态方法 + +```c++ +// 代码位置在 paddle/pir/core/builder.h 和 paddle/pir/core/builder.cc + +StrAttribute Builder::str_attr(const std::string &value) { + return StrAttribute::get(context_, value); +} +BoolAttribute Builder::bool_attr(bool value) { + return BoolAttribute::get(context_, value); +} +FloatAttribute Builder::float_attr(float value) { + return FloatAttribute::get(context_, value); +} +DoubleAttribute Builder::double_attr(double value) { + return DoubleAttribute::get(context_, value); +} +Int32Attribute Builder::int32_attr(int32_t value) { + return Int32Attribute::get(context_, value); +} +``` + +`Builder` 统一了 `Attribute` 类的接口, 我们后续依旧会用到它. + +好的, 接下来我们回忆一下 + +- 创建 `Operation` 需要4个算子信息, 我们目前只介绍了 `Attribute`, `Attribute` 会存在 `AttributeMap` 中, 它是一个 `unordered_map`, key 为属性名, value 是 `Attribute` +- `Attribute`类 与 `AttributeStorage`类, 后者作为该属性对象的真实存储对象,该存储对象通过 `AttributeManager` 由 `IRContext` 负责管理。 +- 可以通过 `Int32Attribute::get` 静态方法来创建 `Attribute` 实例, 也可以通过 `Builder::int32_attr` 来创建实例 + + +接下里我们来看 `OpInfo`, 这是一个使用 `Pimpl` 设计模式构造的类, 关于 `Pimpl` , 家人们可以查一下使用方式. 以下这段摘自 [pfcc/paddle-code-reading/IR_Dialect/ir_program.md](./ir_program.md), 以方便大家了解接下来的内容 + +> PIR 中很多对象是一次构造,多次引用。常规做法是构造一次,然后通过指针的方式进行引用。 +> 比如 `OpInfo` , 目前 Paddle 是在算子注册的时候,构造一个 `OpInfo` 对象,然后在相应的 `OperatorBase` 里面包含了 `OpInfo` 指针作为成员变量。 +> +> 本文大部分场景采用 `Pimpl` 设计模式。将具体实现对象封装在 `Impl` 数据结构中。 +> 采用这种设计模式的数据结构包括 `Type`、`Attribute`、`OpInfo`、`OpResult`、`Value`、`OpOperand` 等等,它们自身都只包含一个 `impl` 指针,可以快速浅拷贝。真正的实现对应 `TypeStorage` 、`AttributeStorage`、`OpInfoImpl`、`OpResultImpl`、`ValueImpl` 和 `OpOperandImpl`。其中, `Impl` 不会对用户暴漏,但是 `Storage` 会。 + +> 这种实现方式的好处在于: +> * 实现部分保持独立,后续可以随意重构升级,不影响其它模块。 +> * 接口部分只包含一个指针,可以高效率拷贝。 + + + +以下是一段单测代码片段, 这里通过 `IrContext::GetRegisteredOpInfo` 传入 `op_name` 字符串来获取相应的 `OpInfo` + +```c++ + // 代码位置在 test/cpp/pir/core/ir_infershape_test.cc + + // (2) Get registered operations. + + std::string op_name = OperationTest::name(); + pir::OpInfo op_info = ctx->GetRegisteredOpInfo(op_name); + + std::vector op_inputs = {}; + std::vector op_output_types = {pir::Float32Type::get(ctx)}; + pir::Operation *op = + pir::Operation::Create(op_inputs, {}, op_output_types, op_info); +``` + + +`GetRegisteredOpInfo` 方法通过 `IrContextImpl` 的 `GetOpInfo` 方法来查找 `OpInfo` + +```c++ +// 代码位置 paddle/pir/core/ir_context.cc 和 paddle/pir/core/ir_context.h + +IrContextImpl &impl() { return *impl_; } + +OpInfo IrContext::GetRegisteredOpInfo(const std::string &name) { + return impl().GetOpInfo(name); +} +``` + +`impl` 方法返回私有变量 `IrContextImpl *impl_`, 所以此处也是 `Pimpl` 设计模式. + +在 `IrContextImpl` 中, 所有的 `OpInfo` 都注册在 `OpInfoMap registed_op_infos_` 中, `OpInfoMap` 是类型 `std::unordered_map` + + +`GetOpInfo` 方法搜索注册在 `registed_op_infos_` 中的 `OpInfo`, 如果有则返回, 未查找到则返回一个空 `OpInfo` 对象. + + +```c++ + // 代码位置 paddle/pir/core/ir_context.cc + OpInfo GetOpInfo(const std::string &name) { + + // lock_guard对象生命周期内一直上锁 + std::lock_guard guard(registed_op_infos_lock_); + + // 在 registed_op_infos_ 中找对应的 `OpInfo`, 有则返回, 没有则 return 空的 `OpInfo` + auto iter = registed_op_infos_.find(name); + if (iter != registed_op_infos_.end()) { + VLOG(8) << "Found a cached OpInfo of: [name=" << name + << ", OpInfo: ptr=" << iter->second.AsOpaquePointer() << "]."; + return iter->second; + } + VLOG(8) << "No cache found operation of: [Name=" << name << "]."; + return OpInfo(); + } +``` + +到此, 我们知道了 `Attribute` 和 `OpInfo` 都会在 `IrContext` 中进行注册并由其管理 + + +接下来来看 `OpInfo` 的源代码, 可以看到有两个 `bool` 函数, `HasTrait` 和 `HasInterface`, 二者用来表示是否具有某个特征 `Trait` 和 是否具有某个接口 `Interface`. + + + +```c++ +// 位置 paddle/pir/core/op_info.h + +class IR_API OpInfo { + public: + + // ...... + + template + bool HasTrait() const { + return HasTrait(TypeId::get()); + } + + bool HasTrait(TypeId trait_id) const; + + template + bool HasInterface() const { + return HasInterface(TypeId::get()); + } + + bool HasInterface(TypeId interface_id) const; + + // ...... + + private: + /// The internal implementation of the operation name. + /// Not owned. + OpInfoImpl *impl_{nullptr}; +}; +``` + +以下这段摘自 [pfcc/paddle-code-reading/IR_Dialect/ir_program.md](./ir_program.md), 以方便大家了解接下来的内容 + + +> 特征 `Trait` 和接口 `Interface` 用来对算子进行另一维度的划分。 是对算子所拥有的类型信息的抽象描述。 +> +> +> 比如我们用 `ReadOnly` 特征来描述一个算子不会修改它的输入操作数。这个信息对图优化和并行调度很重要。 +> 有些算子具有 `InferShape` 函数,有些则没有。可以通过 `InferShape` 接口来标志算子是否存在 `InferShape` 函数。 +> 再比如,有些算子有反向算子构建函数,有些则没有。可以通过 `GradOpMaker` 接口来标志算子是否存在 `GradOpMaker` 函数。 +> +> +> 在当前的 Paddle 中, `OpInfo` 包含了很多接口,比如 `GradOpMakerFN` 、`InferVarTypeFN`、`InferShapeFN`、`InferInplaceOpFN` 、`InferNoNeedBufferVarsFN` 等等。但实际上,每一项信息只是针对某些算子有效。把它们全部以成员变量的形式放在 `OpInfo` 结构体,意味着我们可能只是想要给某几个算子类型增加的一个接口信息,却导致所有的算子类型都增加了一块内存占用。 +> +> +> 因此,我们将这些信息统一抽象为接口,然后在 `OpInfo` 里面包含一个 `InterfaceMap` ;由具体算子在定义的时候决定自己都包含了哪些接口,进而构造对应的 `InterfaceMap`。这样后续针对特定算子的接口扩展就不会对其它算子造成影响。 +> +> +> 特征和接口的区别在于: +> + 特征的相关接口实现跟具体算子无关,不需要多态,它对相关算子的接口没有要求,比如 `ReadOnly` ,只是通知使用者,该算子不会修改它的输入的值。 +> + 而接口的实现,是跟具体算子类型相关的,需要考虑多态,比如 `InferShape` 接口,是强行要求了相应算子一定定义了 `InferShape` 接口。 +> +> +> 之所以区分特征和接口,也是从实现上考虑: +> + 特征的实现比较简单,我们用 `TraitId` 对每个特征进行标志,最简单的方式,可以在 `OpInfoImpl` 中,包含一个 `vector` 。来判断该算子是否具有某个特征即可。 +> + 接口的实现则较为复杂,需要根据具体的算子类型的不同,实现多态。 + + + +我们等一下再看 `Trait` 和 `Interface`, 先来看看 `op_info.cc` 的源码, 基本上都是封装调用了 `OpInfoImpl impl_` 的方法, 再次验证 `Pimpl` 设计模式. + +```c++ +// 位置 paddle/pir/core/op_info.cc + +namespace pir { +bool OpInfo::HasTrait(TypeId trait_id) const { + return impl_ && impl_->HasTrait(trait_id); +} + +bool OpInfo::HasInterface(TypeId interface_id) const { + return impl_ && impl_->HasInterface(interface_id); +} + +IrContext *OpInfo::ir_context() const { + return impl_ ? impl_->ir_context() : nullptr; +} +Dialect *OpInfo::dialect() const { return impl_ ? impl_->dialect() : nullptr; } + +const char *OpInfo::name() const { return impl_ ? impl_->name() : nullptr; } + +TypeId OpInfo::id() const { return impl_ ? impl_->id() : TypeId(); } + +void OpInfo::Verify(Operation *operation) const { impl_->verify()(operation); } + +void *OpInfo::GetInterfaceImpl(TypeId interface_id) const { + return impl_ ? impl_->GetInterfaceImpl(interface_id) : nullptr; +} +} // namespace pir +``` + + +`OpInfoImpl` 的构造函数是私有的, 但是内部提供了直接构造 `OpInfo` 的静态方法 `OpInfo Create`, 参数需要传入接口信息 `interface_map` 和 特征信息 `trait_set` , 二者都是 `vector` + +```c++ + // 代码位置 paddle/pir/core/op_info_impl.h + public: + static OpInfo Create(Dialect *dialect, + TypeId op_id, + const char *op_name, + std::vector &&interface_map, + const std::vector &trait_set, + size_t attributes_num, + const char *attributes_name[], + VerifyPtr verify); + + // `OpInfoImpl` 构造函数私有 + private: + OpInfoImpl(pir::Dialect *dialect, + TypeId op_id, + const char *op_name, + uint32_t num_interfaces, + uint32_t num_traits, + uint32_t num_attributes, + const char **p_attributes, + VerifyPtr verify) + : dialect_(dialect), + op_id_(op_id), + op_name_(op_name), + num_interfaces_(num_interfaces), + num_traits_(num_traits), + num_attributes_(num_attributes), + p_attributes_(p_attributes), + verify_(verify) {} +``` + +`trait_set` 的每个元素是 `TypeId`, 用来判断该算子是否具有某个特征, 目前关于 `TypeId` 类的源码暂不赘述, 仅需要知道其用法. + +`TypeId` 是 `Type` 的唯一标识,每个 `Type` 对应一个唯一的 `TypeId`, 相同的id表示同一个 `Type` 类。 `TypeId` 提供了一个实例化接口: `TypeId::get`, 以下是其用法demo: + +```c++ +// 创建一个类 TypeA +class TypeA {}; +// 获取其唯一的 TypeId +TypeId type_a_id = TypeId::get(); +``` + +`interface_map` 的每个元素是 `InterfaceValue`, 由具体算子在定义的时候决定自己都包含了哪些接口. 我们看一眼 `InterfaceValue` 的私有变量, `type_id_` 实现不同接口的区分, 而 `model_` 是一个函数指针, 根据具体的算子类型的不同,实现多态. + + +```c++ +// 代码位置 paddle/pir/core/interface_value.h +class IR_API InterfaceValue { + // ...... + private: + TypeId type_id_; + void *model_{nullptr}; +}; +``` + +接下来查看 `OpInfoImpl` 的成员函数 `Create`, 其返回一个私有变量 `impl_` 指向自己的 `OpInfo` 对象. 代码可以细分为4部分: + +这里插一句, 因为不同类型的 `Opration` 对应的 `OpInfoImpl` 的接口个数 `interfaces_num`、特征个数 `traits_num` 和属性个数 `attributes_num` 都是不同的, 但都是设置后不会再修改. 因此, 我们采用类似 `Operation` 的设计方式, 下图是其 layout 可视化结果(该图可能存在一些问题, 家人们只需关注其前三部分的 layout 即可). + + +

+ +

+ +代码第一部分用来申请 `interfaces`, `traits` 和 `opinfo_impl` 三者的总内存, 内存首地址为 `char *base_ptr` + + + +```c++ +// 代码位置 paddle/pir/core/op_info_impl.cc +OpInfo OpInfoImpl::Create(Dialect *dialect, + TypeId op_id, + const char *op_name, + std::vector &&interface_map, + const std::vector &trait_set, + size_t attributes_num, + const char *attributes_name[], // NOLINT // <---- 规避代码风格检查 + VerifyPtr verify) { + // (1) Malloc memory for interfaces, traits, opinfo_impl. + size_t interfaces_num = interface_map.size(); + size_t traits_num = trait_set.size(); + VLOG(6) << "Create OpInfoImpl with: " << interfaces_num << " interfaces, " + << traits_num << " traits, " << attributes_num << " attributes."; + size_t base_size = sizeof(InterfaceValue) * interfaces_num + + sizeof(TypeId) * traits_num + sizeof(OpInfoImpl); + char *base_ptr = static_cast(::operator new(base_size)); + VLOG(6) << "Malloc " << base_size << " Bytes at " + << static_cast(base_ptr); +``` + + +第二部分是将 `interface_map` 中的代表接口的 `InterfaceValue` 对象依次放入刚申请的内存中 + +```c++ + if (interfaces_num > 0) { + std::sort(interface_map.begin(), interface_map.end()); // 这里有排序方便后续二分查找 + for (size_t index = 0; index < interfaces_num; ++index) { + new (base_ptr + index * sizeof(InterfaceValue)) + InterfaceValue(std::move(interface_map[index])); + } + base_ptr += interfaces_num * sizeof(InterfaceValue); + } +``` + + +第三部分是将 `trait_set` 中的代表特征的 `TypeId` 对象依次放入刚申请的内存中 + +```c++ + if (traits_num > 0) { + auto p_first_trait = reinterpret_cast(base_ptr); + memcpy(base_ptr, trait_set.data(), sizeof(TypeId) * traits_num); + std::sort(p_first_trait, p_first_trait + traits_num); // 这里有排序方便后续二分查找 + base_ptr += traits_num * sizeof(TypeId); + } +``` + + +最后一部分, 创建 `OpInfoImpl` 对象, 并将对象地址赋值给 `base_ptr`, 并创建 `OpInfo` 对象并返回. + +```c++ + // Construct OpInfoImpl. + VLOG(6) << "Construct OpInfoImpl at " << reinterpret_cast(base_ptr) + << " ......"; + OpInfo op_info = OpInfo(new (base_ptr) OpInfoImpl(dialect, + op_id, + op_name, + interfaces_num, + traits_num, + attributes_num, + attributes_name, + verify)); + return op_info; +} +``` + +接下来我们通过一个 demo 来理清楚刚才讲的内容 + +```c++ + // 代码位置 paddle/fluid/pybind/pir.cc + range_block_do( + whole_block, range, [&need_buffer_values](::pir::Operation *op) { + if (op->HasInterface() == false) { + // not a OpYamlInfoInterface, can't have no_need_buffer. + for (const auto &operand : op->operands_source()) { + need_buffer_values.insert(operand); + } + } else { + auto opinfo = + op->dyn_cast().GetOpInfo(); + int counter = 0; + for (const auto &op_input_info : std::get<0>(opinfo)) { + if (!op_input_info.no_need_buffer) { + need_buffer_values.insert(op->operand_source(counter)); + } + counter += 1; + } + } + }); +``` + +在上述代码匿名函数中的 if 判断条件调用了 `Operation` 的成员函数 `HasInterface` + +```c++ +op->HasInterface() +``` + +但我们刚刚明明记得 `HasInterface` 是 `OpInfo` 的成员函数啊, 是的, 在 `Operation` 中, 对 `OpInfo` 的成员函数 `HasTrait` 和 `HasInterface` 做了封装 + +同时每个 `Operation` 都有两个私有变量, `AttributeMap attributes_` 和 `OpInfo info_`, 前者定义Op属性, 后者定义Op类型信息 + +```c++ +// 代码位置 paddle/pir/core/operation.h + +class IR_API alignas(8) Operation final { + public: + // ...... + + template + bool HasTrait() const { + return info_.HasTrait(); + } + + template + bool HasInterface() const { + return info_.HasInterface(); + + private: + // ...... + + AttributeMap attributes_; + + OpInfo info_; + } +``` + +接下来我们继续深入, 再看一次 `OpInfo` 的成员函数 `HasTrait` 和 `HasInterface` 的源码 + +这里 `op_info.h` 头文件对 `HasTrait` 和 `HasInterface` 做了申明 + +```c++ + // paddle/pir/core/op_info.h + template + bool HasTrait() const { + return HasTrait(TypeId::get()); + } + + bool HasTrait(TypeId trait_id) const; + + template + bool HasInterface() const { + return HasInterface(TypeId::get()); + } + + bool HasInterface(TypeId interface_id) const; +``` + +而在 `op_info.cc` 文件中做出了实现, 回忆一下, 依旧是 `pimpl` 的设计模式, 由 `OpInfoImpl` 的成员函数 `HasTrait` 和 `HasInterface` 去实现 + +```c++ +// paddle/pir/core/op_info.cc + +bool OpInfo::HasTrait(TypeId trait_id) const { + return impl_ && impl_->HasTrait(trait_id); +} + +bool OpInfo::HasInterface(TypeId interface_id) const { + return impl_ && impl_->HasInterface(interface_id); +} +``` + + +因此, 我们去看 `OpInfoImpl` 的成员函数 `HasTrait` 和 `HasInterface` 源码即可, 如果 `trait` 或 `interfaces` 的数量小于 0, 直接返回 `false`, 显然就是没有此接口或者特征, 如果数量大于0, 则使用 `std::binary_search` 进行二分法搜索, 回忆一下, 在存 `Trait` 和 `Interface` 时, 是经过 `sort` 排序的 + +如果家人们不清楚 `p_first_trait` 和 `p_first_interface` 这俩个变量的计算方式, 那么就回忆一下 `OpInfoImpl` 内部的创建 `OpInfo` 的函数 `create`, 他是按照何种 layout 来存内容的. 这俩变量的计算过程就是减去对应的偏移量. + + +```c++ +// 代码位置 paddle/pir/core/op_info_impl.cc + +bool OpInfoImpl::HasTrait(TypeId trait_id) const { + if (num_traits_ > 0) { + const TypeId *p_first_trait = + reinterpret_cast(reinterpret_cast(this) - + sizeof(pir::TypeId) * num_traits_); + return std::binary_search( + p_first_trait, p_first_trait + num_traits_, trait_id); + } + return false; +} + +bool OpInfoImpl::HasInterface(TypeId interface_id) const { + if (num_interfaces_ > 0) { + const InterfaceValue *p_first_interface = + reinterpret_cast( + reinterpret_cast(this) - + sizeof(pir::TypeId) * num_traits_ - + sizeof(InterfaceValue) * num_interfaces_); + return std::binary_search(p_first_interface, + p_first_interface + num_interfaces_, + InterfaceValue(interface_id)); + } + return false; +} +``` + + +到此, 家人们对 `OpInfo` 应该有了一定的了解, 我们讲完了算子信息的 `2/4`:属性 `Attribute` 和 类型信息 `OpInfo`. + +接下来看算子信息的输入 `OpOperand/OpOperandImpl` 和输出 `OpResult/OpResultImpl` + +我们之前提到了计算图是个有向图, 其中 `Value` 用来表示边, 他将两个 `Operaton` 关联起来,描述了程序中的 `UD链`. + +> `UD(Use-Definition)链`:用来关联算子,能够通过接口获取到定义该变量的惟一算子,以及使用该变量的算子链表 + + +`OpResult` 表示定义端,定义了一个 `Value`, 类 `OpResult` 继承自 `Value`. 而 `OpOperand` 表示使用端,描述了对一个 `Value` 的使用, 不继承 `value`. + + +`UD链` 似乎有些抽象, 我们依旧从一个单测开始看, 此处是运行结束后, 打印当前 `Value` 对象的 `UD链` + +```c++ + // 代码位置在 test/cpp/pir/core/ir_value_test.cc + + // destroy + VLOG(0) << op1->result(0).PrintUdChain() << std::endl; + op4->Destroy(); + VLOG(0) << op1->result(0).PrintUdChain() << std::endl; + op3->Destroy(); + VLOG(0) << op1->result(0).PrintUdChain() << std::endl; + op2->Destroy(); + VLOG(0) << op1->result(0).PrintUdChain() << std::endl; + op1->Destroy(); +``` + +在 `Value` 的 `PrintUdChain` 函数中, 先使用宏 `CHECK_VALUE_NULL_IMPL` 去校验自己的 `impl_` 变量是否为 `nullptr`, 如果为空则抛异常. 是的, `Value` 依旧是 `Pimpl` 的设计模式, `impl()` 返回 `ValueImpl* impl_` 指针 + +```c++ +// 代码位置在 paddle/pir/core/value.cc 和 paddle/pir/core/value.h +std::string Value::PrintUdChain() { + CHECK_VALUE_NULL_IMPL(PrintUdChain); + return impl()->PrintUdChain(); +``` + + +```c++ +// 用来校验的宏 +// 代码位置在 paddle/pir/core/value.cc 和 paddle/pir/core/value.h +#define CHECK_NULL_IMPL(class_name, func_name) \ + IR_ENFORCE(impl_, \ + "impl_ pointer is null when call func:" #func_name \ + " , in class: " #class_name ".") + +#define CHECK_VALUE_NULL_IMPL(func_name) CHECK_NULL_IMPL(Value, func_name) +``` + + +那我们继续 DFS 来看 `ValueImpl` 的成员函数 `PrintUdChain`, 看到这个函数, 知道链表的宝子们就豁然开朗了, 这就是一个链表的遍历, `first_use` 函数返回链表的头结点指针 `tmp`, `<<` 到 `result` 之后, 通过 `next_use` 走到链表的下一个节点, 直到节点为空. + +```c++ +// paddle/pir/core/value_impl.cc +std::string ValueImpl::PrintUdChain() { + std::stringstream result; + result << "Value[" << this << "] -> "; + OpOperandImpl *tmp = first_use(); + if (tmp) { + result << "OpOperand[" << reinterpret_cast(tmp) << "] -> "; + while (tmp->next_use() != nullptr) { + result << "OpOperand[" << reinterpret_cast(tmp->next_use()) + << "] -> "; + tmp = tmp->next_use(); + } + } + result << "nullptr"; + return result.str(); +} +``` + +需要注意的是, `UD链` 的每一个节点是 `OpOperandImpl` 而不是 `OpOperand`, 在 `OpOperand` 源码中, 基本都是对 `OpOperandImpl` 使用的封装, 若有需要再详细看 + +`OpOperand` 位置在 `paddle/pir/core/op_operand.h` 和 `paddle/pir/core/op_operand.cc` + +接下来看 `OpOperandImpl` 的源码, 主要看与 `UD链` 设计相关的部分, 有两个用来遍历链表的私有变量: + + +```c++ +// paddle/pir/core/op_operand_impl.h +class OpOperandImpl { + + private: + Value source_; + + OpOperandImpl *next_use_ = nullptr; // 下一个节点指针 + + OpOperandImpl **prev_use_addr_ = nullptr; // 上一个节点指针的地址 + +}; + +``` + + +`InsertToUdChain` 和 `RemoveFromUdChain` 函数用来对 `UD链` 进行插入和删除节点操作, `InsertToUdChain` 应该使用的是头插法. + + +```c++ +// paddle/pir/core/op_operand_impl.cc + +void OpOperandImpl::InsertToUdChain() { + prev_use_addr_ = source_.impl()->first_use_addr(); + next_use_ = source_.impl()->first_use(); + if (next_use_) { + next_use_->prev_use_addr_ = &next_use_; + } + source_.impl()->set_first_use(this); +} + +void OpOperandImpl::RemoveFromUdChain() { + if (!source_) return; + if (!prev_use_addr_) return; + if (prev_use_addr_ == source_.impl()->first_use_addr()) { + /// NOTE: In ValueImpl, first_use_offseted_by_index_ use lower three bits + /// storage index information, so need to be updated using the set_first_use + /// method here. + source_.impl()->set_first_use(next_use_); + } else { + *prev_use_addr_ = next_use_; + } + if (next_use_) { + next_use_->prev_use_addr_ = prev_use_addr_; + } + next_use_ = nullptr; + prev_use_addr_ = nullptr; + source_ = nullptr; +} +``` + +我们之前提到过 `OpOperand/OpOperandImpl` 表示使用端,描述了对一个 `Value` 的使用. 私有变量 `source_` 就是一个 `Value`. `InsertToUdChain` 函数用来将自己添加到 `source_` 变量持有的 `UD链` 中. 同时, 由于插入删除节点的操作不安全, 于是设置为私有. + + + +瞅一眼 `OpOperandImpl` 的构造函数, 需要传入 `Value` 和对应的 `Operation*` + +```c++ +// paddle/pir/core/op_operand_impl.cc +OpOperandImpl::OpOperandImpl(pir::Value source, pir::Operation *owner) + : source_(source), owner_(owner) { + if (!source) { + return; + } + InsertToUdChain(); +} +``` + +在 `Operation` 的定义中, 也需要定义对应的 `OpOperandImpl *` 和 `OpResultImpl *`, 但是如果 `ctrl` + 鼠标左键, 你会发现木有 `op_result_impl` 和 `op_operand_impl` 这四个函数, 小朋友, 你是否有很多问号? + + +```c++ + // paddle/pir/core/operation.h + int32_t ComputeOpResultOffset(uint32_t index) const; + detail::OpResultImpl *op_result_impl(uint32_t index); + const detail::OpResultImpl *op_result_impl(uint32_t index) const; + + int32_t ComputeOpOperandOffset(uint32_t index) const; + detail::OpOperandImpl *op_operand_impl(uint32_t index); + const detail::OpOperandImpl *op_operand_impl(uint32_t index) const; +``` + +打开 `operation.cc` 拉到最下面, 可以看到有个宏 `COMPONENT_IMPL`. 是的, `OpResultImpl` 和 `OpOperandImpl` 就是在宏里定一个, `Paddle` 源码中有很多这样的 trick, 可以看到 `OpXXXXtImpl` 都调用 `ComputeOpXXXXOffset` 计算了 `offset` 偏移量, 和 `OpInfo/OpInfoImpl` 的设计方式相同, 通过不同的偏移量来访问对应的地址. 之后详细说 `Operation` 源码时, 会再次提到这个. + +```c++ +// paddle/pir/core/operation.cc + +int32_t Operation::ComputeOpResultOffset(uint32_t index) const { + // ...... +} + +int32_t Operation::ComputeOpOperandOffset(uint32_t index) const { + // ...... +} + +#define COMPONENT_IMPL(component_lower, componnent_upper) \ + componnent_upper##Impl *Operation::component_lower##_impl(uint32_t index) { \ + int32_t offset = Compute##componnent_upper##Offset(index); \ + return reinterpret_cast( \ + reinterpret_cast(this) + offset); \ + } \ + const componnent_upper##Impl *Operation::component_lower##_impl( \ + uint32_t index) const { \ + int32_t offset = Compute##componnent_upper##Offset(index); \ + return reinterpret_cast( \ + reinterpret_cast(this) + offset); \ + } + +COMPONENT_IMPL(op_result, OpResult) +COMPONENT_IMPL(op_operand, OpOperand) +``` + + +到此, 关于 `OpOperand/OpOperandImpl` 的设计细节我们都了解的差不多了, 关于 `Value/ValueImpl` 和 `OpResult/OpResultImpl` 我们也做了铺垫. + + +好了, 我们现在开始回溯, 回到梦开始的地方 `Operation *Create` 函数, 他的输入 `inputs` 使用一个 `vector` 来描述 + +```c++ + // paddle/pir/core/operation.h + static Operation *Create(const std::vector &inputs, + const AttributeMap &attributes, + const std::vector &output_types, + pir::OpInfo op_info, + size_t num_regions = 0, + const std::vector &successors = {}); + static Operation *Create(OperationArgument &&op_argument); +``` + +`Value` 更是重量级, 回忆一下: + +- `Operation` 是计算图中的节点, 是对一个算子的具体描述, `Value` 是计算图的节点, 将两个 `Operaton` 关联起来,描述了程序中的 `UD链` +- `OpResult` 是继承自 `Value`, 表示定义端,定义了一个 `Value` +- `OpOperand` 表示使用端,描述了对一个 `Value` 的使用 +- `Value` 也使用 `pImpl` 的设计模式, 有私有变量 `ValueImpl *` 指针 +- `OpResultImpl` 也继承自 `ValueImpl` +- `ValueImpl` 有 `PrintUdChain` 函数来打印 `UD链`, 使用 `first_use` 来返回 `UD链` 的头指针 + + +`Value` 把之前的东西都串起来了, 好的现在开始看 `Value / ValueImpl` 的源码, `Value` 提供了访问 `UD链` 的迭代器, `use_begin` 和 `use_end` 分别返回迭代的开始位置和结束位置. 在 `Value` 内部的函数(如 `size_t Value::use_count`, `void Value::ReplaceUsesWithIf` 和 `void Value::ReplaceAllUsesWith`) 都有用到该迭代器. + + +```c++ +// 代码位置 paddle/pir/core/value.h + +class IR_API Value { + public: + + using UseIterator = ValueUseIterator; + + UseIterator use_begin() const; + + UseIterator use_end() const; + + protected: + detail::ValueImpl *impl_{nullptr}; +}; +``` + + +而 `UseIterator` 是一个单向迭代器, 因为源码中只有 `operator++()` 和 `operator++(int)` 操作, 而没有相应的 `--` 操作. + + +```c++ +// 代码位置 paddle/pir/core/use_iterator.h + +template +class ValueUseIterator { + public: + ValueUseIterator(OperandType use = nullptr) : current_(use) {} // NOLINT + + ValueUseIterator &operator++() { + current_ = current_.next_use(); + return *this; + } + + ValueUseIterator operator++(int) { + ValueUseIterator tmp = *this; + current_ = current_.next_use(); + return tmp; + } + + protected: + OperandType current_; +}; +``` + + +顺便来看一下 `Value` 的三个成员函数, 都使用自身的迭代器来完成. + +- `use_count` 用来返回 `UD链` 的长度 +- `ReplaceUsesWithIf` 通过传入 `function` 来给满足条件的 `UD链` 节点 `OpOperandImpl` 替换新的 `Value source_` +- `ReplaceAllUsesWith` 直接传入新的 `Value` 对象来替换原有节点的 `Value source_` + +```c++ +// paddle/pir/core/value.h +size_t Value::use_count() const { + size_t count = 0; + for (auto it = use_begin(); it != use_end(); ++it) count++; + return count; +} + +void Value::ReplaceUsesWithIf( + Value new_value, + const std::function &should_replace) const { + for (auto it = use_begin(); it != use_end();) { + if (should_replace(*it)) { + (it++)->set_source(new_value); + } + } +} + +void Value::ReplaceAllUsesWith(Value new_value) const { + for (auto it = use_begin(); it != use_end();) { + (it++)->set_source(new_value); + } +} +``` + + +`Value` 类阅读完毕后, 我们来看 `ValueImpl` 的实现, `ValueImpl` 有一个私有变量 `first_use_offseted_by_kind_` , 类型 `OpOperandImpl *`, 由于 `OpOperandImpl` 类含有 4 个指针, 所以该类是 8-byte 对齐的, 故该类的指针 `OpOperandImpl *` 后三个bit位都是 0, 所以我们可以利于这后三bit位来存额外的数据. + +我们将 `kind` 值存在后三位: + +- kind = 0~5 时, 代表 positions 0 to 5 inline output(OpInlineResultImpl); +- kind = 6 时, 代表 position >=6 outline output(OpOutlineResultImpl) +- kind = 7 为保留位 + +```c++ +// 代码位置 paddle/pir/core/value_impl.h + +class alignas(8) ValueImpl { + + // ...... + protected: + + OpOperandImpl *first_use_offseted_by_kind_ = nullptr; +}; +``` + +因此, 函数 `OpOperandImpl *first_use()` 返回 `UD链` 的头地址时, 需要将 `first_use_offseted_by_kind_` 后3位置0, 即和 `~0x07` 相与. 而获取 `kind` 的 `kind()` 函数只需要 `first_use_offseted_by_kind_` 的后三位. + + +```c++ +// 代码位置 paddle/pir/core/value_impl.h +class alignas(8) ValueImpl { + public: + + OpOperandImpl *first_use() const { + return reinterpret_cast( + reinterpret_cast(first_use_offseted_by_kind_) & (~0x07)); + } + + // ...... + + uint32_t kind() const { + return reinterpret_cast(first_use_offseted_by_kind_) & 0x07; + } + + // ...... +}; +``` + +再来看 `ValueImpl` 的构造函数, 传入 `Type` 类型和 `kind` 数. 初始情况下, `UD链` 为空, 于是用空指针 `nullptr` 初始化 `OpOperandImpl *` 指针加上 `kind` 赋值到 `first_use_offseted_by_kind_`. + +```c++ +// 代码位置 paddle/pir/core/value_impl.cc + +#define OUTLINE_RESULT_IDX 6u +#define MAX_INLINE_RESULT_IDX (OUTLINE_RESULT_IDX - 1u) +#define BLOCK_ARG_IDX (OUTLINE_RESULT_IDX + 1u) + +// ...... + +ValueImpl::ValueImpl(Type type, uint32_t kind) { + if (kind > BLOCK_ARG_IDX) { + LOG(FATAL) << "The kind of value_impl(" << kind + << "), is bigger than BLOCK_ARG_IDX(7)"; + } + type_ = type; + first_use_offseted_by_kind_ = reinterpret_cast( + reinterpret_cast(nullptr) + kind); + VLOG(4) << "Construct a ValueImpl whose's kind is " << kind + << ". The offset first_use address is: " + << first_use_offseted_by_kind_; +} +``` + +`set_first_use` 函数也是直接将 `first_use` + `offset` 之和赋值给 `first_use_offseted_by_kind_`. `offset` 就是前文提到的 `kind`. + +```c++ +// 代码位置 paddle/pir/core/value_impl.cc + +void ValueImpl::set_first_use(OpOperandImpl *first_use) { + uint32_t offset = kind(); + first_use_offseted_by_kind_ = reinterpret_cast( + reinterpret_cast(first_use) + offset); + VLOG(4) << "The index of this value is " << offset + << ". Offset and set first use: " << first_use << " -> " + << first_use_offseted_by_kind_ << "."; +} +``` + +注意, 也许是历史遗留问题, 此处的 `kind` 在其他的代码里也许叫 `offset` 或者 `index`, 需要注意对应关系. (可以的话, 之后提个PR修改一下这个变量的名字) + +接下来我们看一个单测来了解 `kind` 的含义, `builder.Build` 的模板参数为你要传入的 `Op`, 参数是构造该 `Op` 需要传入的参数 + +```c++ + // 代码位置 test/cpp/pir/core/ir_region_test.cc + // (1) Init environment. + pir::IrContext* ctx = pir::IrContext::Instance(); + + // (2) Create an empty program object + pir::Program program(ctx); + pir::Builder builder = pir::Builder(ctx, program.block()); + + // (3) Def a = ConstantOp("2.0"); b = ConstantOp("2.0"); + pir::FloatAttribute fp_attr = builder.float_attr(2.0f); + pir::Float32Type fp32_type = builder.float32_type(); + pir::OpResult a = + builder.Build(fp_attr, fp32_type)->result(0); + pir::OpResult b = + builder.Build(fp_attr, fp32_type)->result(0); +``` + +来简单看一眼 `Builder.Build` 的源码, 通过 `Op` 的名字从 `IrContext` 中获取 `OpInfo` 来创建一个 `OperationArgument` 对象 `argument`. + + +```c++ +// 代码位置在 paddle/pir/core/builder.h +template +OpTy Builder::Build(Args &&...args) { + OperationArgument argument(context_->GetRegisteredOpInfo(OpTy::name())); + OpTy::Build(*this, argument, std::forward(args)...); + Operation *op = Build(std::move(argument)); + return OpTy(op); +} +``` + +然后调用传入 `Op` 自己(此处是 `ConstantOp` )的 `Build` 函数, 实参是当前对象(`*this`), `argument`, 将可变参数 `args` 转发, 此处的可变参数是`(fp_attr, fp32_type)`. + +那就再看一眼 `ConstantOp` 的 `Build` 函数: + + +```c++ +// 代码位置在 paddle/pir/core/builtin_op.cc +void ConstantOp::Build(Builder &builder, + OperationArgument &argument, + Attribute value, + Type output_type) { + argument.AddAttribute("value", value); + argument.output_types.push_back(output_type); +} +``` + +就是向 `argument` 的 `AttributeMap attributes` 添加了属性 `value`(此处是 `fp_attr`), 并向 `argument` 的输出类型 `std::vector output_types` 添加输出类型(此处是 `fp32_type`). + +由于 `argument` 是以引用的形式传入, 所以其值也被修改. 之后调用另一个 `Build` 的重载函数去构建 `Op`, 其内部调用了 `Operation::Create` 来创建, 并返回一个 `Operation *` 指针. 此处不去深究 `Insert` 函数做了什么. + + +```c++ +// 代码位置 paddle/pir/core/builder.cc +Operation *Builder::Build(OperationArgument &&argument) { + return Insert(Operation::Create(std::move(argument))); +} +``` + + +总结一下, `builder.Build(fp_attr, fp32_type)` 会创建一个 `Operation` 并返回一个 `Operation *` 的指针. + +```c++ + // 代码位置 test/cpp/pir/core/ir_region_test.cc + pir::FloatAttribute fp_attr = builder.float_attr(2.0f); + pir::Float32Type fp32_type = builder.float32_type(); + pir::OpResult a = + builder.Build(fp_attr, fp32_type)->result(0); + pir::OpResult b = + builder.Build(fp_attr, fp32_type)->result(0); +``` + +所以上述代码后两行就是在执行 `Operation::result` 函数 + +```c++ + // 代码位置 paddle/pir/core/operation.h + OpResult result(uint32_t index) { return op_result_impl(index); } + // ...... + detail::OpResultImpl *op_result_impl(uint32_t index); +``` + +我们之前提到了 `op_result_impl` 和 `op_operand_impl` 是通过 `COMPONENT_IMPL` 宏来实现的, 我们将宏的部分实现出来: + + +```c++ +// 宏的位置在 paddle/pir/core/operation.cc 尾部 + +OpResultImpl *Operation::op_result_impl(uint32_t index) { + int32_t offset = ComputeOpResultOffset(index); + return reinterpret_cast( + reinterpret_cast(this) + offset); +} +const OpResultImpl *Operation:: ( + uint32_t index) const { + int32_t offset = ComputeOpResultOffset(index); + return reinterpret_cast( + reinterpret_cast(this) + offset); +} + +OpOperandImpl *Operation::op_operand_impl(uint32_t index) { + int32_t offset = ComputeOpOperandOffset(index); + return reinterpret_cast( + reinterpret_cast(this) + offset); +} +const OpOperandImpl *Operation::op_operand_impl( + uint32_t index) const { + int32_t offset = ComputeOpOperandOffset(index); + return reinterpret_cast( + reinterpret_cast(this) + offset); +} +``` + +可以看到第一个 `op_result_impl` 函数的参数就是 `index`, 这个 `index` 就是我们之前在 `ValueImpl` 中提到的低8位 `kind`. 通过传入的 `index` 调用 `ComputeOpResultOffset` 来计算内存中的偏移量 `offset`. + +我们来看 `Operation::ComputeOpResultOffset` 函数, 这种计算偏移的方式和 `OpInfo` 中计算偏移的方式类似, 或者说 `Operation` 创建时, 和 `OpInfo` 的创建方式类似, 其内容都会“紧凑”地排布在内存中. + +```c++ +// 代码位置 paddle/pir/core/operation.cc +int32_t Operation::ComputeOpResultOffset(uint32_t index) const { + if (index >= num_results_) { + LOG(FATAL) << "index exceeds OP op result range."; + } + if (index < OUTLINE_RESULT_IDX) { + return -static_cast((index + 1u) * sizeof(OpInlineResultImpl)); + } + constexpr uint32_t anchor = OUTLINE_RESULT_IDX * sizeof(OpInlineResultImpl); + index = index - MAX_INLINE_RESULT_IDX; + return -static_cast(index * sizeof(OpOutlineResultImpl) + anchor); +} +``` + +到此, 可以剧透一下, `kind` 或者说 `index` 就是 `Operation` 返回的 `OpResult` 索引, 因为一般 `Op` 一般很少有6个以上(0到5)的输出, 所以 `kind` 只用低3位去表示是够的. 但是如果有 `Op` 有6个以上甚至更多的输出怎么办? 这个问题暂时先不解答, 我们来看看 `Operation` 的创建过程. 可以分为8个小代码块. + + +第一部分, 计算一个 `Operation` 对象内容所需要的字节数, `max_inline_result_num` 在代码中是个定值, 为 `6`. 在计算 `OpResult` 类型变量的大小时, 如果输出个数 `<=6` 则开辟对应数量的 `OpInlineResultImpl` 内存大小, 如果输出个数 `>6` 则开辟 `6` 个 `OpInlineResultImpl`, 开辟 `输出个数 - 6` 个 `OpOutlineResultImpl`. + +```c++ +Operation *Operation::Create(const std::vector &inputs, + const AttributeMap &attributes, + const std::vector &output_types, + pir::OpInfo op_info, + size_t num_regions, + const std::vector &successors) { + + // 1. Calculate the required memory size for OpResults + Operation + + // OpOperands. + uint32_t num_results = output_types.size(); + uint32_t num_operands = inputs.size(); + uint32_t num_successors = successors.size(); + uint32_t max_inline_result_num = MAX_INLINE_RESULT_IDX + 1; + size_t result_mem_size = + num_results > max_inline_result_num + ? sizeof(detail::OpOutlineResultImpl) * + (num_results - max_inline_result_num) + + sizeof(detail::OpInlineResultImpl) * max_inline_result_num + : sizeof(detail::OpInlineResultImpl) * num_results; + size_t op_mem_size = sizeof(Operation); + size_t operand_mem_size = sizeof(detail::OpOperandImpl) * num_operands; + size_t block_operand_size = num_successors * sizeof(detail::BlockOperandImpl); + size_t region_mem_size = num_regions * sizeof(Region); + size_t base_size = result_mem_size + op_mem_size + operand_mem_size + + region_mem_size + block_operand_size; +``` + + +申请刚才计算的内存, 要求地址的低3位为0, 可以被8整除 +```c++ + // 2. Malloc memory. + char *base_ptr = reinterpret_cast(aligned_malloc(base_size, 8)); +``` + + + +在指定位置构建对应数量的 `OpOutlineResultImpl` 和 `OpInlineResultImpl` 输出对象. + +```c++ + // 3.1. Construct OpResults. + for (size_t idx = num_results; idx > 0; idx--) { + if (idx > max_inline_result_num) { + new (base_ptr) + detail::OpOutlineResultImpl(output_types[idx - 1], idx - 1); + base_ptr += sizeof(detail::OpOutlineResultImpl); + } else { + new (base_ptr) detail::OpInlineResultImpl(output_types[idx - 1], idx - 1); + base_ptr += sizeof(detail::OpInlineResultImpl); + } + } +``` + + +在指定位置构建一个 `Operation` 对象 + +```c++ + // 3.2. Construct Operation. + Operation *op = new (base_ptr) Operation(attributes, + op_info, + num_results, + num_operands, + num_regions, + num_successors); + base_ptr += sizeof(Operation); +``` + + +在指定位置构建 `OpOperandImpl` 对象, 用来描述输入. + +```c++ + // 3.3. Construct OpOperands. + if ((reinterpret_cast(base_ptr) & 0x7) != 0) { + IR_THROW("The address of OpOperandImpl must be divisible by 8."); + } + for (size_t idx = 0; idx < num_operands; idx++) { + new (base_ptr) detail::OpOperandImpl(inputs[idx], op); + base_ptr += sizeof(detail::OpOperandImpl); + } +``` + + +暂时未知 + +```c++ + // 3.4. Construct BlockOperands. + if (num_successors > 0) { + op->block_operands_ = + reinterpret_cast(base_ptr); + for (size_t idx = 0; idx < num_successors; idx++) { + new (base_ptr) detail::BlockOperandImpl(successors[idx], op); + base_ptr += sizeof(detail::BlockOperandImpl); + } + } +``` + + +暂时未知 + +```c++ + // 3.5. Construct Regions + if (num_regions > 0) { + op->regions_ = reinterpret_cast(base_ptr); + for (size_t idx = 0; idx < num_regions; idx++) { + new (base_ptr) Region(op); + base_ptr += sizeof(Region); + } + } +``` + + +验证 +```c++ + // 0. Verify + if (op_info) { + op_info.Verify(op); + } + return op; +} +``` + + +可以看到 `OpInlineResultImpl` 和 `OpOutlineResultImpl` 在当前版本, 除了 `index` 计算不同, 后者还比前者多了私有变量 `outline_index_`. + +```c++ +class OpInlineResultImpl : public OpResultImpl { + public: + OpInlineResultImpl(Type type, uint32_t result_index) + : OpResultImpl(type, result_index) { + if (result_index > MAX_INLINE_RESULT_IDX) { + throw("Inline result index should not exceed MaxInlineResultIndex(5)"); + } + } + + static bool classof(const ValueImpl &value) { + return value.kind() < OUTLINE_RESULT_IDX; + } + + uint32_t index() const { return kind(); } +}; + +/// +/// \brief OpOutlineResultImpl is the implementation of an operation result +/// whose index > 5. +/// +class OpOutlineResultImpl : public OpResultImpl { + public: + OpOutlineResultImpl(Type type, uint32_t outline_index) + : OpResultImpl(type, OUTLINE_RESULT_IDX), outline_index_(outline_index) {} + + static bool classof(const ValueImpl &value) { + return value.kind() == OUTLINE_RESULT_IDX; + } + + uint32_t index() const { return outline_index_; } + + private: + uint32_t outline_index_; +}; +``` + +到此, 我们回溯一下, 如果某个 `Op` 有6个以上甚至更多的输出怎么办? 前6个 `OpResult` 是 `OpInlineResult`, 其 `protected` 变量`first_use_offseted_by_kind_` 低3位是对应的索引`0~5`, 第7个输出及以上是 `OpOutlineResultImpl`, 其 `first_use_offseted_by_kind_` 低3位都是6, 且有一个私有变量 `outline_index_` 来记录其索引. + + + +到此, 我们新 IR 体系下的计算图小节终于介绍完毕. + + + +## 新 IR 体系下的 `Dialect` + +首先要注意 `dialect` 和 `Dialect` 的区别, 注意d小写的只是一个 `namespace`, D大写的是一个类 + +```c++ +namespace paddle { +namespace dialect { + // ...... +} // namespace dialect +} // namespace paddle +``` + +`Dialect` 基本上可以理解为 `namespace`. 在 `Dialect` 中,我们可以定义一系列类型、属性和 Op 等。 +`Dialect` 对象的实例会被加载到全局IrContext中。 +特定的编译器只需要结合现有的 `Dialect` 并添加自己的扩展或定制即可。 + +以下是 `BuiltinDialect` 的源码: + +```c++ +// 代码位置 paddle/pir/core/builtin_dialect.h +// Dialect 是一个类 +class IR_API BuiltinDialect : public pir::Dialect { + public: + explicit BuiltinDialect(pir::IrContext *context); + /// + /// \brief Each Dialect needs to provide a name function to return the name of + /// the Dialect. + /// + /// \return The name of this Dialect. + /// + static const char *name() { return "builtin"; } + + private: + void initialize(); +}; +``` + +可以看到在 `BuiltinDialect` 的实现中, 注册了一系列的类型、属性和 Op 等: + +```c++ +void BuiltinDialect::initialize() { + // Register all built-in types defined in builtin_type.h. + RegisterTypes(); + + RegisterAttributes(); + + RegisterOps(); +} +``` + + +## 新 IR 体系下的 IrContext + +主要位置在 `paddle/pir/core/ir_context.cc` 和 `paddle/pir/core/ir_context.h` + +`IrContext` 是一个全局无参数类,用于存储和管理Type、Attribute等相关数据结构。 + +## To Be Continued + + diff --git a/pfcc/paddle-code-reading/IR_Dialect/img/while_op_detail.jpg b/pfcc/paddle-code-reading/IR_Dialect/img/while_op_detail.jpg new file mode 100644 index 000000000..7b0e93d62 Binary files /dev/null and b/pfcc/paddle-code-reading/IR_Dialect/img/while_op_detail.jpg differ diff --git a/pfcc/paddle-code-reading/IR_Dialect/img/while_op_detail2.jpg b/pfcc/paddle-code-reading/IR_Dialect/img/while_op_detail2.jpg new file mode 100644 index 000000000..6bb384a15 Binary files /dev/null and b/pfcc/paddle-code-reading/IR_Dialect/img/while_op_detail2.jpg differ diff --git a/pfcc/paddle-code-reading/IR_Dialect/ir_program.md b/pfcc/paddle-code-reading/IR_Dialect/ir_program.md index ecb147472..25b6f38c2 100644 --- a/pfcc/paddle-code-reading/IR_Dialect/ir_program.md +++ b/pfcc/paddle-code-reading/IR_Dialect/ir_program.md @@ -419,7 +419,7 @@ Variable类似于paddle中的Varibale, 它包含: 3. `bool is_mutable_`: 表明数据是否会在模型的执行当中被改变; 4. 数据的大小、对齐等等其他性质。 -对于模型中的对权重的使用,我们定义 G`etParameterOp`、`SetParameterOp`。分别从相应模型的哈希表中, 获取、设置相应的权重内容。 +对于模型中的对权重的使用,我们定义 `GetParameterOp`、`SetParameterOp`。分别从相应模型的哈希表中, 获取、设置相应的权重内容。 其中, `GetParameterOp`接受一个字符串作为属性,意义是从该模型的哈希表中加载该字符串对应的属性,并将其转换为输出。 `SetParameterOp` 接受一个字符串作为属性,一个张量类型的输入,没有输出。 表示用该属性和张量组成的键值对更新模型权重哈希表。 diff --git a/pfcc/paddle-code-reading/IR_Dialect/program_translator.md b/pfcc/paddle-code-reading/IR_Dialect/program_translator.md index f5e44add1..c46d6eed7 100644 --- a/pfcc/paddle-code-reading/IR_Dialect/program_translator.md +++ b/pfcc/paddle-code-reading/IR_Dialect/program_translator.md @@ -295,7 +295,7 @@ void ProgramTranslator::Translate() { for (size_t block_idx = 0; block_idx < legacy_program_->Size(); block_idx++) { const BlockDesc& block = legacy_program_->Block(block_idx); - SetStopGradientAttributeForAllValue(block); + SetStopGradientAttributeForAllValue(block); // 将 OpDesc 逐个翻译为 Operation } for (size_t block_idx = 0; block_idx < legacy_program_->Size(); block_idx++) { @@ -317,7 +317,7 @@ void ProgramTranslator::Translate() { ```c++ void ProgramTranslator::InsertOperationToSingleBlock(const BlockDesc& block) { - auto& op_translator = OpTranslator::instance(); + auto& op_translator = OpTranslator::instance(); // OpTranslator 是一个单例 for (auto op : block.AllOps()) { OpTranslateFn& fn = op_translator[op->Type()]; if (op->Type() == "shadow_output") { @@ -456,7 +456,7 @@ ir::Operation* OpTranscriber::operator()(ir::IrContext* ctx, TranslationContext* param_map, const OpDesc& op_desc, ir::Program* program) { - auto op_info = this->LoopkUpOpInfo(ctx, op_desc); + auto op_info = this->LoopkUpOpInfo(ctx, op_desc); // 根据 Op name 获取 op info auto* op_info_concept = op_info.GetInterfaceImpl(); @@ -469,25 +469,26 @@ ir::Operation* OpTranscriber::operator()(ir::IrContext* ctx, this->InsertSliceOperationForInput( ctx, param_map, op_desc, input_infos, program); - auto op_inputs = this->GenerateOperationInput( + + auto op_inputs = this->GenerateOperationInput( // 获取 Input ctx, param_map, op_desc, op_info.name(), input_infos, program); OpOutputMapping arg_to_idx; OpOutputTypeList op_output_types; std::tie(op_output_types, arg_to_idx) = - this->GenerateOperationOutput(ctx, op_desc, output_infos); + this->GenerateOperationOutput(ctx, op_desc, output_infos); // 获取 output_types auto attribute_map = - this->TranslateOpAttribute(ctx, op_info.name(), attr_infos, op_desc); + this->TranslateOpAttribute(ctx, op_info.name(), attr_infos, op_desc); // 获取 attributes VLOG(4) << "[general op][" << op_desc.Type() << "] preparation end."; ir::Operation* operation = - ir::Operation::Create(op_inputs, attribute_map, op_output_types, op_info); + ir::Operation::Create(op_inputs, attribute_map, op_output_types, op_info); // 创建 operation VLOG(4) << "[general op][" << op_desc.Type() << "] opearation creation end."; - program->block()->push_back(operation); + program->block()->push_back(operation); // 插入到 program 中?为什么在这里插入? VLOG(4) << "[general op][" << op_desc.Type() << "] opearation insertion end."; - this->RecordOpResultMapping(ctx, param_map, op_desc, operation, arg_to_idx); + this->RecordOpResultMapping(ctx, param_map, op_desc, operation, arg_to_idx); // 记录 Var 和 value 的对应关系 return operation; } @@ -608,31 +609,34 @@ tensor<-1x2048x7x7xf32> OpCompatInfo 用于处理动静定义不一致的问题,通过扫描 op_compat.yaml 生成,其接口如下: ```c++ - std::string operator[](const std::string& op_type) { + std::string operator[](const std::string& op_type) { // 从 OpDesc.name 得到规范化的 Op 名字,方便查询 OpInfo if (op_name_mappings.find(op_type) == op_name_mappings.end()) { return op_type; } return op_name_mappings.at(op_type); } - bool HasMutableAttribute(const std::string& op_type) { + bool HasMutableAttribute(const std::string& op_type) { // 查询某个 Op 是否有可变 attribute return (op_mutable_attributes.find(op_type) != op_mutable_attributes.end()); } - const std::unordered_set* GetMutableAttributes( + const std::unordered_set* GetMutableAttributes( // 查询 Op 有哪些可变 attribute const std::string& op_type) { if (!HasMutableAttribute(op_type)) return nullptr; return &op_mutable_attributes.at(op_type); } + // 查询 Op 的可变 attribute 可能对应 OpDesc 里的哪些 Argument, 如 shape 会对应 ShapeTensor ShapeTensorList const MutableAttributeInfo& GetMutableAttributeInfos( const std::string& op_type, const std::string& arg_name) { return op_mutable_attribute_infos.at(op_type).at(arg_name); } + // 查询 Op 的 argument 对应 OpDesc 里的哪个 argument std::string GetLegacyArgName(const std::string& op_type, const std::string& arg_name); + // 查询 Op 的 attribute 对应 OpDesc 里的哪个 attribute std::string GetLegacyAttrName(const std::string& op_type, const std::string& arg_name); ``` @@ -786,13 +790,13 @@ TypeTranslator::TypeTranslator() { ```c++ -class AttributeVisitor; +class AttributeVisitor; // framework::Attribute 是 paddle::variant 类型,因此需要通过 visitor 访问 class AttributeTranslator { private: AttributeTranslator(); - AttributeVisitor* general_visitor; - std::unordered_map special_visitors; + AttributeVisitor* general_visitor; // 通用转换方式 + std::unordered_map special_visitors; // 对特定类型的转换方式 public: AttributeTranslator(const AttributeTranslator&) = delete;