Skip to content

新IR常量折叠

Jump to bottom
zhangyuqin1998 edited this page Oct 18, 2023 · 3 revisions

推理流程梳理

0. 训练阶段参数初始化与 scope 构建

  1. paddle.nn 中的函数会将参数初始化的 op 追加到 start program 中
  2. 执行器执行 start program
  3. 在 NewIRInterpreter 阶段 build scope，把 op 用到的输入输出 var 都构造出来并放到 global scope 里
  4. NewIRInterpreter 按序执行 op kernel。之后，初始化后的参数就在 global scope 里了
  5. 执行器执行 main program，可以从 global scope 里拿到参数
  6. ...

1. 执行器python前端

网络和参数都就绪之后，调用 Executor 的 run() 方法执行，这里以新执行器为例。新执行器会：

  1. 创建 StandaloneExecutor实例
  2. feed data
  3. 调用 StandaloneExecutor::run

2. 执行器c++后端

StandaloneExecutor 的构造函数依次做了几件重要的事情：

  1. TranslateLegacyProgramToProgram，旧 IR 被翻译成新 IR
  2. 常规 pass，如常量折叠、算子融合、死代码消除等
  3. PdOpLowerToKernelPass，新 IR 绑定到 kernel
  4. inplace pass
  5. 创建 InterpreterCore，将 impl 构造为 NewIRInterpreter （直到这一步才用到 scope_，之前都没用到）

其中真正负责执行的是 NewIRInterpreter。

2.1 TranslateLegacyProgramToProgram

  1. GetParameterForSingleBlock a. 遍历 BlockDesc 的所有 OpDesc，再遍历每个 OpDesc 的所有 input_var_name ⅰ. 判断是否 is_parameter，判断是否 is_unseen_variable ⅱ. 如果同时 is_parameter 和 is_unseen_variable，就说明 need_get_parameter_op ⅲ. 如果 need_get_parameter_op，InsertGetParamaterOp。同时维护一个全局从 var_name 到 OpResult 的映射。调用 program_->SetParameter(var_name, nullptr)** 这里第二个参数为空，这也是常量折叠pass有问题的根本原因，因为ir翻译期间并没有引入scope，导致参数数据无法在此时被加载到program里**
  2. TranslateBlock。根据不同的 op type 选择不同的 OpTranslateFn进行转换
  3. SetParameterFromSingleBlock。如果 output 在 parameter_name_mappings_中已经存在了，说明这个 param 需要更新，说明 need_set_parameter_op

2.2 PdOpLowerToKernelPass

主要的处理逻辑是在 ProcessBlock() 方法中，遍历了 block 的每个 op，根据 OpYamlInfoParser 绑定 kernel 信息。

2.3 NewIRInterpreter

NewIRInterpreter 在构造函数中主要做了 BuildScope。scope 除了存储参数，还要存储中间的 tensor 数据。参数已经在之前被放入到 scope 里了，但 tensor 还没有，所以需要为这些 tensor 分配好内存，并用 scope 管理。

在执行 NewIRInterpreter::Run 时：

  1. SolvePersisableVarNames。遍历所有 vars，把 Persisable 的 parameter 记录下来，parameter 不会被 gc 回收
  2. BuildInstruction。遍历所有 op 然后 BuildPhiContext。从 scope 里拿到 var， 再根据 var 拿到 op 的输入和输出 tensor，放到 Context 里。这样 Kernel 在调用时就可以从 Context 里拿到 tensor 了。
  3. PreAnalysis
  4. TraceRunImpl 或 MultiThreadRunImpl，根据拓扑序和依赖关系调用 Instruction。
  5. 根据 fetch list 获取返回值。

codegen

一个很自然的想法是，能不能通过 codegen 的方式为每个 op 生成 fold 方法？为此，我们梳理一下 op codegen 的大概流程 在 gen.py 的 OpGenerator 里：

  1. 根据 op_yaml_files 和 op_compat_yaml_file 解析生成 op_compat_item
  2. 构造 OpInfoParser(op, op_compat_item)
  3. 根据上面解析出的结果，格式化 OP_INFO_TEMPLATE，用于生成算子的 GetOpInfo() 方法。这东西很重要。OpInfoTuple 把所有元信息和 op 打包在一起
  4. 根据模板生成头文件
  5. 根据模板 CC_FILE_TEMPLATE 生成源文件

理论上拿到 OpInfoTuple 就可以找到 kernel 并构建 InferMetaContext 和 KernelContext 了。只要把 KernelContext 传给 kernel 就能执行了。（需要先 BuildScope，为 kernel 的 output 在 scope 里申请出空间） 所以我们可以看出，并不需要额外的 codegen 了。想运行指定的 op，只要拿到它的 OpInfoTuple，然后传入 scope 就可以。可以提供类似如下的接口：

DenseTensor runSingleOp(Pperation* op, Scope* scope) {
  1. BuildScope
  2. 获取 OpInfoTuple
  3. 构建 InferMetaContext
  4. 构建 KernelContext
  5. run infer_meta
  6. run phi_kernel
  7. 从 scope 里找出 result 然后返回
}

新IR常量折叠问题分析

1. 新 IR 常量折叠的主要改写逻辑

  1. 根据 op 创建 temp_program
  2. PdOpLowerToKernelPass，根据 temp_program 构造 kernel_program
  3. 创建 InterpreterCore 并执行
  4. 根据 fetch_list 取出 out_tensor
  5. 根据 out_tensor 构造 pir::Parameter
  6. 在 scope_ 中创建新 param_var，并把 out_tensor 赋值给 param_var
  7. 在原 program 中 SetParameter
  8. 插入 get_parameter_op 并替换子图

2. 问题分析

可以看到，新IR常量折叠里主要存在两个问题：

问题一：需要从原 program 中获取 Parameter，但我们在 TranslateLegacyProgramToProgram 时，只记录了 param_name，并没有创建 Parameter 对象（见上面流程梳理3.1）。目前我们仍然是把 param 以 var 的形式存储在 scope 里了。而在 pass 阶段，并没有传入 scope，也就没法获得原 program 的参数。

问题二：需要单独构造出 temp program，并构造 InterpreterCore。其实可以直接调用 op 的 kernel。

3. 解法

对于问题一，在 pass 阶段获取参数有两种做法：

  • 把 scope 传入 TranslateLegacyProgramToProgram 阶段，这样可以从 scope 里拿到参数，创建 Parameter 对象，再通过 SetParameter 设置到 program 里。这样做的话，新 IR 体系下，就不需要把参数存储在 scope 里了。
  • 把 scope 传入 PassManager，让 pass 在执行时也可以拿到参数。这样做就不需要在 program 的 Parameter 里维护 data 指针了。

个人倾向于第二种解法。从逻辑角度讲，既然我们认为 param 和 tensor 都是 variable，那么用同一种方式管理他们的数据才是更符合直觉的。都用 scope 来管理就足够了，没必要在 program 里的 Parameter 维护 data 指针。

对于问题二，由于我们在进入常规 pass 阶段时还没有完成 PdOpLowerToKernelPass 了，所以需要先将 op 绑定到 kernel 进行了绑定，再 SelectKernel 然后 BuildPhiContext，利用 kernel 和 scope 中的 param 完成计算，并将计算结果作为新的 param 存储在 scope 里。之后和原来一样插入 get_parameter_op 并替换子图就 OK 了。（其实这个逻辑已经有类似的实现了，即 PhiKernelAdaptor。我们可以构造临时 program，然后复用 PhiKernelAdaptor 算出结果）

MLIR 调研

MLIR 基础设施提供两种规范化的方法：

一般的Pattern Rewriter

mlir 的 pattern rewiter 提供了常量折叠的 pass，但这块我还不是很理解怎么根据 GenericOp 执行计算（用 llvm 吗？），后面再仔细看看

template <typename ConcreteType>
class FoldConstantBase : public OpRewritePattern<GenericOp> {
	LogicalResult matchAndRewrite(GenericOp genericOp,
                                PatternRewriter &rewriter) const override {
        // 拿到 kernel ?
        RegionComputationFn computeFn =
        	static_cast<const ConcreteType *>(this)->getRegionComputeFn(genericOp);
        // 计算
        // return ConstantOp
    }
};

fold

MLIR 的Op可能会覆盖fold来实现，与一般的 DAG 到 DAG Pattern匹配器相比，它暴露了一个更简单的 API，并适用于通用的匹配器不适用的情况。 例如，DAG 重写可以删除当前函数中的任意节点，这可能会使迭代器无效。 作为 API 的常量折叠则不会删除任何节点，它只是提供一个常量值（列表）并允许客户端根据需要更新其数据结构。在规范化pass之外 ,fold在Dialect Conversion基础架构中用作合法化机制，并且可以通过OpBuilder::createOrFold在任何地方使用OpBuilder直接调用。fold 的限制是不能创建新的Op，只能替换根Op（但不能删除）。 它允许原地更新Op，或返回一组预先存在的值（或属性）以替换Op。

TVM 调研

tvm 和 paddle 的做法比较类似，也是构造了一个 dlcontext 并设置设备为 cpu，创建 llvm target 然后构造了一个 CreateInterpreter，执行了 CreateInterpreter::Eval

Expr FoldConstant(const Expr& expr, const IRModule& mod) {
	using tvm::transform::PassContext;
	DLContext ctx;
	ctx.device_type = kDLCPU;
	ctx.device_id = 0;
	Target target = Target::Create("llvm");
	// use a fresh build context
	// in case we are already in a build context.
	With<PassContext> fresh_build_ctx(PassContext::Create());

 return ConstantFolder(CreateInterpreter(mod, ctx, target), mod).Mutate(expr);
}

OneFlow 调研

看起来 oneflow 是通过 mlir 的 fold 的一套方法实现的常量折叠 oneflow 主要是实现了 UnaryFold 和 BinaryFold 两种基本的折叠方式。上层 op 根据自己的参数数目直接复用 UnaryFold 和 BinaryFold，并且 op 在 fold 里硬编码自己的 kernel。

OpFoldResult UnaryFold(MLIRContext* ctx, ArrayRef<Attribute> operands,
                       const std::function<MaybeTensor(const TensorPtr&)>& f) {
  ::oneflow::LazyMode::Guard guard{false};
  if (!operands.front()) { return {}; }  // Important!

  const auto attr_dict = operands.front().cast<mlir::DictionaryAttr>();
  auto attrs = NamedAttrList(attr_dict);
  const auto tensor = support::DenseElementsAttrToTensor(
      attr_dict.get("value"), attr_dict.get(OpTrait::IsOpConfCompatible<void>::getDeviceTagAttr()),
      attr_dict.get(OpTrait::IsOpConfCompatible<void>::getDeviceNameAttr()));
  const auto result = f(tensor).GetPtrOrThrow();
  attrs.set("value", support::TensorToDenseElementsAttr(result, ctx));
  attrs.set(OpTrait::IsOpConfCompatible<void>::getOpNameAttr(), GenNewVariableOpName(ctx));
  attrs.set(OpTrait::TensorSource<void>::getDataTypeAttrName(),
            attr_dict.get(OpTrait::TensorSource<void>::getDataTypeAttrName()));

  return attrs.getDictionary(ctx);
}

OpFoldResult BinaryFold(MLIRContext* ctx, ArrayRef<Attribute> operands,
                        const std::function<MaybeTensor(const TensorPtr&, const TensorPtr&)>& f) {
  ::oneflow::LazyMode::Guard guard{false};
  if (!(operands.front() && operands.back())) { return {}; }  // Important!
  auto lhs_attr_dict = operands.front().cast<mlir::DictionaryAttr>();
  auto rhs_attr_dict = operands.back().cast<mlir::DictionaryAttr>();
  if (!DictionaryAttrsHaveSameDataType({lhs_attr_dict, rhs_attr_dict})) {
    llvm::errs()
        << "Input tensors should have same data type in binary operation of constant folding."
        << "\n";
    return nullptr;
  }
}

//////////////////////////////////////////


OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();
  return UnaryFold(getContext(), operands, [this](const auto& tensor) {
    std::vector<int32_t> perm_;
    for (auto& x : getPerm().getValue()) { perm_.emplace_back(x.cast<IntegerAttr>().getSInt()); }
    return functional::Transpose(tensor, perm_);
  });
}

OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();
  return UnaryFold(getContext(), operands, [this](const auto& tensor) {
    std::vector<int64_t> shape_vec;
    for (auto& x : getShape().getValue()) {
      shape_vec.emplace_back(x.cast<mlir::IntegerAttr>().getValue().getSExtValue());
    }
    return functional::Reshape(
        tensor, ::oneflow::Shape(::oneflow::DimVector(shape_vec.begin(), shape_vec.end())));
  });
}

OpFoldResult ScalarAddOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();
  return UnaryFold(getContext(), operands, [this](const auto& tensor) -> MaybeTensor {
    if (getHasIntOperand()) { return functional::ScalarAdd(tensor, getIntOperand(), 1, false); }
    if (getHasFloatOperand()) {
      return functional::ScalarAdd(tensor, getFloatOperand().convertToDouble(), 1, false);
    }
    emitError("Scalar op must has a int operand or a float operand.");
    return TensorPtr();
  });
}

OpFoldResult SqrtOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();
  return UnaryFold(getContext(), operands, functional::Sqrt);
}

OpFoldResult BroadcastMulOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();
  return BinaryFold(getContext(), operands, functional::Mul);
}

OpFoldResult BroadcastDivOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();
  return BinaryFold(getContext(), operands, functional::Div);
}

OpFoldResult BroadcastSubOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();
  return BinaryFold(getContext(), operands, [](const auto& lhs, const auto& rhs) -> MaybeTensor {
    return functional::Sub(lhs, rhs, /*alpha=*/1.0, false);
  });
}

但比较奇怪的是，只有少量几个 op 实现了 fold 方法，其它 op 难道无法折叠吗？

进行折叠的控制逻辑还是在 mlir 里。调用链如下：

  1. mlir::applyOpPatternsAndFold
  2. MultiOpPatternRewriteDriver::simplify
  3. GreedyPatternRewriteDriver::processWorklist
  4. OperationFolder::tryToFold
  5. op 覆盖的 fold 方法
Clone this wiki locally