A Peek into TFRT

TensorFlow Runtime, TFRT, is a new runtime for TensorFlow. TFRT has some plausible designs we can spend some time on.

1. Koan-Sin Tan,
[email protected]
COSCUP, Aug 2nd, 2020
TensorFlow Runtime
A Peek into the Future of TensorFlow
1

2. •disclaimer: opinions are my own
•feel free to interrupt me if you have any questions during the presentation
•questions could be Taiwanese, English, or Mandarin
•most of TFRT materials are adapted from TFRT deep dive in MLIR design meeting [1] and TFRT docs [2]
•code around Aug 1, 2020 (git commit ecf1c20 [3])
[1] TFRT Deep Dive,   slides - recording, https://mlir.llvm.org/talks/
[2] https://github.com/tensorflow/runtime/tree/master/documents
[3] https://github.com/tensorflow/runtime/commit/ecf1c20
2

3. •Used open source before the term “open
source” is used
•A software guy, learned to use Unix and open
source software on VAX-11/780 running 4.3BSD
•Used to be a programming language junkie
•Worked on various system software, e.g., CPU
scheduling and power management of non-
CPU components
•Recently, on NN performance on edge devices
related stuff
•Contributed from time to time to TensorFlow Lite
•started a command line label_image for TFLite
who i am
https://gunkies.org/w/images/c/c1/DEC-VAX-11-780.jpg
3

4. What is TFRT
•TensorFlow Runtime (TFRT) is one of the two new MLIR runtimes emerged in 2020 so far.
•The other one is Intermediate Representation Execution Environment, IREE. It seems
so far tfrt has better design documentation
•Both of them have mobile / edge environment in mind.
•I didn’t see mobile accelerated code in TFRT yet.
•IREE has some Vulkan related code and some simple code works on Android already
•ResNet GPU inference is 28% faster with TFRT
•https://github.com/tensorflow/runtime, https://youtu.be/15tiQoPpuZ8
4

5. Build it
•if you follow the instructions described in README.md, it should just work. At least on x86_64 linux.
•however, it’s not tested for non Linux environment yet
•ssize_t and int64_t
•on Mac OS X: ssize_t: long, int64_t: long long
•current code mixed the use of ssize_t and int64_t
•test: one the acclaimed features of TFRT, like MLIR, is its use of 

LLVM FileCheck
•my hacks, shape related (ssize_t) tests not fixed yet
•it’s not tested on non-x86 platforms, such as aarch64, either 

•
5

6. •The three key directories under the TFRT root directory are
•lib: Contains core TFRT infrastructure code
•backends: Contains device specific infrastructure and op/kernel implementations
•include: Contains public header files for core TFRT infrastructure
6

7. Walking thru the tutorial
•unfortunately, it seems it’s not easy to jump directly into source code without having
some background knowledge
•so we’ll walk thru the tutorial [1]
•What are in the tutorial
•print hello world
•print integer
•adding kernels
[1] https://github.com/tensorflow/runtime/blob/master/documents/tutorial.md
7

8. using tfrt and tfrt_test
hello.mlir
func @hello() {
%chain = tfrt.new.chain
// Create a string containing "hello world" and store it in %hello.
%hello = "tfrt_test.get_string" () { string_attr = "hello world" } : () -> !tfrt.string
// Print the string in %hello.
"tfrt_test.print_string" (%hello, %chain) : (!tfrt.string, !tfrt.chain) -> !tfrt.chain
tfrt.return
}
The ‘@hello function above shows how to create and print a string. The text after each ‘:’ specifies the types involved:
•()->!tfrt.string means that tfrt_test.get_string takes no arguments and returns a !tfrt.string. tfrt is a
MLIR dialect prefix (or namespace) for TFRT
•(!tfrt.string, !tfrt.chain) -> !tfrt.chain means that tfrt_test.print_string takes two arguments (!
tfrt.string and !tfrt.chain) and returns a !tfrt.chain. chain [1] is a TFRT abstraction to manage dependencies
[1] https://github.com/tensorflow/runtime/blob/master/documents/explicit_dependency.md
8

9. hello world in MLIR
func @stringconstant() -> !llvm<"[12 x i8]"> {
%1 = llvm.constant("Hello world!") : !llvm<"i8*">
// CHECK: ret [12 x i8] c"Hello world!"
llvm.return %1 : !llvm<"i8*">
}
func @main() {
%0 = llvm.constant(0) : !llvm.i64
%1 = call @stringconstant() : () -> !llvm<"[12 x i8]">
%2 = llvm.getelementptr %1[%0] : (!llvm<"[12 x i8]">, !llvm.i64) -> !llvm<"i8*">
%3 = llvm.bitcast %2 : !llvm<"i8*"> to !llvm<"i8*">
%32 = llvm.call @puts(%2) : (!llvm<"i8*">) -> !llvm.i32
return
}
func @puts(!llvm<"i8*">) -> !llvm.i32
•MLIR “standard dialect” doesn’t have I/O functions
•there is LLVM dialect, of course we can use LLVM to call standard libc
function
9

10. Hello integer
func @hello_integers() {
%chain = tfrt.new.chain
// Create an integer containing 42.
%forty_two = tfrt.constant.i32 42
// Print 42.
tfrt.print.i32 %forty_two, %chain
tfrt.return
}
•as stated in the tutorial, we can run other functions in the same modular
•we can turn to more basic ones, such as integers or floating point numbers
•@hello_integers shows how to create and print integers
•This example does not have the verbose type information we saw in @hello because there are
custom parsers for the tfrt.constant.i32 and tfrt.print.32 kernels in
basic_kernels.td
10

11. basic_kernels.td
•.td (table description?) files are for LLVM TableGen
[1] TableGen, https://llvm.org/docs/TableGen/
class ConstantOp<string suffix, Type baseType, Attr attr>
: TFRT_Op<"constant." # suffix, [NoSideEffect]> {
let summary = "host executor constant value constructor" ;
let arguments = (ins attr:$value);
let results = (outs baseType);
}
class PrintOp<string suffix, Type type> : TFRT_Op< "print." # suffix> {
let summary = "tfrt.print operation" ;
let description = [{
An operation takes a number input and a chain input.
It prints the number to stdout and returns a chain output.
The chain input must be the second operand.
Example:
%2 = tfrt.print.i32 %0, %1
}];
let arguments = (ins type, TFRT_ChainType);
let results = (outs TFRT_ChainType);
let assemblyFormat = "operands attr-dict" ;
let verifier = ?;
}
https://github.com/tensorflow/runtime/blob/master/include/tfrt/basic_kernels/opdefs/basic_kernels.td#L376-L390
https://github.com/tensorflow/runtime/blob/master/include/tfrt/basic_kernels/opdefs/basic_kernels.td#L58-L64
11

12. Define kernels
12

13. user defined kernels
func @print_coordinate() {
%chain = tfrt.new.chain
%two = tfrt.constant.i32 2
%four = tfrt.constant.i32 4
%coordinate = "my.create_coordinate" (%two, %four) : (i32, i32) -> !my.coordinate
"my.print_coordinate" (%coordinate, %chain) : (!my.coordinate, !tfrt.chain) -> !tfrt.chain
tfrt.return
}
coordinate.mlir shows several TFRT features:
•MLIR types that begin with exclamation mark ( !) are user-defined types like !my.coordinate,
compared to built-in types like i32
•Kernels are just C++ functions with a name in MLIR: my.print_coordinate is the MLIR name for
the C++ PrintCoordinate function
•Kernels may pass arbitrary user-defined types: my.create_coordinate passes a custom
Coordinate struct to my.print_coordinate
13

14. to dig into some code we need
more system information
14

15. Host Runtime
15

16. •TensorFlow user passes into TFRT a
TensorFlow graph created via high-level
TensorFlow APIs, and
•TFRT then calls the  MLIR-based graph
compiler to optimize and lower the
graph into  BEF, a Binary Executable
Format for TFRT graph execution (MLIR
is the compiler infrastructure that we
use to represent TFRT host programs).
•The blue arrows in the simplified
TensorFlow training stack diagram
show this flow.
16

17. •In the README.md we are told to build two
binaries: tfrt_translate and bef_excutor
•tfrt_translate
•The tfrt_translate program does round trip
translation between MLIR and BEF, similar
to an assembler and disassembler.
•bef_executor
•The bef_executor program is the
execution driver of BEF files. It reads in a
BEF file, sets up runtime, and
asynchronously executes function(s) in
that file.
17

18. TFRT Host Runtime
•Foundation of TFRT: schedules work on the host and devices
•Clean separation between host and device runtimes:
•Host runtime does not know anything about devices, just their runtimes (sets of kernels)
•Key design points:
•Fully asynchronous - kernel executions can not block
•Excellent error propagation in the presence of asynchrony
•Performance as a first-class concern, for graph and eager
•Outline:
•Common runtime infrastructure
•Graph execution
•Op-by-op execution (“eager”)
18

19. •Container for data or resources
•Not Tensor specific
•A “future” type, fulfilled with exactly one value, or an error
•Lock-free, low memory overhead, type erased, reference
counted& "
•Helper class AsyncValueRef<T> provides type safety when
contained type is known
•AsyncValues enable efficient asynchronous compute
•Asynchronous functions return unavailable AsyncValues
•Caller can schedule dependent
computations with AsyncValue::AndThen()
•Caller need not block until AsyncValue
becomes available
Key Abstraction: AsyncValue
https://github.com/tensorflow/runtime/blob/master/include/tfrt/host_context/async_value.h
19

20. Kernels
•Kernel: unit of computation scheduled by the runtime
•Similar to kernel concept in current TensorFlow
•Kernels accept AsyncValue inputs and produce AsyncValue output
•Runtime coordinates dataflow of AsyncValues between kernels
•Outputs may not be immediately available, unlike current TensorFlow
•Runtime generally does not understand kernel semantics
// Kernel that adds two integers.
// AsyncKernelFrame holds the kernel’s arguments and results.
static void TFRTAdd(AsyncKernelFrame* frame) {
// Fetch the kernel’s 0th argument.
AsyncValue* arg1 = frame->GetArgAt(0);
// Fetch the kernel’s 1st argument.
AsyncValue* arg2 = frame->GetArgAt(1);
int v1 = arg1->get<int>();
int v2 = arg2->get<int>();
// Set the kernel’s 0th result.
frame->EmplaceResultAt<int>(0, v1 + v2);
}
https://github.com/tensorflow/runtime/blob/master/documents/tfrt_host_runtime_design.md
https://github.com/tensorflow/runtime/blob/master/lib/basic_kernels/integer_kernels.cc#L39-L45
https://github.com/tensorflow/runtime/blob/master/include/tfrt/host_context/kernel_utils.h#L61-L149
20

21. Host Program
•Host programs encode a dataflow graph
•Similar to GraphDef in current TensorFlow
•Expressed in MLIR. Typically compiler generated
•Designed for low-level dispatch efficiency
•Designed for compiler transformations and analysis, e.g.,
•Use dataflow analysis for buffer reuse



func @sample_function() -> i32 {
%one = tfrt.constant.i32 1 // Make AsyncValue with value 1
%two = tfrt.constant.i32 2 // Make AsyncValue with value 2
%three = tfrt.add.i32 %one, %two // Make AsyncValue with value 3 (1+2)
%ch0 = tfrt.new.chain
tfrt.print.i32 %three, %ch0 // Print AsyncValue %three
tfrt.return %three : i32 // Return AsyncValue %three
}
21

22. TFRT Binary Executable Format (BEF)
•BEF encodes a hardware-specific lowered graph
function
•Primary interface between compiler and runtime 

•Designed for efficient execution
•Low overhead: execute program by reading mmap’d
byte array 

•Persistent and stable: Compile once offline, run
many times 

online. Great for inference use-cases 

•Composed of sections, similar to ELF. Each section
has its own format 

•Extensible: BEF is versioned, reader ignores unknown
sections, new versions may define new sections 
 https://github.com/tensorflow/runtime/blob/master/documents/binary_executable_format.md
22

23. BEF Executor
•BEF Executor evaluates a BEF dataflow graph “executor” style:
•Not a bytecode-like interpreter: no concept of program counter
•“Strict” execution by default: run a kernel only when all its inputs are available
•Executor features:
•Lock-free: atomics instead of mutexes
•Non-blocking: defer dependent work with AsyncValue::AndThen
•Supports “non-strict” execution: may run a kernel when some of its
inputs are available
•Good for efficiently forwarding unavailable inputs to outputs
•Key concepts:
•BEF: dataflow graph
•Kernel: dataflow node
•AsyncValues: dataflow edge
https://github.com/tensorflow/runtime/blob/master/lib/bef_executor/bef_interpreter.cc#L223-L25423

24. Host Runtime Summary
24

25. How about Core Runtime?
•Surely, we can do similar walkthrough, but that will takes more time
•Two things
•Op Execution API, Execute()
•BEF Executor can handle it too



void CoreRuntime::Impl::Execute(const ExecutionContext& exec_ctx,
string_view op_name, OpHandler* op_handler,
MutableArrayRef<TensorHandle> arguments,
const OpAttrsRef& attrs,
MutableArrayRef<TensorHandle> results,
AsyncValueRef<Chain>* chain) {
// Ask the op_handler to execute the op. If successful, we're done.
auto op_handle = op_handler->MakeOp(op_name);
if (op_handle) {
op_handle.get()(exec_ctx, arguments, attrs, results, chain);
return;
}
// Otherwise, we fail with an 'unknown op' error.
auto err =
EmitErrorAsync(exec_ctx, "op '" + op_name.str() + "' is not supported" );
for (auto& result : results) result = TensorHandle(err.CopyRef());
if (chain) *chain = std::move(err);
}
25
https://github.com/tensorflow/runtime/blob/master/lib/core_runtime/core_runtime.cc#L124-L143
https://github.com/tensorflow/runtime/blob/master/documents/
tfrt_op_by_op_execution_design.md

26. BEF Executor for “op” graph
•corert.executeop
•sample



26

https://github.com/tensorflow/runtime/blob/master/lib/core_runtime/kernels.cc
func @example() -> !tfrt.chain {
%cpu = corert.get_op_handler( "cpu")
// Create TensorHandles
%lhs = corert.executeop( %cpu)
"test.create_dense_tensor" () { shape = [1, 1], values = [-1.0 : f32] }
%rhs = corert.executeop( %cpu)
"test.create_dense_tensor" () { shape = [1, 1], values = [-2.0 : f32] }
%result = corert.executeop( %cpu) "test.add" (%lhs, %rhs)
%ch0 = tfrt.new.chain
%ch1 = corert.print_tensorhandle( %result, %ch0)
tfrt.return %ch1 : !tfrt.chain
}
func @example() -> !tfrt.chain {
%ch0 = tfrt.new.chain
%cpu = corert.get_op_handler %ch0 "cpu"
// Create TensorHandles
%lhs = corert.executeop( %cpu)
"test.create_dense_tensor" () { shape = [1, 1], values = [-1.0 : f32] } : 1
%rhs = corert.executeop( %cpu)
"test.create_dense_tensor" () { shape = [1, 1], values = [-2.0 : f32] } : 1
%result = corert.executeop( %cpu) "test.add" (%lhs, %rhs) : 1
%ch1 = "corert.print_tensorhandle" (%result, %ch0) : (!corert.tensorhandle, !tfrt.chain) -> !tfrt.chain
tfrt.return %ch1 : !tfrt.chain
}

27. Device Runtime
CPU
27
//===----------------------------------------------------------------------===//
// CPU Relu kernels
//===----------------------------------------------------------------------===//
// Computes B = Relu(A).
template <typename T>
static AsyncValueRef<Chain> Relu(const DenseHostTensor& A, DenseHostTensor* B,
const ExecutionContext& exec_ctx) {
auto fn = [](auto& a, auto& b) { return a.cwiseMax(static_cast<T>(0)); };
return ::tfrt::compat::UnaryEigenKernelAsync< T, T>(A, B, std::move(fn),
exec_ctx);
}
//===----------------------------------------------------------------------===//
// CPU BiasAdd kernels
//===----------------------------------------------------------------------===//
// A special case of tf.add where bias is restricted to be 1-D.
// Currently only support NHWC data format.
template <typename T, size_t RANK>
static AsyncValueRef<Chain> BiasAdd(const DenseHostTensor& input,
const DenseHostTensor& bias,
DenseHostTensor* output,
const ExecutionContext& exec_ctx) {
DHTIndexableView<T, RANK> input_view(&input);
MutableDHTIndexableView <T, RANK> output_view(output);
DHTIndexableView<T, 1> bias_view(&bias);
const auto& shape_input = input_view.FixedShape();
const auto& shape_bias = bias_view.FixedShape();
const auto& shape_output = output_view.FixedShape();
if (shape_input != shape_output) {
return EmitErrorAsync(exec_ctx, "unexpected output shape" );
}
if (shape_bias[0] != shape_input[RANK - 1]) {
return EmitErrorAsync(exec_ctx, "bias shape does not match input shape" );
}
// Reshape bias to the shape of input. Broadcast along the last axis of input.
Eigen::array<Eigen::Index, RANK> reshape_dims;
Eigen::array<Eigen::Index, RANK> broadcast_dims;
for (size_t i = 0; i < RANK - 1; ++i) {
reshape_dims[i] = static_cast<Eigen::Index>(1);
broadcast_dims[i] = static_cast<Eigen::Index>(shape_input[i]);
}
reshape_dims[RANK - 1] = static_cast<Eigen::Index>(shape_bias[0]);
broadcast_dims[RANK - 1] = static_cast<Eigen::Index>(1);
auto input_t = AsEigenConstTensor(input_view);
auto bias_t = AsEigenConstTensor(bias_view);
auto output_t = AsEigenTensor(output_view);
auto expr = input_t + bias_t.reshape(reshape_dims).broadcast(broadcast_dims);
return AsyncAssign(
exec_ctx.host()->GetOrCreateSharedContext<EigenHostContext>(),
std::move(output_t), std::move(expr),
KeepBuffers::alive(&input, &bias, output));
}
https://github.com/tensorflow/runtime/blob/master/backends/cpu/lib/kernels/cpu_kernels.h

28. Dialects we can see now
•tfrt: we know what this is for
•tfrt_test: to test tfrt
•tfrt_data: tf.data, to deal with input pipeline
•tfrt_dht: dense host tensor
•corert: Core Runtime, eager execution
•ts: tensor shape
•coo: COOrdinate list sparse tensor
•eigen: wrapper around the eigen library
•btf: binary tensor format
•cuda: you know what cuda means :-)
28

29. Concluding Remarks
•MLIR related talks and publications, https://mlir.llvm.org/talks/
•We scratched the surface of TFRT host runtime and core runtime. There are more details
•threading model: thread pool / work queue,
•memory allocation: tcmalloc for server, other small allocators for embedded systems,
•non-strict execution, and
•registers: BEF executor is a register machine
•we didn’t touch other important components such as device runtimes, eps. the GPU
part, and distributed environment
29

30. Fin
30

31. Device Runtime Design Principles
•A thin wrapper of low-level (driver) APIs, exposing device capabilities to graph compiler
•Memory Allocation
•Async host <-> device transfer, and kernel execution
•Dependency management
•Focus on mechanism instead of policy
•E.g. No built-in special-purpose streams for GPU support:
•For pure eager execution, can default to one stream for everything
•For tf.function execution, compiler can pick streams
31