tensorflow xla 符号执行

tensorflow xla里的tf2xla部分是通过符号执行来完成的。

什么是符号执行

符号执行是一种程序分析技术，使用符号执行分析一个程序时，该程序会使用符号值作为输入，而非一般执行程序时使用的具体值。

比如下面的程序（来自wiki）

int f() {
  ...
  y = read();
  z = y * 2;
  if (z == 12) {
    fail();
  } else {
    printf("OK");
  }
}

一般我们执行这个程序，会读取一个y的值，比如 y = 3，那么z=6，此时程序进入fail分支，如果y=6，则最终打印OK，每次程序运行都会读取一个具体的数值，最终也会得到一个具体的结果。但对于符号执行来说，我们会给y一个值，比如y = λ，其中λ只是代表输入的一个符号，所以到了if判断的地方，没法具体得到这个值是true还是false，执行路径就会分叉，最终形成一颗树的结构。当程序最终执行完，从这个生成树的叶子节点到跟节点就会有很多路径，利用这些路径信息就可以对程序分析做很多事情。

参考：

https://en.wikipedia.org/wiki/Symbolic_execution

tensorflow中的符号执行

从定义可以看出，符号执行首先需要一个符号输入，最终得到的也是一个（或多个）带符号的输出，所以这里就有符号值（symbolic value）的概念。这也正好对应xla最终生成的函数的输入和输出，正好对应了tf2xla中CompileGraph的主流程：

BuildArguments 处理输入
ExecuteGraph 符号执行
BuildComputation 处理输出

BuildArguments

xla的参数有几种类型，符号输入主要涉及的kResource和kParameter；主要就是通过 xla::Parameter来构造输入，示例代码：

tuple = xla::Parameter(builder, 0, (*input_shapes)[0], "arg_tuple");

arg_handles[i] = xla::Parameter(builder, i, (*input_shapes)[i], absl::StrCat("arg", i));

ExecuteGraph

这个函数里主要就是调用了GraphCompiler，对图中每个op进行编译（符号执行）。这里的执行分为2步：

创建kernel
编译op（执行）：执行的过程中会构造xla的图

据一个 tf.roll的例子，截取这个op在xla中的kernel实现如下：

void Compile(XlaOpKernelContext* ctx) override {
    const TensorShape input_shape = ctx->InputShape(0);
    xla::XlaOp shift = ctx->Input(1);
    const TensorShape shift_shape = ctx->InputShape(1);
    const TensorShape axis_shape = ctx->InputShape(2)

    xla::Literal axis;

    xla::XlaOp output = ctx->Input(0);
    xla::PrimitiveType shift_type = ctx->input_xla_type(1);
    int64 num_axes = axis_shape.dims() == 0 ? 1 : axis_shape.dim_size(0);
    for (int64 i = 0; i != num_axes; ++i) {
      auto cur_axis_status = axis_shape.dims() == 0
                                 ? axis.GetIntegralAsS64({})
                                 : axis.GetIntegralAsS64({i});
      int64 cur_axis = cur_axis_status.ValueOrDie();

       xla::XlaOp offset =
          shift_shape.dims() == 0
              ? shift
              : xla::Reshape(xla::SliceInDim(shift, /*start_index=*/i,
                                             /*limit_index=*/i + 1,
                                             /*stride=*/1, /*dimno=*/0),
                             {});
      xla::XlaOp axis_size = xla::ConstantR0WithType(
          ctx->builder(), shift_type, input_shape.dim_size(cur_axis));
      // Adjust large offsets into [0, axis_size). This also makes negative
      // offsets positive.
      offset = ((offset % axis_size) + axis_size) % axis_size;

       // Stack two copies of the dimension, then slice from the calculated
      // offset.
      xla::XlaOp concat =
          xla::ConcatInDim(ctx->builder(), {output, output}, cur_axis);
      std::vector start_indices(
          input_shape.dims(), xla::Zero(ctx->builder(), shift_type));
      start_indices[cur_axis] = axis_size - offset;
      output =
          xla::DynamicSlice(concat, start_indices, input_shape.dim_sizes());
    }
    ctx->SetOutput(0, output);
  }

参考：

[TF:XLA] Implement support for tf.roll · tensorflow/tensorflow@cc3bd2c

这个op的作用就是对向量或矩阵按某个纬度做平移变换：

# 't' is [0, 1, 2, 3, 4]
roll(t, shift=2, axis=0) ==> [3, 4, 0, 1, 2]

# shifting along multiple dimensions
# 't' is [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]
roll(t, shift=[1, -2], axis=[0, 1]) ==> [[7, 8, 9, 5, 6], [2, 3, 4, 0, 1]]

# shifting along the same axis multiple times
# 't' is [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]
roll(t, shift=[2, -3], axis=[1, 1]) ==> [[1, 2, 3, 4, 0], [6, 7, 8, 9, 5]]

参考

https://www.tensorflow.org/api_docs/python/tf/roll

里面最主要的3行代码在这里：

xla::XlaOp concat = xla::ConcatInDim(ctx->builder(), {output, output}, cur_axis);

output = xla::DynamicSlice(concat, start_indices, input_shape.dim_sizes());

ctx->SetOutput(0, output);

前2行是用xla的builder 接口构造xla的图，每个接口代表xla hlo图中的一个op，但这个图的构造就是通过对原图进行执行（编译）得到的，最后把输出作为xla中op的输出。

BuildComputation

包含了对输出节点以及kResource的处理，保证输出节点一定在hlo computation的最后。可以参考源码注释：

  // Builds the XLA computation. We *always* form a tuple here to ensure that
  // the output value is the last thing added into the XLA computation, even
  // if there is only one output value.

original refer

Pete Yuan