Working with SYCL for a portable and performant library adding reactions to particle codes used in plasma physics.

The UK Atomic Energy Authority’s mission is to lead the delivery of sustainable fusion energy and maximise scientific and economic benefit Find out more www.gov.uk/ukaea This work has been funded by the EPSRC Energy Programme [grant number EP/W006839/1] VANTAGEReactions GitHub Repo NESO-Particles GitHub Repo Working with SYCL for a portable and performant library adding reactions to particle codes used in plasma physics. Sanket Gadgil 1Stefan Mijin 1Will Saunders 1 1UK Atomic Energy Authority, Culham Campus, Abingdon, Oxfordshire, UK, OX14 3DB Introduction In magnetically confined fusion plasmas, such as those found in tokamaks, the atomic and molecular reactions found in the exhaust are important in understanding a variety of phenomena. For example, interactions between the plasma and the inner wall of the reactor. As current codes for handling this type of investigation are fairly old and lack features such as robust performance scaling and the ability to run on a GPU, we decided to write a modern performance portable C++ code (VANTAGE-Reactions) that improves on the current state-of-the-art (EIRENE[2], DEGAS2[3]). We decided to build on NESO-Particles, which is a performance portable framework for building particle codes on meshes (both structured and unstructured) developed at UKAEA(UK Atomic Energy Authority). NESO-Particles is built on SYCL[4] to provide said performance portability but whilst SYCL offers some key advantages, it also presents some challenges that must be overcome. What is SYCL? In the past, when writing code for different devices like CPUs, GPUs or FPGAs, it was necessary to write a separate optimized source code for each variant of the software you wanted to run. This required writing for a separate API and potentially running into divergent development paths. SYCL aims to resolve this for C++ by allowing developers to use a single API to write one version of the source code and compile it for different devices, which results in distinct binaries. In theory, it can go even further and allow for more expansive heterogeneous computing with multiple non-CPU devices used in the same source code (ie. different sections of the code being run on different devices). Note that, some device-specific optimization can be necessary– SYCL should give executable code but it doesn’t know about specific hardware variations. However, for our purposes, we focussed on designing around distinct parallel sections in the code that can be dispatched to a single external device based on the compilation (specifically we tested on NVidia A100 GPUs). Here’s a basic for loop as written for use with the SYCL API: auto defaultQueue = sycl::queue{}; // Wrappers for std::array<int, data_size> objects auto buf0 = sycl::buffer{array_0.data(), sycl::range{data_size}}; auto buf1 = sycl::buffer{array_1.data(), sycl::range{data_size}}; auto buf2 = sycl::buffer{array_2.data(), sycl::range{data_size}}; // Submission to the SYCL queue. defaultQueue.submit( [&](sycl::handler& cgh) { // Accessors that allow SYCL to decide how to access the data // on the device. auto acc0 = sycl::accessor{buf0, cgh, sycl::read_only}; auto acc1 = sycl::accessor{buf1, cgh, sycl::read_only}; auto acc2 = sycl::accessor{buf2, cgh, sycl::write_only}; // The main parallel_for loop. cgh.parallel_for<>( data_size, // number of iterations to loop over [=] (auto idx) { acc2[idx] = acc0[idx] + acc1[idx]; } ); } ); Challenges for VANTAGE-Reactions For the purposes of our code, we wanted to utilise dynamic polymorphism in C++ to allow us to reuse a great deal of functionality when creating objects to store data and functions related to the reactions we want to apply to particles. This presents a challenge when working with SYCL as it does not have the functionality to handle virtual functions or virtual inheritance. Additionally, there are also restrictions on what variable types are deemed to be device-copyable. The actual SYCL functionality (ie. loops, buffers, etc.) is wrapped by NESO-Particles to provide a cleaner interface for us but the limitations of SYCL still remain and we still have to be careful. Any trivially copyable types either built-in to C++ or user-defined are acceptable, so a method to allow for polymorphism via static means (using templates) could be useful. Potential solution for polymorphism via CRTP CRTP stands for Curiously Recurring Template Pattern, and is a viable alternative to dynamic polymorphism (with some limits). The idea is to construct derived class by inheriting from a template class and then to have the derived class as a template parameter for the base class. Below are two examples of inheritance in C++. On the left is the CRTP variant and on the right is a virtual member function overload based variant. Both are functionally identical but only the CRTP version is trivially copyable and is therefore device-copyable according to SYCL. template <typename Derived> class Animal_CRTP { public: void sound() { auto& underlying = static_cast<Derived&>(*this);,→ return underlying.template sound(); } }; class Cat_CRTP : public Animal_CRTP<Cat_CRTP> {,→ public: void sound() { std::cout << "Cats meow!" << std::endl; } }; class Animal_Virt { public: virtual void sound() {} }; class Cat_Virt : public Animal_Virt { public: void sound() override { std::cout << "Cats meow!" << std::endl; } }; Specifically, an object of type Cat_CRTP can be wrapped in a SYCL buffer and accessed within a SYCL loop that would be run on an external device, whereas Cat_Virt cannot. The key point is that any function that needs an Animal_CRTP derived object, could simply use a template parameter for the function like T_Animal and just specify the type of the argument in the function signature as: Animal_CRTP<T_Animal> and when called the function would accept a Cat_CRTP object. Host-Device separation for VANTAGE-Reactions An alternative solution for the specific needs of VANTAGE-Reactions was to separate out host-compatible and device-compatible objects and enforce a strict methodology for defining the device-compatible objects. Chiefly, the aim was to ensure that all device-compatible objects are device-copyable so that they are able to be accessed inside a SYCL loop and are the primary point of interaction with a SYCL loop. An example of the layout of the framework is given below (with placeholder names): HostBaseClass DeviceBaseClass HostDerivedClass_0 - device_obj: DeviceDerivedClass_0 + get_device_obj(): DeviceDerivedClass_0 HostDerivedClass_1 - device_obj: DeviceDerivedClass_1 + get_device_obj(): DeviceDerivedClass_1 HostDerivedClass_2 - device_obj: DeviceDerivedClass_2 + get_device_obj(): DeviceDerivedClass_2 DeviceDerivedClass_0 DeviceDerivedClass_1 DeviceDerivedClass_2 The general rule is that any HostDerivedClass type is mostly free to use non-SYCL compliant features (chiefly, virtual functions) but any DeviceDerivedClass type must be fully device-copyable and use only SYCL compliant features. Whenever there’s a need to run a SYCL loop with data from a DeviceDerivedClass, a copy of the object is retrieved from the HostDerivedClass object, stored in a temporary variable and then passed by value to the SYCL loop. The DeviceDerivedClass types are not implemented via CRTP due to some unique templating limitations in our code, they instead use non-virtual function overloads, since they are not ever passed to polymorphic functions. Acknowledgments & References Work on VANTAGE-Reactions was done at UKAEA (UK Atomic Energy Authority) with funding from UKRI (UK Research and Innovation). [1] SYCL-CRTP details: https://codeplay.com/portal/blogs/2019/07/12/ enabling-polymorphism-in-sycl-using-the-cpp-idiom-crtp [2] EIRENE homepage: https://www.eirene.de/ [3] DEGAS2 homepage: https://w3.pppl.gov/degas2/ [4] SYCL 2020 specification: https://registry.khronos.org/SYCL/specs/ sycl-2020/html/sycl-2020.html