Tensorflow XlaOpKernel | tf2xla 机制详解

compiler/aot/ 以AOT的方式将tf2xla/接入TF引擎
compiler/jit/以JIT的方式将tf2xla/接入TF引擎, 核心是7个优化器和3个tfop,其中XlaCompileOp调用tf2xla的“编译”入口完成功能封装,XlaRunOp调用xla/client完成“运行”功能。
compiler/tf2xla/对上提供xla_compiler.cc:XlaCompiler::CompileFunction()供jit:compile_fn()使用将cluster转化为XlaComputation。核心是利用xla/client提供的接口,实现对XlaOpKernel的“Symbolic Execution”功能。每个XlaOpKernel子类均做的以下工作: **从XlaOpKernelContext中取出XlaExpression或XlaOp, 调用xla/client/xla_buidler.h提供的方法完成计算, 将计算结果的XlaOp存入XlaKernelContext.**
compiler/xla/client/ 对上提供xla_builder.cc:Builder等供CompileFunction()使用,将Graph由Op表达转化为HloModuleProto:HloComputationProto:HloInstructionProto表达并保存在XlaComputation中。
对上提供local_client.cc:LocalClient::Compile(),作为编译入口供jit:BuildExecutable()使用,将已经得到的XlaComputation交给service并进一步编译为二进制。
对上提供local_client.cc:LocalExecutable::Run(),作为运行入口供jit/kernels/xla_ops.cc:XlaRunOp使用,通过Key找到相应的二进制交给service层处理
compiler/xla/service/ 对上提供local_service.cc:LocalService::BuildExecutable()供LocalClient::Compile()使用实现真正的编译,承接XlaComputation封装的HloProto, 将其转化为HloModule:HloComputation:HloInstruction表达, 对其进行优化之后, 使用LLVM后端将其编译为相应Executable后端的二进制代码
对上提供executable.cc:Executable::ExecuteOnStream()供LocalExecutable::Run()使用实现真正的执行二进制。

从Kernel的视角, XLA并不会新增Op, 而是针对已有的Op, 新增了基于XLA的另一个版本的Kernel: XlaOpKerne。在TF引擎中, OpKernel在软件栈上已是底层, 即最终的计算过程都要在OpKernel中实现. 但在XLA中, XlaOpKernel只是编译的入口, 大量的实际工作都交给了更下层的XLA引擎去完成.XLA相关的代码在tensorflow/compiler中.

tf2xla/负责XlaOpKernel的构造, 注册. 虽然XLA与TF引擎不在一层, 但二者面临的问题有很多有相似之处, 比如都需要对Kernel和Device保持易扩展性, 都需要维持前驱/后继Kernel的数据流和控制流关系. 基于类似的种种原因, XLA内部实现的注册XlaOpKernel的接口与TF引擎中注册OpKernel的风格十分相似, 同时, 其内部实现又有本质的不同, 而这些”不同”, 正是我们需要关注的.

要理解XlaOpKernel与OpKernel的不同, 关键在于了解”Symbolic Execution“.
先来看TF引擎, 它的OpKernel::Compute()方法要: OpKernelContext.Input()取输入数据 ==> 计算 ==> OpKernelContext.SetOutput()存输出数据, 计算结果继续通过OpKernelContext流入后继Opkernel, 其中流动的是真正的训练数据, 暂且将这个过程称之为”Execution”.
对比之下, XLA中的”Symbolic Execution”中的”Symbolic”即是说, XlaOpKernel的设计目的不在于去处理训练数据, 而在于去生成能够正确的处理数据的代码.” 这类似于C++模板, SWIG等生成代码的设计思想. 具体地, 在XlaOpKernel::Compile()中: XlaOpKernelContext.Input()以XlaOp形式取输入 ==> 调用xla/client/xla_buidler.h提供的方法实现Op该有的功能, 实际上是生成一组能处理数据的HloInstruction ==> XlaOpKernelContext.SetOutput()存储XlaOp形式的结果, 计算结果继续通过XlaOpKernelContext流入后继XlaOpkernel, 其中流动的都是以XlaOp表征的对训练数据的处理方法.

至于真正处理数据的时机, 就要交给XLA引擎, 它来负责后续的”编译”和”执行”, 具体地, 在JIT中, XlaCompileOp会在所有的XlaOpKernel::Compile()执行完毕之后, 继续调用xla/service中相应的方法将这些所有生成的HloInstruction编译生成二进制并进一步交给XlaRunOp来执行.

接口形式

"XlaOpKernel(compiler/tf2xla/xla_op_kernel.h)"继承自"OpKernel"(当前XlaOpKernel不支持AsynOpKernel类似的异步机制), 通过这种继承, XlaOpKernel与OpKernel注册框架有天然的相容性, 同时, 又针对XLA的设计要求作了以下处理, 来实现XLA需要的Symbolic Execution: 一个XlaOpKernel子类不再实现以OpKernelContext为参数的Compute()方法, 而要实现以XlaOpKernelContext为参数的Compile()方法.

同时, 依前文所述, XlaOpKernel的运行上下文, 输入输出与OpKernel有很大不同, 的这里, XLA重新封装了一个上下文类: "XlaOpkernelContext", 要注意到, 生成处理数据的XlaOp本身就和数据的形状, 类型等有关系, 即数据的元数据, 这些信息又存储在"OpKernelContext"中, 所以使用了”关联“OpKernelContext的方式来构造XlaOpKernelContext

XLA已经实现的OpKernel在 tensorflow/compiler/tf2xla/kernels/中, 实现一个新的XlaOpKernel, 子类需要实现"Compile()"方法, 并通过"REGISTER_XLA_OP"注册到系统中, 举个例子:

注册原理

和TF引擎中的OpKernel机制类似, XLA内部也使用了regsitry->registrar->registration->create_fn()的结构管理旗下的XlaOpKernel.

如此, 就添加了新的XlaOpKernel(实际上是XlaOpKernelRegistry::OpRegistration)注册到了XlaOpRegistry中, 但事情还远没有结束, 和TF引擎一样, XLA同样需要检索大量的XlaOpKernel, 但XLA无意重新实现一遍TF引擎已经实现过的Regsitration的管理机制, 而为了复用TF的实现, 除了继承OpKernel, 还需要在保留create_fn()的基础上, 将XlaOpKernelRegitry::OpRegistration转换为KernelRegistration, 如此就可以使用TF引擎的OpKernel管理机制. 具体的, 在JIT中, 由Optimization: MarkForCompilationPass中完成这种转换:

-9- 将device_type信息存储在KernelDef, XlaOpKernel就是靠device_type的不同才能与"GlobalKernelRegistry()"中OpKernel区分开, 在当前版本(1.14)中, XLA共注册3个backend, 所以此处的kdef取"DEVICE_GPU_XLA_JIT"等下述3个值之一, 关于"::tensorflow::XlaBackendRegistrar "我另文详述, 这里只需了解这些Backend的类型最终将作为XlaOpKernel的kdef.device_type_.

-12- 根据“Tensorflow OpKernel机制详解“一文, "kernel_factory::OpKernelRegistrar()"会调用"InitInternal()"将KEY和作为VALUE的"KernelRegistration"注册到"GlobalKernelRegistry()", 这里, factory即是REGISTER_XLA_OP`时create_fn(): 一个用于new 一个XlaOpKernel实例的方法. 至此, 我们注册了的XlaOpKernel就进入到了了GlobalKernelRegistry(), 这些XlaOpKernel可以通过TF引擎的通用接口“FindKernelRegistration()”来构造并获取. 此时, 在TF引擎中, 一个Op就有个多个Kernel: OpKernel + 多个适配了不同device的XlaOpKernel, Kernel之间彼此通过device_type进行区分.

Find原理

JIT将所有的XlaOpKernel注册到TF引擎, 那真正运行的时候如何找到相应的XlaOpKernel呢?这就涉及到了刚才一直在强调的问题: 注册到TF引擎时, 每一个OpKernelRegistrar都使用了JIT内部的设备类型. JIT在系统初始化的时候即注册了3个DeviceFactory:

根据Tensorflow的设计, 所有注册的DeviceFactory最终都会被Tensorflow执行引擎在初始化阶段调用"CreateDevices()"工厂模式”生产”相应的device实例, JIT注册的这三个也不例外, 以JIT内的GPU为例, 可以看到, 每一个device.device_type_都是"DEVICE_GPU_XLA_JIT".在OpKernel执行阶段, XlaCompileOp执行编译的时候结合上述的DEVICE_GPU_XLA_JIT等属性, 从GlobalKernelRegistry()中检索所需的Kernel;

-1- JIT编译执行入口, 初始化结束后, Tensorflow执行引擎执行XlaCompileOp进而进入”编译”阶段
-4-8- 将device中获取到的DEVICE_GPU_XLA_JIT 存入XlaCompilationCache
-10- 用XlaCompilationCache中的device_type构造XlaCompiler::Options, 进一步用于构造XlaCompiler
-18- 构造XlaCompiler::device_对象, 可以看到, 构造DeviceAttr使用的device_type_就是刚才传入的DEVICE_GPU_XLA_JIT
-19,21,23- 将存有DEVICE_GPU_XLA_JIT的Device laCompiler::device_存入XlaCompiler::flib_runtime_
-33- 构造GraphCompiler, 传入的flib即XlaCompiler::flib_runtime_, 存入GraphCompiler::flib_
-39- 此处使用的device追根溯源, 就是系统初始化CreateDevices()时构造的device, 即 DEVICE_GPU_XLA_JIT
-18,40- 遥相呼应
-46- 根据DEVICE_GPU_XLA_JIT构造检索用的key
-54,52- 执行create_fn(), 在XLA中即是构造XlaOpKernel.
-59- 执行OpKernel的入口
-61- 这里就是XlaOpKernel::Compute(), 如前文所述, 实质就是执行Compile()

另外, 关于使用device区别不同Kernel的讨论, 可以看到, TF引擎为OpKernel只提供了三种设备:"DEVICE_CPU", "DEVICE_GPU"和"DEVICE_SYCL"(core/framework/types.h), 可见, JIT通过构造虚拟设备的方式将两类Kernel巧妙的融合在一个数据结构中, 设计可谓巧妙. 类似的很多Linux内核程序也是借助了设备驱动接口注册虚拟设备实现和用户态的交互, 优秀的代码设计大体相似, 但垃圾的代码却各有各的垃圾法.

XlaOpKernel调试

和KernelRegistry一样, XlaOpRegistry也没有提供很多调试接口, 目前的版本只有这一个, 毕竟, 既然是开发人员, 就不能奢求太多

Leave a Reply

This site uses Akismet to reduce spam. Learn how your comment data is processed.

%d bloggers like this: