5.11. Custom transforms

The concept of a transform comes from the PopART framework and is a means of graph optimisation. Unlike patterns which only match ops in the PopART IR, transforms carry out optimisations of the entire graph in the PopART IR, and therefore change the PopART IR in a more complex way.

For example, the role of the SubgraphOutline transform is to extract duplicate Op structures into new graphs and use CallOps calls to save memory.

This chapter helps users understand how to implement custom PopART transforms and use them in PopRT.

Before reading this chapter, it is necessary to be familiar with the following:

• PopART framework.

• Infrastructure of the PopART IR

5.11.1. Implementing custom PopART transforms

To create a custom transform in PopART, at least one C++ file needs to be written.

This example is used to print the current serialised IR, which can be compiled into a standalone dynamic library and linked when using PopRT:

Listing 5.16 ir_serialise_transform.cpp

// Copyright (c) 2022 Graphcore Ltd. All rights reserved.

#include <iostream>
#include <string>
#include <popart/graph.hpp>
#include <popart/ir.hpp>
#include <popart/op.hpp>
#include <popart/transforms/transform.hpp>

namespace popart {
class Graph;

class IrSerialise : public Transform {
public:
  static std::size_t id();

  IrSerialise() : Transform() {}
  virtual ~IrSerialise() override {}

  virtual bool apply(Graph &graph) const final;

  virtual std::size_t getId() const final { return id(); }

  virtual std::string getName() const final { return "IrSerialise"; }
};

std::size_t IrSerialise::id() { return typeid(IrSerialise).hash_code(); }

bool IrSerialise::apply(Graph &graph) const {
  const auto &ir = graph.getIr();
  std::stringstream ss;
  ir.serialise(Ir::SerialiseFormat::JSON, ss);
  const auto modelStr = ss.str();
  std::cout << "SerializedIr : " << std::endl;
  std::cout << modelStr << std::endl;
  return true;
}

namespace {
bool init = Transform::registerTransform(new IrSerialise);
}

} // namespace popart

In order to implement a custom PopART transform, it is necessary to inherit and override popart::Transform or implement the main methods:

• apply() implements IR conversion and other functions

• getId() is the unique transform ID

• getName() defines the name of the custom transform. It is necessary to avoid conflicts with existing transform names in PopART.

• registerTransform() registers custom transform with PopART

Please refer to the default transform in PopART contained in the public PopART repo on GitHub.

5.11.2. Using custom transforms in PopRT

After the custom transform has been written, the next thing to do is to use it in PopRT.

The source code of the custom transform needs to be compiled into a standalone dynamic link library and linked when using PopRT. The following is an example command for compilation:

g++ \
    -std=c++14 \
    -fPIC \
    -O3 \
    -DONNX_NAMESPACE=onnx \
    ir_serialise_transform.cpp \
    -shared \
    -lpopart \
    -o custom_transform.so

Then, the dynamic library containing the custom transform can be linked with the --custom_library_so_paths PopRT CLI parameter:

poprt  --custom_library_so_paths path/to/shared/library

Since transforms change PopART IR in a more complex way, they need to be executed in a predefined order. Usually, before writing the transform, it is necessary to consider which execution position to put it in.

The execution order of the transform is divided into several stages, and PopART allows a user-defined custom transform to be called at the checkpoint after each stage is completed.

Predefined checkpoints are as follows:

• Fwd0: Initial IR after ONNX Lowering to PopART IR

• Fwd1: After pre-alias patterns are applied to FWD0

• Bwd0: After the backward pass

• Prealias: After pre-alias patterns are applied to BWD0

• MainLoops: After applying the MainLoops transform

• Final: After the final IR of all transforms is applied

See the C++ code in the PopART public GitHub repo for more details

Configure the custom transform with the --compiler_options parameter of the PopRT CLI, as:

poprt \
    --input_model model.onnx \
    --output_model model_export.onnx \
    --export_popef \
    --output_dir model \
    --custom_library_so_paths build/custom_transforms.so \
    --compiler_options custom_transform_applier_settings="{'Fwd0': ['IrSerialise']}"