Native operations

Motivation:

• Speed up some important common calculations, and potentially reduce memory requirements. Examples:

  • LSTM

  • CTC loss

• Pure TensorFlow implementations can be suboptimal

  • TF ops almost always create copies, even SplitOp etc

    • Not a memory problem, as input tensor will get freed if not used further

    • Performance problem

  • Gradient might be suboptimal

    • Require too much memory (see automatic gradient checkpointing for a solution)

    • No automatic optimization

    • (Could be solved by custom TF gradient)

  • Memory can be too much distributed (tf.TensorArray, TF Stack)

    • Esp. problematic in loop: Separate tensor for every iteration

    • Much better to allocate it as consecutive / contiguous block

  • Overhead (calling individual TF ops, etc) (minor compared to the other points) (XLA can partially also solve this)

Solution: Write native (C++/CUDA) code

Why is native code faster?

• Operate inplace on tensors

  • Solves all problems mentioned, no unnecessary copies

  • Can use consecutive tensor / memory

• Enforces custom gradient implementation

Problems with native code:

• Can be difficult, memory unsafe, needs more debugging

• Need multiple implementations: CPU (C++), GPU (CUDA)

Our Approach in RETURNN:

The NativeOp framework. See returnn.native_op, returnn.tf.native_op, returnn.theano.native_op.

• Some wrapper / helper code to simplify writing custom native op

• Abstractions to allow single code for CPU & GPU

  • Write kernel CUDA style, using threadIdx, blockIdx, etc

    • Kernel code must be flexible for amount of threads

    • Example, LSTM kernel, loop over dimensions, executed per time-frame:

      int idx = threadIdx.x + blockDim.x ∗ blockIdx.x;
      while (idx < n_cells ∗ n_batch) {
          int batch_idx = idx / n_cells;
          int cell_idx = idx % n_cells;
          ...
          idx += gridDim.x ∗ blockDim.x;
      }
      

  • On CPU

    • Custom gridDim, blockDim

    • Other CUDA-like wrappers

History:

• Already available for the Theano backend

• Ported to TensorFlow

  • Directly support for all already prev. implemented ops (LSTM, Baum Welch aligner, …)

• Easy to port to other frameworks

Examples:

• NativeLstm (LstmGenericBase)

• NativeLstm2

• TwoDLSTM

• FastBaumWelch

• FastViterbi

• OptimalCompletionEditDistance

• EditDistance

• Chunking, UnChunking

See also TensorFlow LSTM Benchmark.