2024-12-11.md

marp	theme	headingDivider
true	ginkgo	2

Custom LinOps

Early-Adopter Session 11.12.2024

Topics

• Non-Native Operators
• Custom gko::LinOp Subtypes
• Precision Dispatch

Non-Native Operators

Scenario: A user wants to solve linear systems with Ginkgo, but the system operator has no equivalent Ginkgo type.

Happens if:

• Different interpretation of data storage (e.g. storing the diagonal element of each row first in the CSR format)
• Uncommon matrix format is used (e.g. LDU format in OpenFOAM)
• Operator application is implemented matrix-free
• ...

Using Non-Native Operators in Ginkgo

1. Converting to Ginkgo Type

   Extract operator entries and store them in gko::matrix_data, or gko::device_matrix_data.

   Read entries into a Ginkgo matrix with mtx->read(data).

2. Wrapping in Custom gko::LinOp

   Create new subtype of gko::LinOp that uses non-native operator internally.

   Pass subtype to Ginkgo solvers.

Creating gko::LinOp Subtypes

class MyLinOp: public gko::EnableLinOp<MyLinOp, gko::LinOp> {
  // required constructor
  MyLinOp(std::shared_ptr<const gko::Executor> exec)
    : gko::EnableLinOp<MyLinOp, gko::LinOp>(exec) { ... }

  // required method implementations
  void apply_impl(const gko::LinOp* b, gko::LinOp* x) const override;
  void apply_impl(const gko::LinOp* alpha, const gko::LinOp* b, 
                  const gko::LinOp* beta, gko::LinOp* x) const override;
};

gko::EnableLinOp<MyLinOp, gko::LinOp> can be shortened to gko::EnableLinOp<MyLinOp>, since gko::LinOp is the default base type.

It is an error to derive from gko::LinOp directly.

Creating gko::LinOp Subtypes

class MyLinOp: public gko::EnableLinOp<MyLinOp, gko::LinOp> {
  // required constructor
  MyLinOp(std::shared_ptr<const gko::Executor> exec)
    : gko::EnableLinOp<MyLinOp, gko::LinOp>(exec) { ... }

  // Passing the LinOp size to the base class
  MyLinOp(std::shared_ptr<const gko::Executor> exec, ...)
    : gko::EnableLinOp<MyLinOp, gko::LinOp>(exec, gko::dim<2>(...)) { ... }

  ...
};

Why is EnableLinOp the Base Class?

The gko::LinOp class itself is derived from gko::PolymorphicObject. Thus, it is necessary to implement more pure virtual functions than the apply_impl.

The gko::EnableLinOp provides sensible default implementations for those virtual functions.

Note: Most implementations will rely on the copy/move assignments, so custom classes without those will fail to compile.

Handling Constructors

Polymorphic Ginkgo types don't offer public constructors and instead require to use the static create method.

Custom subtypes may follow this, or just provide a public constructor.

Approach Using create

class MyLinOp: public gko::EnableLinOp<MyLinOp> {
  friend gko::EnablePolymorphicObject<MyLinOp, gko::LinOp>;

public:
  static std::unique_ptr<MyLinOp> create(...); 

private:
  MyLinOp(std::shared_ptr<const gko::Executor> exec);
  // more constructors
};

auto obj = MyLinOp::create(args...);

friend declaration is required so that EnablePolymorphicObject can access the constructor.

Approach Using Public Constructors

class MyLinOp: public gko::EnableLinOp<MyLinOp> {
public:
  MyLinOp(std::shared_ptr<const gko::Executor> exec);
  // more constructors
};

auto obj = std::make_unique<MyLinOp>(args...);

Implementing a Custom Apply

void MyLinop::apply_impl(const gko::LinOp* b, gko::LinOp* x) const{
  // get gko::matrix::Dense from gko::LinOp inputs
  auto dense_b = gko::as<gko::matrix::Dense<ValueType>>(b);
  auto dense_x = gko::as<gko::matrix::Dense<ValueType>>(x);

  my_operator(extra_data, dense_b->get_const_values(), x->get_values());
}

Note: gko::as<...> is a more flexible version of dynamic_cast<...>, but throws on failure.

Precision Dispatch

Ginkgo provides utilities to handle the dispatch to the runtime gko::matrix::Dense<T> more easily. The gko::precision_dispatch<T> function will either

• cast the inputs to gko::matrix::Dense<T> if that is their runtime type
• convert the inputs to gko::matrix::Dense<T> if that conversion exists
• throw an error

Available conversions:

• Dense<T1> <--> Dense<T2>
• Dense<std::complex<T1>> <--> Dense<std::complex<T2>>

with T1, T2 in {float, double, gko::half}

Using Precision Dispatch

#include <ginkgo/core/base/precision_dispatch.hpp>

void MyLinop::apply_impl(const gko::LinOp* b, gko::LinOp* x) const{
  // let ginkgo handle the casting to the runtime type
  gko::precision_dispatch<ValueType>(
    [this](auto dense_b, auto dense_x){
      my_operator(extra_data, dense_b->get_const_values(), x->get_values());
    },
    b, x
  );
}

Flavors of Precision Dispatch

The precision dispatch is also available as:

• gko::mixed_precision_dispatch: allows different precision types for b and x without conversion
• gko::[mixed_]precision_dispatch_real_complex: allows converting complex vectors to real vectors (with 2x columns)
• gko::experimental::precision_dispatch_real_complex_distributed: same as above but with distributed vectors

Key Take Aways

• Derive from gko::EnableLinOp<MyLinOp>, not from gko::LinOp
• The constructor MyLinOp(std::shared_ptr<const Executor>) has to be available to the class gko::EnablePolymorphicObject<MyLinOp, gko::LinOp>
• Use precision dispatch to handle runtime types of input/output vectors of apply_impl