Overview

As Uchen.ML is heading towards the public announcement and first demos, some low-hanging fruit needs to be picked in terms of optimizations. The most often used piece of any ML library is the linear layer as it is the most basic building block for any neural net. This post details the process of optimizing the code.

Requirements

Uchen.ML is designed for implementing ML solutions that can be easily integrated into existing systems, with specific goals on Web Assembly, embedded and video games.

To maintain velocity and to avoid overcomplicating build and validation process, following constraints are in place:

• Only the C++20 standard library is used. ABSL dependency is there (logging, asserts and some utilities) but it is under consideration if its inclusion will remain mandatory.
• No compiler-specific optimizations, including pragmas, conditional compilations or intrinsics.
• No CPU architecture-specific optimizations. Particularly, no optimizations for one architecture that may be detrimental for others. Apple M2 and Intel Core CPUs are used to inform and direct the optimization efforts.
• Uchen.ML is and will remain a CPU-only ML framework. There are no plans at this point to implement GPU or other acceleration support.

This constraints will be lifted as the deployment targets and actual requirements are better understood.

Benchmark code

The benchmark runs inference through the linear layers of different configurations. Inputs are initialized to 0, parameters are initialized to random values outside the benchmark loop. Range of parameter values is between -1 and 1. Output values are not checked. float datatype is used for inputs and outputs

template <size_t Is, size_t Os, typename D = float>
static void BM_Linear(benchmark::State& state) {
  // Linear layer with Is inputs and Os outputs
  uchen::Linear<uchen::Vector<D, Is>, Os> layer;
  // zero-initialized input vector. This operation is O(n) to the number
  // of the inputs and may have a negligible impact on benchmark.
  uchen::Vector<D, Is> input;
  // Parameters are using the store filled with random values outside
  // the loop This operation is O(1) to the number of parameters and has
  // no impact on benchmark.
  uchen::parameters_t<decltype(layer)> parameters(params->store());
  for (auto _ : state) {
    benchmark::DoNotOptimize(layer(input, parameters));
  }
}

Hardware

Regular PCs are used. Note that the numbers can not be compared across the architectures and this paper is only concerned with the relative gains and not the absolute values.

Apple M2 Pro
  10 Cores
  L1 Data 64 KiB
  L1 Instruction 128 KiB
  L2 Unified 4096 KiB (x10)

Intel(R) Core(TM) i7-10700KF CPU @ 3.80GHz   3.79 GHz
  L1 Data 32 KiB (x8)
  L1 Instruction 32 KiB (x8)
  L2 Unified 256 KiB (x8)
  L3 Unified 16384 KiB (x1)

Naive version

Linear layer runs the following operation to produce the output: $y_j=b_j+\sum _{i=\; 0}^n{w}_{ji}{x}_i$

Which translated into the following C++ code:

  output_t operator()(const input_t& inputs,
                      const Parameters<parameter_count>& parameters) const {
    output_t outputs;
    constexpr size_t Is = input_t::elements;
    for (size_t output = 0; output < Outputs; ++output) {
      outputs[output] = parameters[output * (Is + 1) + Is];
      for (size_t input = 0; input < Is; ++input) {
        outputs[output] +=
            inputs[input] * parameters[output * (Is + 1) + input];
      }
    }
    return outputs;
  }

Given n inputs and m outputs, parameters layout is: Parameters are a flat array in the following format (n is a number of inputs, m is the number of outputs): $\{w_{00},\; \dots \; ,w_{0n},b_0,w_{10},\; \dots \; ,w_{mn},b_m\}$

Benchmark Results:

Note that the number of operations (memory reads, stores and arithmetic) is directly correlated to the number of parameters so all the models were set up to have roughly the same number of them.

Intel architecture shows approximately 3.4x spread between the best case scenario (number outputs drastically exceeds number of inputs) and worst case scenario (number of inputs is much greater then the number of outputs). Apple ARM implementation showed 7.5x spread.

Transposed iteration order

The first optimization is to change the iteration order. Instead of iterating over the outputs and then over the inputs, we iterate over the inputs and then over the outputs. This change allows for better cache utilization and reduces the number of cache misses.

  output_t operator()(
      const input_t& inputs,
      const Parameters<(Input::elements + 1) * Outputs>& parameters) const {
    auto it = parameters.begin();
    output_t outputs;
    for (size_t i = 0; i < outputs.size(); i++) {
      outputs[i] = *it++;
    }
    for (auto input : inputs) {
      for (auto& output : outputs) {
        output += (*it++) * input;
      }
    }
    return outputs;
  }

Note that the code uses a linear iteration over the parameters, making the memory access pattern more predictable and cache-friendly.

Parameters layout is as follows (n is a number of inputs, m is the number of outputs): $\{b_0,\; \dots \; ,b_{\mathrm{m}},w_{00},\; \dots \; ,w_{0n},w_{10},\; \dots \; ,w_{m0}\}$

Benchmark Results:

Number of inputs has a significant negative impact on the performance as it adds an extra memory load.

Compile for the specific CPU

By default, the compiler generates the code that works on most CPUs. This precludes usage of some newer instructions that have a significant impact on the performance. To test this, I pass -march=native to the compiler which makes it target my current CPU. I am using Bazel, so the invocation looked like this (-g is for including debugging symbols as I use gdb to look at the assembly code):

bazel run -c opt --cxxopt="-g" --cxxopt="-march=native" //benchmark:linear

Benchmark Results (Intel only):

This “optimization” shows pretty solid across the board. Ultimately, it will be up to embedders to decide the CPU target to compile for.

A bigger layer

With the numbers in single digit microseconds, it seems reasonable to increase the size of the linear layer to see what impact it has on the performance.

Benchmark Results (Intel only):

The worst case is “only” 2x slower than the best case.

Using SIMD intrinsics directly

As mentioned above, SIMD intrinsics are not considered (i.e. Web Assembly still has no support for them). However, it is still interesting to see what the performance gains could be if we tried them. I did not use AVX as data alignment is not supported yet, though this will change soon.

  output_t operator()(
      const input_t& inputs,
      const Parameters<(Input::elements + 1) * Outputs>& parameters) const {
    auto it = parameters.begin();
    output_t outputs;
    for (size_t i = 0; i < outputs.size(); i++) {
      outputs[i] = (*it++);
    }
    for (size_t i = 0; i < Input::elements; i += 4) {
      __m128 input = _mm_loadu_ps(&(*inputs.begin()) + i);
      for (size_t j = 0; j < Outputs; ++j) {
        // Parameters are accessed in the wrong order - this is a bug!
        // This code is for benchmark only.
        __m128 parameters = _mm_load_ps(&(*it));
        __m128 product = _mm_dp_ps(input, parameters, 0xf1);
        it += 4;
        outputs[j] += _mm_cvtss_f32(product);
      }
    }
    return outputs;
  }
};

Benchmark Results (Intel only):

“Good” cases worsened, which shows that the updated code is ran. Shape of the data is known at compile time when using Uchen.ML so the compilers make informed decisions about vectorization and other optimizations and it looks like competing with them is an unnecessary exercise.

Manually unrolling the loop yields similar results:

  output_t operator()(
      const input_t& inputs,
      const Parameters<(Input::elements + 1) * Outputs>& parameters) const {
    auto it = parameters.begin();
    output_t outputs;
    for (size_t i = 0; i < outputs.size(); i++) {
      outputs[i] = (*it++);
    }
    for (auto input : inputs) {
      for (size_t i = 0; i < Outputs; i++) {
        outputs[i++] += (*it++) * input;
        outputs[i++] += (*it++) * input;
        outputs[i++] += (*it++) * input;
        outputs[i] += (*it++) * input;
      }
    }
    return outputs;
  }

Conclusion

At this point it looks like the performance on this granularity is pretty close to the practical limit for a single core. Next article will detail multi-threading and the memory alignment (particularly, dealing with the number of parameters not divisible by 8).

The backpropagation optimizations and Uchen’s approach to the memory management will also be detailed in the future articles.