char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Ian Faustsycl: fix multiple warnings in compiling sycl backend (#26713)6602dd338

char/aliceai-t5-moe
155.8 KiB3782 linesraw
1//
2// MIT license
3// Copyright (C) 2024 Intel Corporation
4// SPDX-License-Identifier: MIT
5//
6
7//
8// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
9// See https://llvm.org/LICENSE.txt for license information.
10// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
11//
12
13#ifndef GGML_SYCL_DPCT_HELPER_HPP
14#define GGML_SYCL_DPCT_HELPER_HPP
15
16#include <cstdlib>
17#include <iostream>
18#include <map>
19
20#include <sycl/sycl.hpp>
21#include <sycl/half_type.hpp>
22#include <oneapi/mkl.hpp>
23
24#if defined(__linux__)
25#include <sys/mman.h>
26#elif defined(_WIN64)
27#ifndef NOMINMAX
28#define NOMINMAX
29#endif
30#include <windows.h>
31#else
32#error "Only support Windows and Linux."
33#endif
34
35#if defined(__linux__)
36#include <unistd.h>
37#include <sys/syscall.h>
38#endif
39#if defined(_WIN64)
40#ifndef NOMINMAX
41#define NOMINMAX
42#endif
43#include <windows.h>
44#endif
45
46
47#define DPCT_COMPATIBILITY_TEMP (900)
48
49#if defined(_MSC_VER)
50#define __dpct_align__(n) __declspec(align(n))
51#define __dpct_inline__ __forceinline
52#else
53#define __dpct_align__(n) __attribute__((aligned(n)))
54#define __dpct_inline__ __inline__ __attribute__((always_inline))
55#endif
56
57#if defined(_MSC_VER)
58#define __dpct_noinline__ __declspec(noinline)
59#else
60#define __dpct_noinline__ __attribute__((noinline))
61#endif
62
63#define DPCT_UNUSED(x) (void)(x)
64
65[[noreturn]] inline void _abort(const char * str) {
66    std::cerr << str << std::endl;
67    std::abort();
68}
69
70inline std::string get_device_type_name(const sycl::device &Device) {
71    auto DeviceType = Device.get_info<sycl::info::device::device_type>();
72    switch (DeviceType) {
73    case sycl::info::device_type::cpu:
74        return "cpu";
75    case sycl::info::device_type::gpu:
76        return "gpu";
77    case sycl::info::device_type::host:
78        return "host";
79    case sycl::info::device_type::accelerator:
80        return "acc";
81    default:
82        return "unknown";
83    }
84}
85
86inline std::string get_device_backend_and_type(const sycl::device &device) {
87    std::stringstream device_type;
88    sycl::backend backend = device.get_backend();
89    device_type <<  backend << ":" << get_device_type_name(device);
90    return device_type.str();
91}
92
93template <typename Ts> struct matrix_info_t {
94    oneapi::mkl::transpose transpose_info[2];
95    Ts                     value_info[2];
96    std::int64_t           size_info[3];
97    std::int64_t           ld_info[3];
98    std::int64_t           groupsize_info;
99};
100
101namespace dpct
102{
103    typedef sycl::queue *queue_ptr;
104    typedef sycl::event *event_ptr;
105    typedef char *device_ptr;
106    typedef uint8_t byte_t;
107    typedef sycl::buffer<byte_t> buffer_t;
108
109    /// SYCL default exception handler
110    inline auto exception_handler = [](sycl::exception_list exceptions)
111    {
112        for (std::exception_ptr const &e : exceptions)
113        {
114            try
115            {
116                std::rethrow_exception(e);
117            }
118            catch (sycl::exception const &e)
119            {
120                std::cerr << "Caught asynchronous SYCL exception:" << std::endl
121                          << e.what() << std::endl
122                          << "Exception caught at file:" << __FILE__
123                          << ", line:" << __LINE__ << std::endl;
124            }
125        }
126    };
127
128    enum error_code
129    {
130        success = 0,
131        default_error = 999
132    };
133
134    enum memcpy_direction
135    {
136        host_to_host,
137        host_to_device,
138        device_to_host,
139        device_to_device,
140        automatic
141    };
142
143    enum memory_region
144    {
145        global = 0, // device global memory
146        constant,   // device constant memory
147        local,      // device local memory
148        shared,     // memory which can be accessed by host and device
149    };
150
151    enum class library_data_t : unsigned char
152    {
153        real_float = 0,
154        complex_float,
155        real_double,
156        complex_double,
157        real_half,
158        complex_half,
159        real_bfloat16,
160        complex_bfloat16,
161        real_int4,
162        complex_int4,
163        real_uint4,
164        complex_uint4,
165        real_int8,
166        complex_int8,
167        real_uint8,
168        complex_uint8,
169        real_int16,
170        complex_int16,
171        real_uint16,
172        complex_uint16,
173        real_int32,
174        complex_int32,
175        real_uint32,
176        complex_uint32,
177        real_int64,
178        complex_int64,
179        real_uint64,
180        complex_uint64,
181        real_int8_4,
182        real_int8_32,
183        real_uint8_4,
184        library_data_t_size
185    };
186
187    template <typename T>
188    struct DataType
189    {
190        using T2 = T;
191    };
192    template <typename T>
193    struct DataType<sycl::vec<T, 2>>
194    {
195        using T2 = std::complex<T>;
196    };
197
198    static void destroy_event(event_ptr event)
199    {
200        delete event;
201    }
202
203    static inline unsigned int get_tid()
204    {
205#if defined(__linux__)
206        return syscall(SYS_gettid);
207#elif defined(_WIN64)
208        return GetCurrentThreadId();
209#else
210#error "Only support Windows and Linux."
211#endif
212    }
213
214    namespace detail
215    {
216        static void get_version(const sycl::device &dev, int &major, int &minor)
217        {
218            // Version string has the following format:
219            // a. OpenCL<space><major.minor><space><vendor-specific-information>
220            // b. <major.minor>
221            // c. <AmdGcnArchName> e.g gfx1030
222            std::string ver;
223            ver = dev.get_info<sycl::info::device::version>();
224            std::string::size_type i = 0;
225            while (i < ver.size()) {
226              if (isdigit(ver[i]))
227                break;
228              i++;
229            }
230            major = std::stoi(&(ver[i]));
231            while (i < ver.size()) {
232              if (ver[i] == '.')
233                break;
234              i++;
235            }
236            if (i < ver.size()) {
237              // a. and b.
238              i++;
239              minor = std::stoi(&(ver[i]));
240            } else {
241              // c.
242              minor = 0;
243            }
244        }
245
246        template <typename tag, typename T>
247        class generic_error_type
248        {
249        public:
250            generic_error_type() = default;
251            generic_error_type(T value) : value{value} {}
252            operator T() const { return value; }
253
254        private:
255            T value;
256        };
257
258    } // namespace detail
259
260    // COPY from DPCT head files
261    /// dim3 is used to store 3 component dimensions.
262    class dim3 {
263        public:
264        unsigned x, y, z;
265
266        constexpr dim3(unsigned x = 1, unsigned y = 1, unsigned z = 1)
267            : x(x), y(y), z(z) {}
268
269        dim3(const sycl::id<3> &r) : dim3(r[2], r[1], r[0]) {}
270
271        operator sycl::range<3>() const { return sycl::range<3>(z, y, x); }
272    }; // namespace dim3
273
274    inline dim3 operator*(const dim3 &a, const dim3 &b) {
275    return dim3{a.x * b.x, a.y * b.y, a.z * b.z};
276    }
277    // COPY from DPCT head files
278
279
280    /// Pitched 2D/3D memory data.
281    class pitched_data
282    {
283    public:
284        pitched_data() : pitched_data(nullptr, 0, 0, 0) {}
285        pitched_data(void *data, size_t pitch, size_t x, size_t y)
286            : _data(data), _pitch(pitch), _x(x), _y(y) {}
287
288        void *get_data_ptr() { return _data; }
289        void set_data_ptr(void *data) { _data = data; }
290
291        size_t get_pitch() { return _pitch; }
292        void set_pitch(size_t pitch) { _pitch = pitch; }
293
294        size_t get_x() { return _x; }
295        void set_x(size_t x) { _x = x; }
296
297        size_t get_y() { return _y; }
298        void set_y(size_t y) { _y = y; }
299
300    private:
301        void *_data;
302        size_t _pitch, _x, _y;
303    };
304
305    class device_info
306    {
307    public:
308        // get interface
309        const char *get_name() const { return _name; }
310        char *get_name() { return _name; }
311        template <typename WorkItemSizesTy = sycl::range<3>,
312                  std::enable_if_t<std::is_same_v<WorkItemSizesTy, sycl::range<3>> ||
313                                       std::is_same_v<WorkItemSizesTy, int *>,
314                                   int> = 0>
315        auto get_max_work_item_sizes() const
316        {
317            if constexpr (std::is_same_v<WorkItemSizesTy, sycl::range<3>>)
318                return sycl::range<3>(_max_work_item_sizes_i[0],
319                                      _max_work_item_sizes_i[1],
320                                      _max_work_item_sizes_i[2]);
321            else
322            {
323                return _max_work_item_sizes_i;
324            }
325        }
326        template <typename WorkItemSizesTy = sycl::range<3>,
327                  std::enable_if_t<std::is_same_v<WorkItemSizesTy, sycl::range<3>> ||
328                                       std::is_same_v<WorkItemSizesTy, int *>,
329                                   int> = 0>
330        auto get_max_work_item_sizes()
331        {
332            if constexpr (std::is_same_v<WorkItemSizesTy, sycl::range<3>>)
333                return sycl::range<3>(_max_work_item_sizes_i[0],
334                                      _max_work_item_sizes_i[1],
335                                      _max_work_item_sizes_i[2]);
336            else
337            {
338                return _max_work_item_sizes_i;
339            }
340        }
341        bool get_host_unified_memory() const { return _host_unified_memory; }
342        int get_major_version() const { return _major; }
343        int get_minor_version() const { return _minor; }
344        int get_integrated() const { return _integrated; }
345        int get_max_clock_frequency() const { return _frequency; }
346        int get_max_compute_units() const { return _max_compute_units; }
347        int get_max_work_group_size() const { return _max_work_group_size; }
348        int get_max_sub_group_size() const { return _max_sub_group_size; }
349        int get_max_work_items_per_compute_unit() const
350        {
351            return _max_work_items_per_compute_unit;
352        }
353        int get_max_register_size_per_work_group() const
354        {
355            return _max_register_size_per_work_group;
356        }
357        template <typename NDRangeSizeTy = size_t *,
358                  std::enable_if_t<std::is_same_v<NDRangeSizeTy, size_t *> ||
359                                       std::is_same_v<NDRangeSizeTy, int *>,
360                                   int> = 0>
361        auto get_max_nd_range_size() const
362        {
363            if constexpr (std::is_same_v<NDRangeSizeTy, size_t *>)
364                return _max_nd_range_size;
365            else
366                return _max_nd_range_size_i;
367        }
368        template <typename NDRangeSizeTy = size_t *,
369                  std::enable_if_t<std::is_same_v<NDRangeSizeTy, size_t *> ||
370                                       std::is_same_v<NDRangeSizeTy, int *>,
371                                   int> = 0>
372        auto get_max_nd_range_size()
373        {
374            if constexpr (std::is_same_v<NDRangeSizeTy, size_t *>)
375                return _max_nd_range_size;
376            else
377                return _max_nd_range_size_i;
378        }
379        size_t get_global_mem_size() const { return _global_mem_size; }
380        size_t get_local_mem_size() const { return _local_mem_size; }
381        size_t get_max_mem_alloc_size() const { return _max_mem_alloc_size; }
382        /// Returns the maximum clock rate of device's global memory in kHz. If
383        /// compiler does not support this API then returns default value 3200000 kHz.
384        unsigned int get_memory_clock_rate() const { return _memory_clock_rate; }
385        /// Returns the maximum bus width between device and memory in bits. If
386        /// compiler does not support this API then returns default value 64 bits.
387        unsigned int get_memory_bus_width() const { return _memory_bus_width; }
388        uint32_t get_device_id() const { return _device_id; }
389        std::array<unsigned char, 16> get_uuid() const { return _uuid; }
390        /// Returns global memory cache size in bytes.
391        unsigned int get_global_mem_cache_size() const
392        {
393            return _global_mem_cache_size;
394        }
395
396        // set interface
397        void set_name(const char *name)
398        {
399            size_t length = strlen(name);
400            if (length < 256)
401            {
402                std::memcpy(_name, name, length + 1);
403            }
404            else
405            {
406                std::memcpy(_name, name, 255);
407                _name[255] = '\0';
408            }
409        }
410        void set_max_work_item_sizes(const sycl::range<3> max_work_item_sizes)
411        {
412            for (int i = 0; i < 3; ++i)
413                _max_work_item_sizes_i[i] = max_work_item_sizes[i];
414        }
415        [[deprecated]] void
416        set_max_work_item_sizes(const sycl::id<3> max_work_item_sizes)
417        {
418            for (int i = 0; i < 3; ++i)
419            {
420                _max_work_item_sizes_i[i] = max_work_item_sizes[i];
421            }
422        }
423        void set_host_unified_memory(bool host_unified_memory)
424        {
425            _host_unified_memory = host_unified_memory;
426        }
427        void set_major_version(int major) { _major = major; }
428        void set_minor_version(int minor) { _minor = minor; }
429        void set_integrated(int integrated) { _integrated = integrated; }
430        void set_max_clock_frequency(int frequency) { _frequency = frequency; }
431        void set_max_compute_units(int max_compute_units)
432        {
433            _max_compute_units = max_compute_units;
434        }
435        void set_global_mem_size(size_t global_mem_size)
436        {
437            _global_mem_size = global_mem_size;
438        }
439        void set_local_mem_size(size_t local_mem_size)
440        {
441            _local_mem_size = local_mem_size;
442        }
443        void set_max_mem_alloc_size(size_t max_mem_alloc_size)
444        {
445            _max_mem_alloc_size = max_mem_alloc_size;
446        }
447        void set_max_work_group_size(int max_work_group_size)
448        {
449            _max_work_group_size = max_work_group_size;
450        }
451        void set_max_sub_group_size(int max_sub_group_size)
452        {
453            _max_sub_group_size = max_sub_group_size;
454        }
455        void
456        set_max_work_items_per_compute_unit(int max_work_items_per_compute_unit)
457        {
458            _max_work_items_per_compute_unit = max_work_items_per_compute_unit;
459        }
460        void set_max_nd_range_size(int max_nd_range_size[])
461        {
462            for (int i = 0; i < 3; i++)
463            {
464                _max_nd_range_size[i] = max_nd_range_size[i];
465                _max_nd_range_size_i[i] = max_nd_range_size[i];
466            }
467        }
468        void set_memory_clock_rate(unsigned int memory_clock_rate)
469        {
470            _memory_clock_rate = memory_clock_rate;
471        }
472        void set_memory_bus_width(unsigned int memory_bus_width)
473        {
474            _memory_bus_width = memory_bus_width;
475        }
476        void
477        set_max_register_size_per_work_group(int max_register_size_per_work_group)
478        {
479            _max_register_size_per_work_group = max_register_size_per_work_group;
480        }
481        void set_device_id(uint32_t device_id)
482        {
483            _device_id = device_id;
484        }
485        void set_uuid(std::array<unsigned char, 16> uuid)
486        {
487            _uuid = std::move(uuid);
488        }
489        void set_global_mem_cache_size(unsigned int global_mem_cache_size)
490        {
491            _global_mem_cache_size = global_mem_cache_size;
492        }
493
494    private:
495        char _name[256];
496        int _max_work_item_sizes_i[3];
497        bool _host_unified_memory = false;
498        int _major;
499        int _minor;
500        int _integrated = 0;
501        int _frequency;
502        // Set estimated value 3200000 kHz as default value.
503        unsigned int _memory_clock_rate = 3200000;
504        // Set estimated value 64 bits as default value.
505        unsigned int _memory_bus_width = 64;
506        unsigned int _global_mem_cache_size;
507        int _max_compute_units;
508        int _max_work_group_size;
509        int _max_sub_group_size;
510        int _max_work_items_per_compute_unit;
511        int _max_register_size_per_work_group;
512        size_t _global_mem_size;
513        size_t _local_mem_size;
514        size_t _max_mem_alloc_size;
515        size_t _max_nd_range_size[3];
516        int _max_nd_range_size_i[3];
517        uint32_t _device_id;
518        std::array<unsigned char, 16> _uuid;
519    };
520
521    static int get_major_version(const sycl::device &dev)
522    {
523        int major, minor;
524        detail::get_version(dev, major, minor);
525        return major;
526    }
527
528    static int get_minor_version(const sycl::device &dev)
529    {
530        int major, minor;
531        detail::get_version(dev, major, minor);
532        return minor;
533    }
534
535    static void get_device_info(device_info &out, const sycl::device &dev)
536    {
537        device_info prop;
538        prop.set_name(dev.get_info<sycl::info::device::name>().c_str());
539
540        int major, minor;
541        detail::get_version(dev, major, minor);
542        prop.set_major_version(major);
543        prop.set_minor_version(minor);
544
545        prop.set_max_work_item_sizes(
546#if (__SYCL_COMPILER_VERSION && __SYCL_COMPILER_VERSION < 20220902)
547            // oneAPI DPC++ compiler older than 2022/09/02, where max_work_item_sizes
548            // is an enum class element
549            dev.get_info<sycl::info::device::max_work_item_sizes>());
550#else
551            // SYCL 2020-conformant code, max_work_item_sizes is a struct templated by
552            // an int
553            dev.get_info<sycl::info::device::max_work_item_sizes<3>>());
554#endif
555        prop.set_host_unified_memory(dev.has(sycl::aspect::usm_host_allocations));
556
557        prop.set_max_clock_frequency(
558            dev.get_info<sycl::info::device::max_clock_frequency>() * 1000);
559
560        prop.set_max_compute_units(
561            dev.get_info<sycl::info::device::max_compute_units>());
562        prop.set_max_work_group_size(
563            dev.get_info<sycl::info::device::max_work_group_size>());
564        prop.set_global_mem_size(dev.get_info<sycl::info::device::global_mem_size>());
565        prop.set_local_mem_size(dev.get_info<sycl::info::device::local_mem_size>());
566        prop.set_max_mem_alloc_size(dev.get_info<sycl::info::device::max_mem_alloc_size>());
567
568#if (defined(SYCL_EXT_INTEL_DEVICE_INFO) && SYCL_EXT_INTEL_DEVICE_INFO >= 6)
569        if (dev.has(sycl::aspect::ext_intel_memory_clock_rate))
570        {
571            unsigned int tmp =
572                dev.get_info<sycl::ext::intel::info::device::memory_clock_rate>();
573            if (tmp != 0)
574                prop.set_memory_clock_rate(1000 * tmp);
575        }
576        if (dev.has(sycl::aspect::ext_intel_memory_bus_width))
577        {
578            prop.set_memory_bus_width(
579                dev.get_info<sycl::ext::intel::info::device::memory_bus_width>());
580        }
581        if (dev.has(sycl::aspect::ext_intel_device_id))
582        {
583            prop.set_device_id(
584                dev.get_info<sycl::ext::intel::info::device::device_id>());
585        }
586        if (dev.has(sycl::aspect::ext_intel_device_info_uuid))
587        {
588            prop.set_uuid(dev.get_info<sycl::ext::intel::info::device::uuid>());
589        }
590#elif defined(_MSC_VER) && !defined(__clang__)
591#pragma message("get_device_info: querying memory_clock_rate and \
592        memory_bus_width are not supported by the compiler used. \
593        Use 3200000 kHz as memory_clock_rate default value. \
594        Use 64 bits as memory_bus_width default value.")
595#else
596#warning "get_device_info: querying memory_clock_rate and \
597        memory_bus_width are not supported by the compiler used. \
598        Use 3200000 kHz as memory_clock_rate default value. \
599        Use 64 bits as memory_bus_width default value."
600#endif
601
602        size_t max_sub_group_size = 1;
603        std::vector<size_t> sub_group_sizes =
604            dev.get_info<sycl::info::device::sub_group_sizes>();
605
606        for (const auto &sub_group_size : sub_group_sizes)
607        {
608            if (max_sub_group_size < sub_group_size)
609                max_sub_group_size = sub_group_size;
610        }
611
612        prop.set_max_sub_group_size(max_sub_group_size);
613
614        prop.set_max_work_items_per_compute_unit(
615            dev.get_info<sycl::info::device::max_work_group_size>());
616        int max_nd_range_size[] = {0x7FFFFFFF, 0x7FFFFFFF, 0x7FFFFFFF};
617        prop.set_max_nd_range_size(max_nd_range_size);
618
619        // Estimates max register size per work group, feel free to update the value
620        // according to device properties.
621        prop.set_max_register_size_per_work_group(65536);
622
623        prop.set_global_mem_cache_size(
624            dev.get_info<sycl::info::device::global_mem_cache_size>());
625        out = prop;
626    }
627
628    /// dpct device extension
629    class device_ext : public sycl::device {
630      typedef std::mutex mutex_type;
631
632     public:
633      device_ext() : sycl::device() {}
634      ~device_ext() {
635        std::lock_guard<mutex_type> lock(m_mutex);
636        clear_queues();
637      }
638      device_ext(const sycl::device &base) : sycl::device(base) {
639        std::lock_guard<mutex_type> lock(m_mutex);
640        init_queues();
641      }
642
643      int is_native_atomic_supported() { return 0; }
644      int get_major_version() const { return dpct::get_major_version(*this); }
645
646      int get_minor_version() const { return dpct::get_minor_version(*this); }
647
648      int get_max_compute_units() const {
649        return get_device_info().get_max_compute_units();
650      }
651
652      /// Return the maximum clock frequency of this device in KHz.
653      int get_max_clock_frequency() const {
654        return get_device_info().get_max_clock_frequency();
655      }
656
657      int get_integrated() const { return get_device_info().get_integrated(); }
658
659      int get_max_sub_group_size() const {
660        return get_device_info().get_max_sub_group_size();
661      }
662
663      int get_max_register_size_per_work_group() const {
664        return get_device_info().get_max_register_size_per_work_group();
665      }
666
667      int get_max_work_group_size() const {
668        return get_device_info().get_max_work_group_size();
669      }
670
671      int get_mem_base_addr_align() const {
672        return get_info<sycl::info::device::mem_base_addr_align>();
673      }
674
675      size_t get_global_mem_size() const {
676        return get_device_info().get_global_mem_size();
677      }
678
679      size_t get_max_mem_alloc_size() const {
680        return get_device_info().get_max_mem_alloc_size();
681      }
682
683      /// Get the number of bytes of free and total memory on the SYCL device.
684      /// \param [out] free_memory The number of bytes of free memory on the
685      /// SYCL device. \param [out] total_memory The number of bytes of total
686      /// memory on the SYCL device.
687      void get_memory_info(size_t &free_memory, size_t &total_memory) {
688        total_memory = get_device_info().get_global_mem_size();
689        const char *warning_info =
690            "get_memory_info: [warning] ext_intel_free_memory is not "
691            "supported (export/set ZES_ENABLE_SYSMAN=1 to support), "
692            "use total memory as free memory";
693#if (defined(__SYCL_COMPILER_VERSION) && __SYCL_COMPILER_VERSION >= 20221105)
694        if (!has(sycl::aspect::ext_intel_free_memory)) {
695          std::cerr << warning_info << std::endl;
696          free_memory = total_memory;
697        } else {
698          free_memory = get_info<sycl::ext::intel::info::device::free_memory>();
699        }
700#else
701        std::cerr << warning_info << std::endl;
702        free_memory = total_memory;
703#if defined(_MSC_VER) && !defined(__clang__)
704#pragma message("Querying the number of bytes of free memory is not supported")
705#else
706#warning "Querying the number of bytes of free memory is not supported"
707#endif
708#endif
709      }
710
711      void get_device_info(device_info &out) const {
712        dpct::get_device_info(out, *this);
713      }
714
715      device_info get_device_info() const {
716        device_info prop;
717        dpct::get_device_info(prop, *this);
718        return prop;
719      }
720
721      void reset() {
722        std::lock_guard<mutex_type> lock(m_mutex);
723        clear_queues();
724        init_queues();
725      }
726
727      sycl::queue &in_order_queue() { return _q_in_order; }
728
729      sycl::queue &out_of_order_queue() { return _q_out_of_order; }
730
731      sycl::queue &default_queue() { return in_order_queue(); }
732
733      void queues_wait_and_throw() {
734        std::unique_lock<mutex_type> lock(m_mutex);
735        lock.unlock();
736        for (auto &q : _queues) {
737            q.wait_and_throw();
738        }
739        // Guard the destruct of current_queues to make sure the ref count is
740        // safe.
741        lock.lock();
742      }
743
744      sycl::queue create_queue(bool enable_exception_handler = false) {
745        return create_in_order_queue(enable_exception_handler);
746      }
747
748      sycl::queue create_queue(sycl::device device,
749                               bool enable_exception_handler = false) {
750        return create_in_order_queue(device, enable_exception_handler);
751      }
752
753      sycl::queue create_in_order_queue(bool enable_exception_handler = false) {
754        std::lock_guard<mutex_type> lock(m_mutex);
755        return create_queue_impl(enable_exception_handler,
756                                 sycl::property::queue::in_order());
757      }
758
759      sycl::queue create_in_order_queue(sycl::device device,
760                                        bool enable_exception_handler = false) {
761        std::lock_guard<mutex_type> lock(m_mutex);
762        return create_queue_impl(device, enable_exception_handler,
763                                 sycl::property::queue::in_order());
764      }
765
766      sycl::queue create_out_of_order_queue(
767          bool enable_exception_handler = false) {
768        std::lock_guard<mutex_type> lock(m_mutex);
769        return create_queue_impl(enable_exception_handler);
770      }
771
772      void destroy_queue(sycl::queue queue) {
773        std::lock_guard<mutex_type> lock(m_mutex);
774        _queues.erase(std::remove_if(_queues.begin(), _queues.end(),
775                                    [=](const sycl::queue &q) -> bool
776                                    {
777                                        return q == queue;
778                                    }),
779                    _queues.end());
780      }
781      void set_saved_queue(sycl::queue q) {
782        std::lock_guard<mutex_type> lock(m_mutex);
783        _saved_queue = q;
784      }
785      sycl::queue get_saved_queue() const {
786        std::lock_guard<mutex_type> lock(m_mutex);
787        return _saved_queue;
788      }
789
790     private:
791      void clear_queues() { _queues.clear(); }
792
793      void init_queues() {
794        _q_in_order =
795            create_queue_impl(true, sycl::property::queue::in_order());
796        _q_out_of_order = create_queue_impl(true);
797        _saved_queue = default_queue();
798      }
799
800      /// Caller should acquire resource \p m_mutex before calling this
801      /// function.
802      template <class... Properties>
803      sycl::queue create_queue_impl(bool enable_exception_handler,
804                                    Properties... properties) {
805        sycl::async_handler eh = {};
806        if (enable_exception_handler) {
807          eh = exception_handler;
808        }
809        _queues.push_back(sycl::queue(
810            *this, eh,
811            sycl::property_list(
812#ifdef DPCT_PROFILING_ENABLED
813                sycl::property::queue::enable_profiling(),
814#endif
815                properties...)));
816
817        return _queues.back();
818      }
819
820      template <class... Properties>
821      sycl::queue create_queue_impl(sycl::device device,
822                                    bool enable_exception_handler,
823                                    Properties... properties) {
824        sycl::async_handler eh = {};
825        if (enable_exception_handler) {
826          eh = exception_handler;
827        }
828        _queues.push_back(sycl::queue(
829            device, eh,
830                        sycl::property_list(
831#ifdef DPCT_PROFILING_ENABLED
832                            sycl::property::queue::enable_profiling(),
833#endif
834                            properties...)));
835
836        return _queues.back();
837      }
838
839      void get_version(int &major, int &minor) const {
840        detail::get_version(*this, major, minor);
841      }
842      sycl::queue _q_in_order, _q_out_of_order;
843      sycl::queue _saved_queue;
844      std::vector<sycl::queue> _queues;
845      mutable mutex_type m_mutex;
846    };
847
848
849    /// device manager
850    class dev_mgr
851    {
852    public:
853        device_ext &current_device()
854        {
855            unsigned int dev_id = current_device_id();
856            check_id(dev_id);
857            return *_devs[dev_id];
858        }
859        device_ext &cpu_device() const
860        {
861            std::lock_guard<std::recursive_mutex> lock(m_mutex);
862            if (_cpu_device == -1)
863            {
864                throw std::runtime_error("no valid cpu device");
865            }
866            else
867            {
868                return *_devs[_cpu_device];
869            }
870        }
871        device_ext &get_device(unsigned int id) const
872        {
873            std::lock_guard<std::recursive_mutex> lock(m_mutex);
874            check_id(id);
875            return *_devs[id];
876        }
877        unsigned int current_device_id() const
878        {
879            std::lock_guard<std::recursive_mutex> lock(m_mutex);
880            auto it = _thread2dev_map.find(get_tid());
881            if (it != _thread2dev_map.end())
882                return it->second;
883            return DEFAULT_DEVICE_ID;
884        }
885
886        /// Select device with a device ID.
887        /// \param [in] id The id of the device which can
888        /// be obtained through get_device_id(const sycl::device).
889        void select_device(unsigned int id)
890        {
891            std::lock_guard<std::recursive_mutex> lock(m_mutex);
892            check_id(id);
893            _thread2dev_map[get_tid()] = id;
894        }
895        unsigned int device_count() { return _devs.size(); }
896
897        unsigned int get_device_id(const sycl::device &dev)
898        {
899            unsigned int id = 0;
900            for (auto &dev_item : _devs)
901            {
902                if (*dev_item == dev)
903                {
904                    return id;
905                }
906                id++;
907            }
908            return -1;
909        }
910
911        inline std::string get_preferred_gpu_platform_name() {
912            std::string result;
913
914            std::string filter = "";
915            char* env = getenv("ONEAPI_DEVICE_SELECTOR");
916            if (env) {
917                if (std::strstr(env, "level_zero")) {
918                    filter = "level-zero";
919                }
920                else if (std::strstr(env, "opencl")) {
921                    filter = "opencl";
922                }
923                else if (std::strstr(env, "cuda")) {
924                    filter = "cuda";
925                }
926                else if (std::strstr(env, "hip")) {
927                    filter = "hip";
928                }
929                else {
930                    throw std::runtime_error("invalid device filter: " + std::string(env));
931                }
932            } else {
933                auto default_device = sycl::device(sycl::default_selector_v);
934                auto default_platform_name = default_device.get_platform().get_info<sycl::info::platform::name>();
935
936                if (std::strstr(default_platform_name.c_str(), "Level-Zero") || default_device.is_cpu()) {
937                    filter = "level-zero";
938                }
939                else if (std::strstr(default_platform_name.c_str(), "CUDA")) {
940                    filter = "cuda";
941                }
942                else if (std::strstr(default_platform_name.c_str(), "HIP")) {
943                    filter = "hip";
944                }
945            }
946
947            auto platform_list = sycl::platform::get_platforms();
948
949            for (const auto& platform : platform_list) {
950                auto devices = platform.get_devices();
951                auto gpu_dev = std::find_if(devices.begin(), devices.end(), [](const sycl::device& d) {
952                    return d.is_gpu();
953                });
954
955                if (gpu_dev == devices.end()) {
956                    // cout << "platform [" << platform_name
957                    //      << "] does not contain GPU devices, skipping\n";
958                    continue;
959                }
960
961                auto platform_name = platform.get_info<sycl::info::platform::name>();
962                std::string platform_name_low_case;
963                platform_name_low_case.resize(platform_name.size());
964
965                std::transform(
966                    platform_name.begin(), platform_name.end(), platform_name_low_case.begin(), ::tolower);
967
968                if (platform_name_low_case.find(filter) == std::string::npos) {
969                    // cout << "platform [" << platform_name
970                    //      << "] does not match with requested "
971                    //      << filter << ", skipping\n";
972                    continue;
973                }
974
975                result = platform_name;
976            }
977
978            if (result.empty())
979                throw std::runtime_error("can not find preferred GPU platform");
980
981            return result;
982        }
983
984        template <class DeviceSelector>
985        std::enable_if_t<
986            std::is_invocable_r_v<int, DeviceSelector, const sycl::device &>>
987        select_device(const DeviceSelector &selector = sycl::gpu_selector_v)
988        {
989            sycl::device selected_device = sycl::device(selector);
990            unsigned int selected_device_id = get_device_id(selected_device);
991            select_device(selected_device_id);
992        }
993
994        /// Returns the instance of device manager singleton.
995        static dev_mgr &instance()
996        {
997            static dev_mgr d_m;
998            return d_m;
999        }
1000        dev_mgr(const dev_mgr &) = delete;
1001        dev_mgr &operator=(const dev_mgr &) = delete;
1002        dev_mgr(dev_mgr &&) = delete;
1003        dev_mgr &operator=(dev_mgr &&) = delete;
1004
1005    private:
1006        mutable std::recursive_mutex m_mutex;
1007        static bool compare_dev(sycl::device &device1, sycl::device &device2)
1008        {
1009            sycl::backend backend1 = device1.get_backend();
1010            sycl::backend backend2 = device2.get_backend();
1011            // levelzero backends always come first
1012            if(backend1 == sycl::backend::ext_oneapi_level_zero && backend2 != sycl::backend::ext_oneapi_level_zero) return true;
1013            if(backend1 != sycl::backend::ext_oneapi_level_zero && backend2 == sycl::backend::ext_oneapi_level_zero) return false;
1014            dpct::device_info prop1;
1015            dpct::get_device_info(prop1, device1);
1016            dpct::device_info prop2;
1017            dpct::get_device_info(prop2, device2);
1018            return prop1.get_max_compute_units() > prop2.get_max_compute_units();
1019        }
1020        static int convert_backend_index(std::string & backend) {
1021            if (backend == "ext_oneapi_level_zero:gpu") return 0;
1022            if (backend == "opencl:gpu") return 1;
1023            if (backend == "ext_oneapi_cuda:gpu") return 2;
1024            if (backend == "ext_oneapi_hip:gpu") return 3;
1025            if (backend == "opencl:cpu") return 4;
1026            if (backend == "opencl:acc") return 5;
1027            printf("convert_backend_index: can't handle backend=%s\n", backend.c_str());
1028            _abort("fatal error");
1029        }
1030        static bool compare_backend(std::string &backend1, std::string &backend2) {
1031            return convert_backend_index(backend1) < convert_backend_index(backend2);
1032        }
1033        dev_mgr()
1034        {
1035            sycl::device default_device =
1036                sycl::device(sycl::default_selector_v);
1037            _devs.push_back(std::make_shared<device_ext>(default_device));
1038
1039            std::vector<sycl::device> sycl_all_devs;
1040            // Collect other devices except for the default device.
1041            if (default_device.is_cpu())
1042                _cpu_device = 0;
1043
1044            auto Platforms = sycl::platform::get_platforms();
1045            // Keep track of the number of devices per backend
1046            std::map<sycl::backend, size_t> DeviceNums;
1047            std::map<std::string, std::vector<sycl::device>> backend_devices;
1048            auto preferred_platform_name = get_preferred_gpu_platform_name();
1049
1050            while (!Platforms.empty()) {
1051                auto Platform = Platforms.back();
1052                Platforms.pop_back();
1053                auto platform_name = Platform.get_info<sycl::info::platform::name>();
1054                if (platform_name.compare(preferred_platform_name) != 0) {
1055                    continue;
1056                }
1057                auto devices = Platform.get_devices();
1058                std::string backend_type = get_device_backend_and_type(devices[0]);
1059                for (const auto &device : devices) {
1060                    backend_devices[backend_type].push_back(device);
1061                }
1062            }
1063
1064            std::vector<std::string> keys;
1065            for(auto it = backend_devices.begin(); it != backend_devices.end(); ++it) {
1066                keys.push_back(it->first);
1067            }
1068            std::sort(keys.begin(), keys.end(), compare_backend);
1069
1070            for (auto &key : keys) {
1071                std::vector<sycl::device> devs = backend_devices[key];
1072                std::sort(devs.begin(), devs.end(), compare_dev);
1073                for (const auto &dev : devs) {
1074                    sycl_all_devs.push_back(dev);
1075                }
1076            }
1077
1078            for (auto &dev : sycl_all_devs)
1079            {
1080                if (dev == default_device)
1081                {
1082                    continue;
1083                }
1084                _devs.push_back(std::make_shared<device_ext>(dev));
1085                if (_cpu_device == -1 && dev.is_cpu())
1086                {
1087                    _cpu_device = _devs.size() - 1;
1088                }
1089            }
1090        }
1091        void check_id(unsigned int id) const
1092        {
1093            if (id >= _devs.size())
1094            {
1095                throw std::runtime_error("invalid device id");
1096            }
1097        }
1098        std::vector<std::shared_ptr<device_ext>> _devs;
1099        /// DEFAULT_DEVICE_ID is used, if current_device_id() can not find current
1100        /// thread id in _thread2dev_map, which means default device should be used
1101        /// for the current thread.
1102        const unsigned int DEFAULT_DEVICE_ID = 0;
1103        /// thread-id to device-id map.
1104        std::map<unsigned int, unsigned int> _thread2dev_map;
1105        int _cpu_device = -1;
1106    };
1107
1108    static inline sycl::queue &get_default_queue()
1109    {
1110        return dev_mgr::instance().current_device().default_queue();
1111    }
1112
1113    namespace detail
1114    {
1115        enum class pointer_access_attribute
1116        {
1117            host_only = 0,
1118            device_only,
1119            host_device,
1120            end
1121        };
1122
1123        static pointer_access_attribute get_pointer_attribute(sycl::queue &q,
1124                                                              const void *ptr)
1125        {
1126            switch (sycl::get_pointer_type(ptr, q.get_context()))
1127            {
1128            case sycl::usm::alloc::unknown:
1129                return pointer_access_attribute::host_only;
1130            case sycl::usm::alloc::device:
1131                return pointer_access_attribute::device_only;
1132            case sycl::usm::alloc::shared:
1133            case sycl::usm::alloc::host:
1134                return pointer_access_attribute::host_device;
1135            }
1136        }
1137
1138        template <typename ArgT>
1139        inline constexpr std::uint64_t get_type_combination_id(ArgT Val)
1140        {
1141            static_assert((unsigned char)library_data_t::library_data_t_size <=
1142                              std::numeric_limits<unsigned char>::max() &&
1143                          "library_data_t size exceeds limit.");
1144            static_assert(std::is_same_v<ArgT, library_data_t>, "Unsupported ArgT");
1145            return (std::uint64_t)Val;
1146        }
1147
1148        template <typename FirstT, typename... RestT>
1149        inline constexpr std::uint64_t get_type_combination_id(FirstT FirstVal,
1150                                                               RestT... RestVal)
1151        {
1152            static_assert((std::uint8_t)library_data_t::library_data_t_size <=
1153                              std::numeric_limits<unsigned char>::max() &&
1154                          "library_data_t size exceeds limit.");
1155            static_assert(sizeof...(RestT) <= 8 && "Too many parameters");
1156            static_assert(std::is_same_v<FirstT, library_data_t>, "Unsupported FirstT");
1157            return get_type_combination_id(RestVal...) << 8 | ((std::uint64_t)FirstVal);
1158        }
1159
1160        class mem_mgr
1161        {
1162            mem_mgr()
1163            {
1164                // Reserved address space, no real memory allocation happens here.
1165#if defined(__linux__)
1166                mapped_address_space =
1167                    (byte_t *)mmap(nullptr, mapped_region_size, PROT_NONE,
1168                                   MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
1169#elif defined(_WIN64)
1170                mapped_address_space = (byte_t *)VirtualAlloc(
1171                    NULL,               // NULL specified as the base address parameter
1172                    mapped_region_size, // Size of allocation
1173                    MEM_RESERVE,        // Allocate reserved pages
1174                    PAGE_NOACCESS);     // Protection = no access
1175#else
1176#error "Only support Windows and Linux."
1177#endif
1178                next_free = mapped_address_space;
1179            }
1180
1181        public:
1182            using buffer_id_t = int;
1183
1184            struct allocation
1185            {
1186                buffer_t buffer;
1187                byte_t *alloc_ptr;
1188                size_t size;
1189            };
1190
1191            ~mem_mgr()
1192            {
1193#if defined(__linux__)
1194                munmap(mapped_address_space, mapped_region_size);
1195#elif defined(_WIN64)
1196                VirtualFree(mapped_address_space, 0, MEM_RELEASE);
1197#else
1198#error "Only support Windows and Linux."
1199#endif
1200            }
1201
1202            mem_mgr(const mem_mgr &) = delete;
1203            mem_mgr &operator=(const mem_mgr &) = delete;
1204            mem_mgr(mem_mgr &&) = delete;
1205            mem_mgr &operator=(mem_mgr &&) = delete;
1206
1207            /// Allocate
1208            void *mem_alloc(size_t size)
1209            {
1210                if (!size)
1211                    return nullptr;
1212                std::lock_guard<std::mutex> lock(m_mutex);
1213                if (next_free + size > mapped_address_space + mapped_region_size)
1214                {
1215                    throw std::runtime_error("dpct_malloc: out of memory for virtual memory pool");
1216                }
1217                // Allocation
1218                sycl::range<1> r(size);
1219                buffer_t buf(r);
1220                allocation A{buf, next_free, size};
1221                // Map allocation to device pointer
1222                void *result = next_free;
1223                m_map.emplace(next_free + size, A);
1224                // Update pointer to the next free space.
1225                next_free += (size + extra_padding + alignment - 1) & ~(alignment - 1);
1226
1227                return result;
1228            }
1229
1230            /// Deallocate
1231            void mem_free(const void *ptr)
1232            {
1233                if (!ptr)
1234                    return;
1235                std::lock_guard<std::mutex> lock(m_mutex);
1236                auto it = get_map_iterator(ptr);
1237                m_map.erase(it);
1238            }
1239
1240            /// map: device pointer -> allocation(buffer, alloc_ptr, size)
1241            allocation translate_ptr(const void *ptr)
1242            {
1243                std::lock_guard<std::mutex> lock(m_mutex);
1244                auto it = get_map_iterator(ptr);
1245                return it->second;
1246            }
1247
1248            /// Check if the pointer represents device pointer or not.
1249            bool is_device_ptr(const void *ptr) const
1250            {
1251                std::lock_guard<std::mutex> lock(m_mutex);
1252                return (mapped_address_space <= ptr) &&
1253                       (ptr < mapped_address_space + mapped_region_size);
1254            }
1255
1256            /// Returns the instance of memory manager singleton.
1257            static mem_mgr &instance()
1258            {
1259                static mem_mgr m;
1260                return m;
1261            }
1262
1263        private:
1264            std::map<byte_t *, allocation> m_map;
1265            mutable std::mutex m_mutex;
1266            byte_t *mapped_address_space;
1267            byte_t *next_free;
1268            const size_t mapped_region_size = 128ull * 1024 * 1024 * 1024;
1269            const size_t alignment = 256;
1270            /// This padding may be defined to some positive value to debug
1271            /// out of bound accesses.
1272            const size_t extra_padding = 0;
1273
1274            std::map<byte_t *, allocation>::iterator get_map_iterator(const void *ptr)
1275            {
1276                auto it = m_map.upper_bound(const_cast<byte_t *>(reinterpret_cast<const byte_t *>(ptr)));
1277                if (it == m_map.end())
1278                {
1279                    // Not a virtual pointer.
1280                    throw std::runtime_error("can not get buffer from non-virtual pointer");
1281                }
1282                const allocation &alloc = it->second;
1283                if (ptr < alloc.alloc_ptr)
1284                {
1285                    // Out of bound.
1286                    // This may happen if there's a gap between allocations due to alignment
1287                    // or extra padding and pointer points to this gap.
1288                    throw std::runtime_error("invalid virtual pointer");
1289                }
1290                return it;
1291            }
1292        };
1293
1294        template <class T, memory_region Memory, size_t Dimension>
1295        class accessor;
1296        template <memory_region Memory, class T = byte_t>
1297        class memory_traits
1298        {
1299        public:
1300            static constexpr sycl::access::target target =
1301                sycl::access::target::device;
1302            static constexpr sycl::access_mode mode =
1303                (Memory == constant) ? sycl::access_mode::read
1304                                     : sycl::access_mode::read_write;
1305            static constexpr size_t type_size = sizeof(T);
1306            using element_t =
1307                typename std::conditional<Memory == constant, const T, T>::type;
1308            using value_t = typename std::remove_cv<T>::type;
1309            template <size_t Dimension = 1>
1310            using accessor_t = typename std::conditional<
1311                Memory == local, sycl::local_accessor<value_t, Dimension>,
1312                sycl::accessor<T, Dimension, mode, target>>::type;
1313            using pointer_t = T *;
1314        };
1315
1316        static inline void *dpct_malloc(size_t size, sycl::queue &q)
1317        {
1318            return sycl::malloc_device(size, q.get_device(), q.get_context());
1319        }
1320
1321#define PITCH_DEFAULT_ALIGN(x) (((x) + 31) & ~(0x1F))
1322        static inline void *dpct_malloc(size_t &pitch, size_t x, size_t y, size_t z,
1323                                        sycl::queue &q)
1324        {
1325            pitch = PITCH_DEFAULT_ALIGN(x);
1326            return dpct_malloc(pitch * y * z, q);
1327        }
1328
1329        /**
1330         * @brief Sets \p value to the first \p size elements starting from \p dev_ptr in \p q.
1331         * @tparam valueT The type of the element to be set.
1332         * @param [in] q The queue in which the operation is done.
1333         * @param [in] dev_ptr Pointer to the virtual device memory address.
1334         * @param [in] value The value to be set.
1335         * @param [in] size Number of elements to be set to the value.
1336         * @return An event representing the memset operation.
1337         */
1338        template <typename valueT>
1339        static inline sycl::event dpct_memset(sycl::queue &q, void *dev_ptr,
1340                                              valueT value, size_t size)
1341        {
1342            return q.fill(dev_ptr, value, size);
1343        }
1344
1345        /**
1346         * @brief Sets \p value to the 3D memory region pointed by \p data in \p q.
1347         * @tparam valueT The type of the element to be set.
1348         * @param [in] q The queue in which the operation is done.
1349         * @param [in] data Pointer to the pitched device memory region.
1350         * @param [in] value The value to be set.
1351         * @param [in] size 3D memory region by number of elements.
1352         * @return An event list representing the memset operations.
1353         */
1354        template <typename valueT>
1355        static inline std::vector<sycl::event>
1356        dpct_memset(sycl::queue &q, pitched_data data, valueT value,
1357                    sycl::range<3> size)
1358        {
1359            std::vector<sycl::event> event_list;
1360            size_t slice = data.get_pitch() * data.get_y();
1361            unsigned char *data_surface = (unsigned char *)data.get_data_ptr();
1362            for (size_t z = 0; z < size.get(2); ++z)
1363            {
1364                unsigned char *data_ptr = data_surface;
1365                for (size_t y = 0; y < size.get(1); ++y)
1366                {
1367                    event_list.push_back(dpct_memset(q, data_ptr, value, size.get(0)));
1368                    data_ptr += data.get_pitch();
1369                }
1370                data_surface += slice;
1371            }
1372            return event_list;
1373        }
1374
1375        /**
1376         * @brief Sets \p val to the pitched 2D memory region pointed by \p ptr in \p q.
1377         * @tparam valueT The type of the element to be set.
1378         * @param [in] q The queue in which the operation is done.
1379         * @param [in] ptr Pointer to the virtual device memory.
1380         * @param [in] pitch The pitch size by number of elements, including padding.
1381         * @param [in] val The value to be set.
1382         * @param [in] x The width of memory region by number of elements.
1383         * @param [in] y The height of memory region by number of elements.
1384         * @return An event list representing the memset operations.
1385         */
1386        template <typename valueT>
1387        static inline std::vector<sycl::event>
1388        dpct_memset(sycl::queue &q, void *ptr, size_t pitch, valueT val, size_t x,
1389                    size_t y)
1390        {
1391            return dpct_memset(q, pitched_data(ptr, pitch, x, 1), val,
1392                               sycl::range<3>(x, y, 1));
1393        }
1394
1395        static memcpy_direction deduce_memcpy_direction(sycl::queue &q, void *to_ptr,
1396                                                        const void *from_ptr,
1397                                                        memcpy_direction dir)
1398        {
1399            switch (dir)
1400            {
1401            case memcpy_direction::host_to_host:
1402            case memcpy_direction::host_to_device:
1403            case memcpy_direction::device_to_host:
1404            case memcpy_direction::device_to_device:
1405                return dir;
1406            case memcpy_direction::automatic:
1407            {
1408                // table[to_attribute][from_attribute]
1409                static const memcpy_direction
1410                    direction_table[static_cast<unsigned>(pointer_access_attribute::end)]
1411                                   [static_cast<unsigned>(pointer_access_attribute::end)] =
1412                                       {{memcpy_direction::host_to_host,
1413                                         memcpy_direction::device_to_host,
1414                                         memcpy_direction::host_to_host},
1415                                        {memcpy_direction::host_to_device,
1416                                         memcpy_direction::device_to_device,
1417                                         memcpy_direction::device_to_device},
1418                                        {memcpy_direction::host_to_host,
1419                                         memcpy_direction::device_to_device,
1420                                         memcpy_direction::device_to_device}};
1421                return direction_table[static_cast<unsigned>(get_pointer_attribute(
1422                    q, to_ptr))][static_cast<unsigned>(get_pointer_attribute(q, from_ptr))];
1423            }
1424            default:
1425                throw std::runtime_error("dpct_memcpy: invalid direction value");
1426            }
1427        }
1428
1429        static sycl::event
1430        dpct_memcpy(sycl::queue &q, void *to_ptr, const void *from_ptr, size_t size,
1431                    memcpy_direction direction,
1432                    const std::vector<sycl::event> &dep_events = {})
1433        {
1434            if (!size)
1435                return sycl::event{};
1436            return q.memcpy(to_ptr, from_ptr, size, dep_events);
1437            DPCT_UNUSED(direction);
1438        }
1439
1440        // Get actual copy range and make sure it will not exceed range.
1441        static inline size_t get_copy_range(sycl::range<3> size, size_t slice,
1442                                            size_t pitch)
1443        {
1444            return slice * (size.get(2) - 1) + pitch * (size.get(1) - 1) + size.get(0);
1445        }
1446
1447        static inline size_t get_offset(sycl::id<3> id, size_t slice,
1448                                        size_t pitch)
1449        {
1450            return slice * id.get(2) + pitch * id.get(1) + id.get(0);
1451        }
1452
1453        /// copy 3D matrix specified by \p size from 3D matrix specified by \p from_ptr
1454        /// and \p from_range to another specified by \p to_ptr and \p to_range.
1455        static inline std::vector<sycl::event>
1456        dpct_memcpy(sycl::queue &q, void *to_ptr, const void *from_ptr,
1457                    sycl::range<3> to_range, sycl::range<3> from_range,
1458                    sycl::id<3> to_id, sycl::id<3> from_id,
1459                    sycl::range<3> size, memcpy_direction direction,
1460                    const std::vector<sycl::event> &dep_events = {})
1461        {
1462            // RAII for host pointer
1463            class host_buffer
1464            {
1465                void *_buf;
1466                size_t _size;
1467                sycl::queue &_q;
1468                const std::vector<sycl::event> &_deps; // free operation depends
1469
1470            public:
1471                host_buffer(size_t size, sycl::queue &q,
1472                            const std::vector<sycl::event> &deps)
1473                    : _buf(std::malloc(size)), _size(size), _q(q), _deps(deps) {}
1474                void *get_ptr() const { return _buf; }
1475                size_t get_size() const { return _size; }
1476                ~host_buffer()
1477                {
1478                    if (_buf)
1479                    {
1480                        _q.submit([&](sycl::handler &cgh)
1481                                  {
1482        cgh.depends_on(_deps);
1483        cgh.host_task([buf = _buf] { std::free(buf); }); });
1484                    }
1485                }
1486            };
1487            std::vector<sycl::event> event_list;
1488
1489            size_t to_slice = to_range.get(1) * to_range.get(0),
1490                   from_slice = from_range.get(1) * from_range.get(0);
1491            unsigned char *to_surface =
1492                (unsigned char *)to_ptr + get_offset(to_id, to_slice, to_range.get(0));
1493            const unsigned char *from_surface =
1494                (const unsigned char *)from_ptr +
1495                get_offset(from_id, from_slice, from_range.get(0));
1496
1497            if (to_slice == from_slice && to_slice == size.get(1) * size.get(0))
1498            {
1499                return {dpct_memcpy(q, to_surface, from_surface, to_slice * size.get(2),
1500                                    direction, dep_events)};
1501            }
1502            direction = deduce_memcpy_direction(q, to_ptr, from_ptr, direction);
1503            size_t size_slice = size.get(1) * size.get(0);
1504            switch (direction)
1505            {
1506            case host_to_host:
1507                for (size_t z = 0; z < size.get(2); ++z)
1508                {
1509                    unsigned char *to_ptr = to_surface;
1510                    const unsigned char *from_ptr = from_surface;
1511                    if (to_range.get(0) == from_range.get(0) &&
1512                        to_range.get(0) == size.get(0))
1513                    {
1514                        event_list.push_back(dpct_memcpy(q, to_ptr, from_ptr, size_slice,
1515                                                         direction, dep_events));
1516                    }
1517                    else
1518                    {
1519                        for (size_t y = 0; y < size.get(1); ++y)
1520                        {
1521                            event_list.push_back(dpct_memcpy(q, to_ptr, from_ptr, size.get(0),
1522                                                             direction, dep_events));
1523                            to_ptr += to_range.get(0);
1524                            from_ptr += from_range.get(0);
1525                        }
1526                    }
1527                    to_surface += to_slice;
1528                    from_surface += from_slice;
1529                }
1530                break;
1531            case host_to_device:
1532            {
1533                host_buffer buf(get_copy_range(size, to_slice, to_range.get(0)), q,
1534                                event_list);
1535                std::vector<sycl::event> host_events;
1536                if (to_slice == size_slice)
1537                {
1538                    // Copy host data to a temp host buffer with the shape of target.
1539                    host_events =
1540                        dpct_memcpy(q, buf.get_ptr(), from_surface, to_range, from_range,
1541                                    sycl::id<3>(0, 0, 0), sycl::id<3>(0, 0, 0), size,
1542                                    host_to_host, dep_events);
1543                }
1544                else
1545                {
1546                    // Copy host data to a temp host buffer with the shape of target.
1547                    host_events = dpct_memcpy(
1548                        q, buf.get_ptr(), from_surface, to_range, from_range,
1549                        sycl::id<3>(0, 0, 0), sycl::id<3>(0, 0, 0), size, host_to_host,
1550                        // If has padding data, not sure whether it is useless. So fill temp
1551                        // buffer with it.
1552                        std::vector<sycl::event>{
1553                            dpct_memcpy(q, buf.get_ptr(), to_surface, buf.get_size(),
1554                                        device_to_host, dep_events)});
1555                }
1556                // Copy from temp host buffer to device with only one submit.
1557                event_list.push_back(dpct_memcpy(q, to_surface, buf.get_ptr(),
1558                                                 buf.get_size(), host_to_device,
1559                                                 host_events));
1560                break;
1561            }
1562            case device_to_host:
1563            {
1564                host_buffer buf(get_copy_range(size, from_slice, from_range.get(0)), q,
1565                                event_list);
1566                // Copy from host temp buffer to host target with reshaping.
1567                event_list = dpct_memcpy(
1568                    q, to_surface, buf.get_ptr(), to_range, from_range, sycl::id<3>(0, 0, 0),
1569                    sycl::id<3>(0, 0, 0), size, host_to_host,
1570                    // Copy from device to temp host buffer with only one submit.
1571                    std::vector<sycl::event>{dpct_memcpy(q, buf.get_ptr(), from_surface,
1572                                                         buf.get_size(),
1573                                                         device_to_host, dep_events)});
1574                break;
1575            }
1576            case device_to_device:
1577                event_list.push_back(q.submit([&](sycl::handler &cgh){
1578                cgh.depends_on(dep_events);
1579                cgh.parallel_for<class dpct_memcpy_3d_detail>(
1580                    size,
1581                    [=](sycl::id<3> id) {
1582                        to_surface[get_offset(id, to_slice, to_range.get(0))] =
1583                            from_surface[get_offset(id, from_slice, from_range.get(0))];
1584                    }); }));
1585                break;
1586            default:
1587                throw std::runtime_error("dpct_memcpy: invalid direction value");
1588            }
1589            return event_list;
1590        }
1591
1592        /// memcpy 2D/3D matrix specified by pitched_data.
1593        static inline std::vector<sycl::event>
1594        dpct_memcpy(sycl::queue &q, pitched_data to, sycl::id<3> to_id,
1595                    pitched_data from, sycl::id<3> from_id, sycl::range<3> size,
1596                    memcpy_direction direction = automatic)
1597        {
1598            return dpct_memcpy(q, to.get_data_ptr(), from.get_data_ptr(),
1599                               sycl::range<3>(to.get_pitch(), to.get_y(), 1),
1600                               sycl::range<3>(from.get_pitch(), from.get_y(), 1), to_id, from_id,
1601                               size, direction);
1602        }
1603
1604        /// memcpy 2D matrix with pitch.
1605        static inline std::vector<sycl::event>
1606        dpct_memcpy(sycl::queue &q, void *to_ptr, const void *from_ptr,
1607                    size_t to_pitch, size_t from_pitch, size_t x, size_t y,
1608                    memcpy_direction direction = automatic)
1609        {
1610            return dpct_memcpy(q, to_ptr, from_ptr, sycl::range<3>(to_pitch, y, 1),
1611                               sycl::range<3>(from_pitch, y, 1),
1612                               sycl::id<3>(0, 0, 0), sycl::id<3>(0, 0, 0),
1613                               sycl::range<3>(x, y, 1), direction);
1614        }
1615
1616        namespace deprecated
1617        {
1618
1619            template <typename T, sycl::usm::alloc AllocKind>
1620            class usm_allocator
1621            {
1622            private:
1623                using Alloc = sycl::usm_allocator<T, AllocKind>;
1624                Alloc _impl;
1625
1626            public:
1627                using value_type = typename std::allocator_traits<Alloc>::value_type;
1628                using pointer = typename std::allocator_traits<Alloc>::pointer;
1629                using const_pointer = typename std::allocator_traits<Alloc>::const_pointer;
1630                using void_pointer = typename std::allocator_traits<Alloc>::void_pointer;
1631                using const_void_pointer =
1632                    typename std::allocator_traits<Alloc>::const_void_pointer;
1633                using reference = typename std::allocator_traits<Alloc>::value_type &;
1634                using const_reference =
1635                    const typename std::allocator_traits<Alloc>::value_type &;
1636                using difference_type =
1637                    typename std::allocator_traits<Alloc>::difference_type;
1638                using size_type = typename std::allocator_traits<Alloc>::size_type;
1639                using propagate_on_container_copy_assignment = typename std::allocator_traits<
1640                    Alloc>::propagate_on_container_copy_assignment;
1641                using propagate_on_container_move_assignment = typename std::allocator_traits<
1642                    Alloc>::propagate_on_container_move_assignment;
1643                using propagate_on_container_swap =
1644                    typename std::allocator_traits<Alloc>::propagate_on_container_swap;
1645                using is_always_equal =
1646                    typename std::allocator_traits<Alloc>::is_always_equal;
1647
1648                template <typename U>
1649                struct rebind
1650                {
1651                    typedef usm_allocator<U, AllocKind> other;
1652                };
1653
1654                usm_allocator() : _impl(dpct::get_default_queue()) {}
1655                ~usm_allocator() {}
1656                usm_allocator(const usm_allocator &other) : _impl(other._impl) {}
1657                usm_allocator(usm_allocator &&other) : _impl(std::move(other._impl)) {}
1658                pointer address(reference r) { return &r; }
1659                const_pointer address(const_reference r) { return &r; }
1660                pointer allocate(size_type cnt, const_void_pointer hint = nullptr)
1661                {
1662                    return std::allocator_traits<Alloc>::allocate(_impl, cnt, hint);
1663                }
1664                void deallocate(pointer p, size_type cnt)
1665                {
1666                    std::allocator_traits<Alloc>::deallocate(_impl, p, cnt);
1667                }
1668                size_type max_size() const
1669                {
1670                    return std::allocator_traits<Alloc>::max_size(_impl);
1671                }
1672                bool operator==(const usm_allocator &other) const { return _impl == other._impl; }
1673                bool operator!=(const usm_allocator &other) const { return _impl != other._impl; }
1674            };
1675
1676        } // namespace deprecated
1677
1678        inline void dpct_free(void *ptr,
1679                              const sycl::queue &q)
1680        {
1681            if (ptr)
1682            {
1683                sycl::free(ptr, q.get_context());
1684            }
1685        }
1686
1687        template <typename T>
1688        inline auto get_memory(const void *x)
1689        {
1690            T *new_x = reinterpret_cast<T *>(const_cast<void *>(x));
1691            return new_x;
1692        }
1693
1694        template <typename T>
1695        inline typename DataType<T>::T2 get_value(const T *s, sycl::queue &q)
1696        {
1697            using Ty = typename DataType<T>::T2;
1698            Ty s_h;
1699            if (get_pointer_attribute(q, s) == pointer_access_attribute::device_only)
1700                detail::dpct_memcpy(q, (void *)&s_h, (const void *)s, sizeof(T), device_to_host)
1701                    .wait();
1702            else
1703                s_h = *reinterpret_cast<const Ty *>(s);
1704            return s_h;
1705        }
1706
1707    } // namespace detail
1708
1709    template <typename T>
1710    inline auto get_value(const T *s, sycl::queue &q)
1711    {
1712        return detail::get_value(s, q);
1713    }
1714
1715    namespace detail
1716    {
1717    template <class Ta, class Tb, class Tc, class Ts>
1718    inline void gemm_impl(sycl::queue & q, oneapi::mkl::transpose a_trans, oneapi::mkl::transpose b_trans, int m,
1719                          int n, int k, const void * alpha, const void * a, int lda, const void * b, int ldb,
1720                          const void * beta, void * c, int ldc) {
1721        Ts   alpha_value = dpct::get_value(reinterpret_cast<const Ts *>(alpha), q);
1722        Ts   beta_value  = dpct::get_value(reinterpret_cast<const Ts *>(beta), q);
1723        auto data_a      = get_memory<const Ta>(a);
1724        auto data_b      = get_memory<const Tb>(b);
1725        auto data_c      = get_memory<Tc>(c);
1726        oneapi::mkl::blas::column_major::gemm(q, a_trans, b_trans, m, n, k, alpha_value, data_a,
1727                                               lda, data_b, ldb, beta_value, data_c, ldc);
1728    }
1729
1730        template <typename VecT, class BinaryOperation, class = void>
1731        class vectorized_binary
1732        {
1733        public:
1734            inline VecT operator()(VecT a, VecT b, const BinaryOperation binary_op)
1735            {
1736                VecT v4;
1737                for (size_t i = 0; i < v4.size(); ++i)
1738                {
1739                    v4[i] = binary_op(a[i], b[i]);
1740                }
1741                return v4;
1742            }
1743        };
1744
1745        template <typename VecT, class BinaryOperation>
1746        class vectorized_binary<
1747            VecT, BinaryOperation,
1748            std::void_t<std::invoke_result_t<BinaryOperation, VecT, VecT>>>
1749        {
1750        public:
1751            inline VecT operator()(VecT a, VecT b, const BinaryOperation binary_op)
1752            {
1753                return binary_op(a, b).template as<VecT>();
1754            }
1755        };
1756
1757        template <class Ta, class Tb, class Tc, class Ts>
1758        inline void gemm_batch_impl(sycl::queue & q, oneapi::mkl::transpose a_trans, oneapi::mkl::transpose b_trans,
1759                                    int m, int n, int k, const void * alpha, const void ** a, int lda, const void ** b,
1760                                    int ldb, const void * beta, void ** c, int ldc, int batch_size,
1761                                    matrix_info_t<float> * matrix_info) {
1762            Ts alpha_value = dpct::get_value(reinterpret_cast<const Ts *>(alpha), q);
1763            Ts beta_value = dpct::get_value(reinterpret_cast<const Ts *>(beta), q);
1764
1765            matrix_info->transpose_info[0] = a_trans;
1766            matrix_info->transpose_info[1] = b_trans;
1767            matrix_info->value_info[0] = alpha_value;
1768            matrix_info->value_info[1] = beta_value;
1769            matrix_info->size_info[0] = m;
1770            matrix_info->size_info[1] = n;
1771            matrix_info->size_info[2] = k;
1772            matrix_info->ld_info[0] = lda;
1773            matrix_info->ld_info[1] = ldb;
1774            matrix_info->ld_info[2] = ldc;
1775            matrix_info->groupsize_info = batch_size;
1776
1777            sycl::event e = oneapi::mkl::blas::column_major::gemm_batch(
1778                q, matrix_info->transpose_info, matrix_info->transpose_info + 1,
1779                matrix_info->size_info, matrix_info->size_info + 1, matrix_info->size_info + 2,
1780                reinterpret_cast<Ts *>(matrix_info->value_info), reinterpret_cast<const Ta **>(a), matrix_info->ld_info,
1781                reinterpret_cast<const Tb **>(b), matrix_info->ld_info + 1,
1782                reinterpret_cast<Ts *>(matrix_info->value_info + 1), reinterpret_cast<Tc **>(c),
1783                matrix_info->ld_info + 2, 1, &(matrix_info->groupsize_info));
1784        }
1785
1786        template <class Ta, class Tb, class Tc, class Ts>
1787        inline void gemm_batch_impl(sycl::queue & q, oneapi::mkl::transpose a_trans, oneapi::mkl::transpose b_trans,
1788                                    int m, int n, int k, const void * alpha, const void * a, int lda,
1789                                    long long int stride_a, const void * b, int ldb, long long int stride_b,
1790                                    const void * beta, void * c, int ldc, long long int stride_c, int batch_size) {
1791            Ts alpha_value = dpct::get_value(reinterpret_cast<const Ts *>(alpha), q);
1792            Ts beta_value = dpct::get_value(reinterpret_cast<const Ts *>(beta), q);
1793            auto data_a = get_memory<const Ta>(a);
1794            auto data_b = get_memory<const Tb>(b);
1795            auto data_c = get_memory<Tc>(c);
1796            oneapi::mkl::blas::column_major::gemm_batch(q, a_trans, b_trans, m, n, k, alpha_value,
1797                                                         data_a, lda, stride_a, data_b, ldb, stride_b, beta_value,
1798                                                         data_c, ldc, stride_c, batch_size);
1799        }
1800
1801    } // namespace detail
1802
1803    template <typename VecT, class BinaryOperation>
1804    inline unsigned vectorized_binary(unsigned a, unsigned b,
1805                                      const BinaryOperation binary_op)
1806    {
1807        sycl::vec<unsigned, 1> v0{a}, v1{b};
1808        auto v2 = v0.as<VecT>();
1809        auto v3 = v1.as<VecT>();
1810        auto v4 =
1811            detail::vectorized_binary<VecT, BinaryOperation>()(v2, v3, binary_op);
1812        v0 = v4.template as<sycl::vec<unsigned, 1>>();
1813        return v0;
1814    }
1815
1816    static void async_dpct_memcpy(void *to_ptr, const void *from_ptr, size_t size,
1817                                  memcpy_direction direction = automatic,
1818                                  sycl::queue &q = dpct::get_default_queue())
1819    {
1820        detail::dpct_memcpy(q, to_ptr, from_ptr, size, direction);
1821    }
1822
1823    static inline unsigned int select_device(unsigned int id)
1824    {
1825        dev_mgr::instance().select_device(id);
1826        return id;
1827    }
1828
1829    template <typename T>
1830    T permute_sub_group_by_xor(sycl::sub_group g, T x, unsigned int mask,
1831                               unsigned int logical_sub_group_size = 32)
1832    {
1833        unsigned int id = g.get_local_linear_id();
1834        unsigned int start_index =
1835            id / logical_sub_group_size * logical_sub_group_size;
1836        unsigned int target_offset = (id % logical_sub_group_size) ^ mask;
1837        return sycl::select_from_group(g, x,
1838                                       target_offset < logical_sub_group_size
1839                                           ? start_index + target_offset
1840                                           : id);
1841    }
1842
1843    template <typename T1, typename T2>
1844    using dot_product_acc_t = std::conditional_t<
1845        std::is_unsigned_v<T1> && std::is_unsigned_v<T2>,
1846        uint32_t,
1847        int32_t>;
1848
1849    template <typename T>
1850    sycl::vec<T, 4> extract_and_sign_or_zero_extend4(T val) {
1851      return sycl::vec<T, 1>(val)
1852          .template as<sycl::vec<
1853              std::conditional_t<std::is_signed_v<T>, int8_t, uint8_t>,
1854              4>>()
1855          .template convert<T>();
1856    }
1857
1858    template <typename T1, typename T2, typename T3>
1859    inline auto dp4a(T1 a, T2 b, T3 c) {
1860      dot_product_acc_t<T1, T2> res = c;
1861      auto va = extract_and_sign_or_zero_extend4(a);
1862      auto vb = extract_and_sign_or_zero_extend4(b);
1863      res += va[0] * vb[0];
1864      res += va[1] * vb[1];
1865      res += va[2] * vb[2];
1866      res += va[3] * vb[3];
1867      return res;
1868    }
1869
1870    struct sub_sat
1871    {
1872        template <typename T>
1873        auto operator()(const T x, const T y) const
1874        {
1875            return sycl::sub_sat(x, y);
1876        }
1877    };
1878
1879    template <typename S, typename T>
1880    inline T vectorized_min(T a, T b)
1881    {
1882        sycl::vec<T, 1> v0{a}, v1{b};
1883        auto v2 = v0.template as<S>();
1884        auto v3 = v1.template as<S>();
1885        auto v4 = sycl::min(v2, v3);
1886        v0 = v4.template as<sycl::vec<T, 1>>();
1887        return v0;
1888    }
1889
1890    inline float pow(const float a, const int b) { return sycl::pown(a, b); }
1891    inline double pow(const double a, const int b) { return sycl::pown(a, b); }
1892    inline float pow(const float a, const float b) { return sycl::pow(a, b); }
1893    inline double pow(const double a, const double b) { return sycl::pow(a, b); }
1894    template <typename T, typename U>
1895    inline typename std::enable_if_t<std::is_floating_point_v<T>, T>
1896    pow(const T a, const U b)
1897    {
1898        return sycl::pow(a, static_cast<T>(b));
1899    }
1900    template <typename T, typename U>
1901    inline typename std::enable_if_t<!std::is_floating_point_v<T>, double>
1902    pow(const T a, const U b)
1903    {
1904        return sycl::pow(static_cast<double>(a), static_cast<double>(b));
1905    }
1906
1907    inline double min(const double a, const float b)
1908    {
1909        return sycl::fmin(a, static_cast<double>(b));
1910    }
1911    inline double min(const float a, const double b)
1912    {
1913        return sycl::fmin(static_cast<double>(a), b);
1914    }
1915    inline float min(const float a, const float b) { return sycl::fmin(a, b); }
1916    inline double min(const double a, const double b) { return sycl::fmin(a, b); }
1917    inline std::uint32_t min(const std::uint32_t a, const std::int32_t b)
1918    {
1919        return sycl::min(a, static_cast<std::uint32_t>(b));
1920    }
1921    inline std::uint32_t min(const std::int32_t a, const std::uint32_t b)
1922    {
1923        return sycl::min(static_cast<std::uint32_t>(a), b);
1924    }
1925    inline std::int32_t min(const std::int32_t a, const std::int32_t b)
1926    {
1927        return sycl::min(a, b);
1928    }
1929    inline std::uint32_t min(const std::uint32_t a, const std::uint32_t b)
1930    {
1931        return sycl::min(a, b);
1932    }
1933    inline std::uint64_t min(const std::uint64_t a, const std::int64_t b)
1934    {
1935        return sycl::min(a, static_cast<std::uint64_t>(b));
1936    }
1937    inline std::uint64_t min(const std::int64_t a, const std::uint64_t b)
1938    {
1939        return sycl::min(static_cast<std::uint64_t>(a), b);
1940    }
1941    inline std::int64_t min(const std::int64_t a, const std::int64_t b)
1942    {
1943        return sycl::min(a, b);
1944    }
1945    inline std::uint64_t min(const std::uint64_t a, const std::uint64_t b)
1946    {
1947        return sycl::min(a, b);
1948    }
1949    inline std::uint64_t min(const std::uint64_t a, const std::int32_t b)
1950    {
1951        return sycl::min(a, static_cast<std::uint64_t>(b));
1952    }
1953    inline std::uint64_t min(const std::int32_t a, const std::uint64_t b)
1954    {
1955        return sycl::min(static_cast<std::uint64_t>(a), b);
1956    }
1957    inline std::uint64_t min(const std::uint64_t a, const std::uint32_t b)
1958    {
1959        return sycl::min(a, static_cast<std::uint64_t>(b));
1960    }
1961    inline std::uint64_t min(const std::uint32_t a, const std::uint64_t b)
1962    {
1963        return sycl::min(static_cast<std::uint64_t>(a), b);
1964    }
1965    // max function overloads.
1966    // For floating-point types, `float` or `double` arguments are acceptable.
1967    // For integer types, `std::uint32_t`, `std::int32_t`, `std::uint64_t` or
1968    // `std::int64_t` type arguments are acceptable.
1969    inline double max(const double a, const float b)
1970    {
1971        return sycl::fmax(a, static_cast<double>(b));
1972    }
1973    inline double max(const float a, const double b)
1974    {
1975        return sycl::fmax(static_cast<double>(a), b);
1976    }
1977    inline float max(const float a, const float b) { return sycl::fmax(a, b); }
1978    inline double max(const double a, const double b) { return sycl::fmax(a, b); }
1979    inline std::uint32_t max(const std::uint32_t a, const std::int32_t b)
1980    {
1981        return sycl::max(a, static_cast<std::uint32_t>(b));
1982    }
1983    inline std::uint32_t max(const std::int32_t a, const std::uint32_t b)
1984    {
1985        return sycl::max(static_cast<std::uint32_t>(a), b);
1986    }
1987    inline std::int32_t max(const std::int32_t a, const std::int32_t b)
1988    {
1989        return sycl::max(a, b);
1990    }
1991    inline std::uint32_t max(const std::uint32_t a, const std::uint32_t b)
1992    {
1993        return sycl::max(a, b);
1994    }
1995    inline std::uint64_t max(const std::uint64_t a, const std::int64_t b)
1996    {
1997        return sycl::max(a, static_cast<std::uint64_t>(b));
1998    }
1999    inline std::uint64_t max(const std::int64_t a, const std::uint64_t b)
2000    {
2001        return sycl::max(static_cast<std::uint64_t>(a), b);
2002    }
2003    inline std::int64_t max(const std::int64_t a, const std::int64_t b)
2004    {
2005        return sycl::max(a, b);
2006    }
2007    inline std::uint64_t max(const std::uint64_t a, const std::uint64_t b)
2008    {
2009        return sycl::max(a, b);
2010    }
2011    inline std::uint64_t max(const std::uint64_t a, const std::int32_t b)
2012    {
2013        return sycl::max(a, static_cast<std::uint64_t>(b));
2014    }
2015    inline std::uint64_t max(const std::int32_t a, const std::uint64_t b)
2016    {
2017        return sycl::max(static_cast<std::uint64_t>(a), b);
2018    }
2019    inline std::uint64_t max(const std::uint64_t a, const std::uint32_t b)
2020    {
2021        return sycl::max(a, static_cast<std::uint64_t>(b));
2022    }
2023    inline std::uint64_t max(const std::uint32_t a, const std::uint64_t b)
2024    {
2025        return sycl::max(static_cast<std::uint64_t>(a), b);
2026    }
2027
2028    inline void
2029    has_capability_or_fail(const sycl::device &dev,
2030                           const std::initializer_list<sycl::aspect> &props)
2031    {
2032        for (const auto &it : props)
2033        {
2034            if (dev.has(it))
2035                continue;
2036            switch (it)
2037            {
2038            case sycl::aspect::fp64:
2039                throw std::runtime_error("'double' is not supported in '" +
2040                                         dev.get_info<sycl::info::device::name>() +
2041                                         "' device");
2042                break;
2043            case sycl::aspect::fp16:
2044                throw std::runtime_error("'half' is not supported in '" +
2045                                         dev.get_info<sycl::info::device::name>() +
2046                                         "' device");
2047                break;
2048            default:
2049#define __SYCL_ASPECT(ASPECT, ID) \
2050    case sycl::aspect::ASPECT:    \
2051        return #ASPECT;
2052#define __SYCL_ASPECT_DEPRECATED(ASPECT, ID, MESSAGE) __SYCL_ASPECT(ASPECT, ID)
2053#define __SYCL_ASPECT_DEPRECATED_ALIAS(ASPECT, ID, MESSAGE)
2054                auto getAspectNameStr = [](sycl::aspect AspectNum) -> std::string
2055                {
2056                    switch (AspectNum)
2057                    {
2058#include <sycl/info/aspects.def>
2059#include <sycl/info/aspects_deprecated.def>
2060                    default:
2061                        return "unknown aspect";
2062                    }
2063                };
2064#undef __SYCL_ASPECT_DEPRECATED_ALIAS
2065#undef __SYCL_ASPECT_DEPRECATED
2066#undef __SYCL_ASPECT
2067                throw std::runtime_error(
2068                    "'" + getAspectNameStr(it) + "' is not supported in '" +
2069                    dev.get_info<sycl::info::device::name>() + "' device");
2070            }
2071            break;
2072        }
2073    }
2074
2075    static inline unsigned int get_current_device_id()
2076    {
2077        return dev_mgr::instance().current_device_id();
2078    }
2079
2080    static inline device_ext &get_current_device()
2081    {
2082        return dev_mgr::instance().current_device();
2083    }
2084
2085    static inline device_ext &get_device(unsigned int id)
2086    {
2087        return dev_mgr::instance().get_device(id);
2088    }
2089
2090    static inline sycl::queue &get_in_order_queue()
2091    {
2092        return dev_mgr::instance().current_device().in_order_queue();
2093    }
2094
2095    static sycl::event
2096    dpct_memcpy(sycl::queue &q, void *to_ptr, const void *from_ptr, size_t size,
2097                memcpy_direction direction,
2098                const std::vector<sycl::event> &dep_events = {})
2099    {
2100        if (!size)
2101            return sycl::event{};
2102        return q.memcpy(to_ptr, from_ptr, size, dep_events);
2103        DPCT_UNUSED(direction);
2104    }
2105
2106    // Get actual copy range and make sure it will not exceed range.
2107    static inline size_t get_copy_range(sycl::range<3> size, size_t slice,
2108                                        size_t pitch)
2109    {
2110        return slice * (size.get(2) - 1) + pitch * (size.get(1) - 1) + size.get(0);
2111    }
2112
2113    static inline size_t get_offset(sycl::id<3> id, size_t slice,
2114                                    size_t pitch)
2115    {
2116        return slice * id.get(2) + pitch * id.get(1) + id.get(0);
2117    }
2118
2119    /// copy 3D matrix specified by \p size from 3D matrix specified by \p from_ptr
2120    /// and \p from_range to another specified by \p to_ptr and \p to_range.
2121    static inline std::vector<sycl::event>
2122    dpct_memcpy(sycl::queue &q, void *to_ptr, const void *from_ptr,
2123                sycl::range<3> to_range, sycl::range<3> from_range,
2124                sycl::id<3> to_id, sycl::id<3> from_id,
2125                sycl::range<3> size, memcpy_direction direction,
2126                const std::vector<sycl::event> &dep_events = {})
2127    {
2128        // RAII for host pointer
2129        class host_buffer
2130        {
2131            void *_buf;
2132            size_t _size;
2133            sycl::queue &_q;
2134            const std::vector<sycl::event> &_deps; // free operation depends
2135
2136        public:
2137            host_buffer(size_t size, sycl::queue &q,
2138                        const std::vector<sycl::event> &deps)
2139                : _buf(std::malloc(size)), _size(size), _q(q), _deps(deps) {}
2140            void *get_ptr() const { return _buf; }
2141            size_t get_size() const { return _size; }
2142            ~host_buffer()
2143            {
2144                if (_buf)
2145                {
2146                    _q.submit([&](sycl::handler &cgh)
2147                              {
2148            cgh.depends_on(_deps);
2149            cgh.host_task([buf = _buf] { std::free(buf); }); });
2150                }
2151            }
2152        };
2153        std::vector<sycl::event> event_list;
2154
2155        size_t to_slice = to_range.get(1) * to_range.get(0),
2156               from_slice = from_range.get(1) * from_range.get(0);
2157        unsigned char *to_surface =
2158            (unsigned char *)to_ptr + get_offset(to_id, to_slice, to_range.get(0));
2159        const unsigned char *from_surface =
2160            (const unsigned char *)from_ptr +
2161            get_offset(from_id, from_slice, from_range.get(0));
2162
2163        if (to_slice == from_slice && to_slice == size.get(1) * size.get(0))
2164        {
2165            return {dpct_memcpy(q, to_surface, from_surface, to_slice * size.get(2),
2166                                direction, dep_events)};
2167        }
2168        direction = detail::deduce_memcpy_direction(q, to_ptr, from_ptr, direction);
2169        size_t size_slice = size.get(1) * size.get(0);
2170        switch (direction)
2171        {
2172        case host_to_host:
2173            for (size_t z = 0; z < size.get(2); ++z)
2174            {
2175                unsigned char *to_ptr = to_surface;
2176                const unsigned char *from_ptr = from_surface;
2177                if (to_range.get(0) == from_range.get(0) &&
2178                    to_range.get(0) == size.get(0))
2179                {
2180                    event_list.push_back(dpct_memcpy(q, to_ptr, from_ptr, size_slice,
2181                                                     direction, dep_events));
2182                }
2183                else
2184                {
2185                    for (size_t y = 0; y < size.get(1); ++y)
2186                    {
2187                        event_list.push_back(dpct_memcpy(q, to_ptr, from_ptr, size.get(0),
2188                                                         direction, dep_events));
2189                        to_ptr += to_range.get(0);
2190                        from_ptr += from_range.get(0);
2191                    }
2192                }
2193                to_surface += to_slice;
2194                from_surface += from_slice;
2195            }
2196            break;
2197        case host_to_device:
2198        {
2199            host_buffer buf(get_copy_range(size, to_slice, to_range.get(0)), q,
2200                            event_list);
2201            std::vector<sycl::event> host_events;
2202            if (to_slice == size_slice)
2203            {
2204                // Copy host data to a temp host buffer with the shape of target.
2205                host_events =
2206                    dpct_memcpy(q, buf.get_ptr(), from_surface, to_range, from_range,
2207                                sycl::id<3>(0, 0, 0), sycl::id<3>(0, 0, 0), size,
2208                                host_to_host, dep_events);
2209            }
2210            else
2211            {
2212                // Copy host data to a temp host buffer with the shape of target.
2213                host_events = dpct_memcpy(
2214                    q, buf.get_ptr(), from_surface, to_range, from_range,
2215                    sycl::id<3>(0, 0, 0), sycl::id<3>(0, 0, 0), size, host_to_host,
2216                    // If has padding data, not sure whether it is useless. So fill temp
2217                    // buffer with it.
2218                    std::vector<sycl::event>{
2219                        dpct_memcpy(q, buf.get_ptr(), to_surface, buf.get_size(),
2220                                    device_to_host, dep_events)});
2221            }
2222            // Copy from temp host buffer to device with only one submit.
2223            event_list.push_back(dpct_memcpy(q, to_surface, buf.get_ptr(),
2224                                             buf.get_size(), host_to_device,
2225                                             host_events));
2226            break;
2227        }
2228        case device_to_host:
2229        {
2230            host_buffer buf(get_copy_range(size, from_slice, from_range.get(0)), q,
2231                            event_list);
2232            // Copy from host temp buffer to host target with reshaping.
2233            event_list = dpct_memcpy(
2234                q, to_surface, buf.get_ptr(), to_range, from_range, sycl::id<3>(0, 0, 0),
2235                sycl::id<3>(0, 0, 0), size, host_to_host,
2236                // Copy from device to temp host buffer with only one submit.
2237                std::vector<sycl::event>{dpct_memcpy(q, buf.get_ptr(), from_surface,
2238                                                     buf.get_size(),
2239                                                     device_to_host, dep_events)});
2240            break;
2241        }
2242        case device_to_device:
2243            event_list.push_back(q.submit([&](sycl::handler &cgh)
2244                                          {
2245        cgh.depends_on(dep_events);
2246        cgh.parallel_for<class dpct_memcpy_3d_detail>(
2247            size,
2248            [=](sycl::id<3> id) {
2249                to_surface[get_offset(id, to_slice, to_range.get(0))] =
2250                    from_surface[get_offset(id, from_slice, from_range.get(0))];
2251            }); }));
2252        break;
2253        default:
2254            throw std::runtime_error("dpct_memcpy: invalid direction value");
2255        }
2256        return event_list;
2257    }
2258
2259    /// memcpy 2D/3D matrix specified by pitched_data.
2260    static inline std::vector<sycl::event>
2261    dpct_memcpy(sycl::queue &q, pitched_data to, sycl::id<3> to_id,
2262                pitched_data from, sycl::id<3> from_id, sycl::range<3> size,
2263                memcpy_direction direction = automatic)
2264    {
2265        return dpct_memcpy(q, to.get_data_ptr(), from.get_data_ptr(),
2266                           sycl::range<3>(to.get_pitch(), to.get_y(), 1),
2267                           sycl::range<3>(from.get_pitch(), from.get_y(), 1), to_id, from_id,
2268                           size, direction);
2269    }
2270
2271    /// memcpy 2D matrix with pitch.
2272    static inline std::vector<sycl::event>
2273    dpct_memcpy(sycl::queue &q, void *to_ptr, const void *from_ptr,
2274                size_t to_pitch, size_t from_pitch, size_t x, size_t y,
2275                memcpy_direction direction = automatic)
2276    {
2277        return dpct_memcpy(q, to_ptr, from_ptr, sycl::range<3>(to_pitch, y, 1),
2278                           sycl::range<3>(from_pitch, y, 1),
2279                           sycl::id<3>(0, 0, 0), sycl::id<3>(0, 0, 0),
2280                           sycl::range<3>(x, y, 1), direction);
2281    }
2282
2283    inline void gemm(sycl::queue & q, oneapi::mkl::transpose a_trans, oneapi::mkl::transpose b_trans, int m, int n,
2284                     int k, const void * alpha, const void * a, library_data_t a_type, int lda, const void * b,
2285                     library_data_t b_type, int ldb, const void * beta, void * c, library_data_t c_type, int ldc,
2286                     library_data_t scaling_type) {
2287        if (scaling_type == library_data_t::real_float &&
2288            c_type == library_data_t::complex_float)
2289        {
2290            scaling_type = library_data_t::complex_float;
2291        }
2292        else if (scaling_type == library_data_t::real_double &&
2293                 c_type == library_data_t::complex_double)
2294        {
2295            scaling_type = library_data_t::complex_double;
2296        }
2297
2298        std::uint64_t key =
2299            detail::get_type_combination_id(a_type, b_type, c_type, scaling_type);
2300        switch (key)
2301        {
2302        case detail::get_type_combination_id(
2303            library_data_t::real_float, library_data_t::real_float,
2304            library_data_t::real_float, library_data_t::real_float):
2305        {
2306            detail::gemm_impl<float, float, float, float>(
2307                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
2308            break;
2309        }
2310        case detail::get_type_combination_id(
2311            library_data_t::real_double, library_data_t::real_double,
2312            library_data_t::real_double, library_data_t::real_double):
2313        {
2314            detail::gemm_impl<double, double, double, double>(
2315                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
2316            break;
2317        }
2318        case detail::get_type_combination_id(
2319            library_data_t::complex_float, library_data_t::complex_float,
2320            library_data_t::complex_float, library_data_t::complex_float):
2321        {
2322            detail::gemm_impl<std::complex<float>, std::complex<float>,
2323                              std::complex<float>, std::complex<float>>(
2324                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
2325            break;
2326        }
2327        case detail::get_type_combination_id(
2328            library_data_t::complex_double, library_data_t::complex_double,
2329            library_data_t::complex_double, library_data_t::complex_double):
2330        {
2331            detail::gemm_impl<std::complex<double>, std::complex<double>,
2332                              std::complex<double>, std::complex<double>>(
2333                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
2334            break;
2335        }
2336        case detail::get_type_combination_id(
2337            library_data_t::real_half, library_data_t::real_half,
2338            library_data_t::real_half, library_data_t::real_half):
2339        {
2340            detail::gemm_impl<sycl::half, sycl::half, sycl::half,
2341                              sycl::half>(q, a_trans, b_trans, m, n, k, alpha, a,
2342                                          lda, b, ldb, beta, c, ldc);
2343            break;
2344        }
2345#ifdef __INTEL_MKL__
2346        case detail::get_type_combination_id(
2347            library_data_t::real_bfloat16, library_data_t::real_bfloat16,
2348            library_data_t::real_float, library_data_t::real_float):
2349        {
2350            detail::gemm_impl<oneapi::mkl::bfloat16, oneapi::mkl::bfloat16, float, float>(
2351                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
2352            break;
2353        }
2354        case detail::get_type_combination_id(
2355            library_data_t::real_half, library_data_t::real_half,
2356            library_data_t::real_float, library_data_t::real_float):
2357        {
2358            detail::gemm_impl<sycl::half, sycl::half, float, float>(
2359                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
2360            break;
2361        }
2362        case detail::get_type_combination_id(
2363            library_data_t::real_half, library_data_t::real_half,
2364            library_data_t::real_half, library_data_t::real_float):
2365        {
2366            float alpha_value =
2367                dpct::get_value(reinterpret_cast<const float *>(alpha), q);
2368            float beta_value =
2369                dpct::get_value(reinterpret_cast<const float *>(beta), q);
2370            sycl::half alpha_half(alpha_value);
2371            sycl::half beta_half(beta_value);
2372            detail::gemm_impl<sycl::half, sycl::half, sycl::half,
2373                              sycl::half>(q, a_trans, b_trans, m, n, k, &alpha_half,
2374                                          a, lda, b, ldb, &beta_half, c, ldc);
2375            break;
2376        }
2377        case detail::get_type_combination_id(
2378            library_data_t::real_int8, library_data_t::real_int8,
2379            library_data_t::real_float, library_data_t::real_float):
2380        {
2381            detail::gemm_impl<std::int8_t, std::int8_t, float, float>(
2382                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
2383            break;
2384        }
2385        case detail::get_type_combination_id(
2386            library_data_t::real_bfloat16, library_data_t::real_bfloat16,
2387            library_data_t::real_bfloat16, library_data_t::real_float):
2388        {
2389            detail::gemm_impl<oneapi::mkl::bfloat16, oneapi::mkl::bfloat16, oneapi::mkl::bfloat16, float>(
2390                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
2391            break;
2392        }
2393        case detail::get_type_combination_id(
2394            library_data_t::real_int8, library_data_t::real_int8,
2395            library_data_t::real_int32, library_data_t::real_int32):
2396        {
2397            float alpha_float =
2398                dpct::get_value(reinterpret_cast<const std::int32_t *>(alpha), q);
2399            float beta_float =
2400                dpct::get_value(reinterpret_cast<const std::int32_t *>(beta), q);
2401            detail::gemm_impl<std::int8_t, std::int8_t, std::int32_t, float>(
2402                q, a_trans, b_trans, m, n, k, &alpha_float, a, lda, b, ldb, &beta_float, c, ldc);
2403            break;
2404        }
2405#endif // __INTEL_MKL__
2406        default:
2407            throw std::runtime_error("the combination of data type is unsupported");
2408        }
2409    }  // gemm()
2410
2411    /// Computes a batch of matrix-matrix product with general matrices.
2412    /// \param [in] q The queue where the routine should be executed.
2413    /// \param [in] a_trans Specifies the operation applied to A.
2414    /// \param [in] b_trans Specifies the operation applied to B.
2415    /// \param [in] m Specifies the number of rows of the matrix op(A) and of the matrix C.
2416    /// \param [in] n Specifies the number of columns of the matrix op(B) and of the matrix C.
2417    /// \param [in] k Specifies the number of columns of the matrix op(A) and the number of rows of the matrix op(B).
2418    /// \param [in] alpha Scaling factor for the matrix-matrix product.
2419    /// \param [in] a Input matrix A.
2420    /// \param [in] a_type Data type of the matrix A.
2421    /// \param [in] lda Leading dimension of A.
2422    /// \param [in] b Input matrix B.
2423    /// \param [in] b_type Data type of the matrix B.
2424    /// \param [in] ldb Leading dimension of B.
2425    /// \param [in] beta Scaling factor for matrix C.
2426    /// \param [in, out] c Input/Output matrix C.
2427    /// \param [in] c_type Data type of the matrix C.
2428    /// \param [in] ldc Leading dimension of C.
2429    /// \param [in] batch_size Specifies the number of matrix multiply operations to perform.
2430    /// \param [in] scaling_type Data type of the scaling factors.
2431    inline void gemm_batch(sycl::queue & q, oneapi::mkl::transpose a_trans, oneapi::mkl::transpose b_trans, int m,
2432                           int n, int k, const void * alpha, const void * a[], library_data_t a_type, int lda,
2433                           const void * b[], library_data_t b_type, int ldb, const void * beta, void * c[],
2434                           library_data_t c_type, int ldc, int batch_size, library_data_t scaling_type,
2435                           matrix_info_t<float> * matrix_info) {
2436        std::uint64_t key =
2437            detail::get_type_combination_id(a_type, b_type, c_type, scaling_type);
2438        switch (key)
2439        {
2440        case detail::get_type_combination_id(
2441            library_data_t::real_float, library_data_t::real_float,
2442            library_data_t::real_float, library_data_t::real_float):
2443        {
2444            detail::gemm_batch_impl<float, float, float, float>(q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb,
2445                                                                beta, c, ldc, batch_size, matrix_info);
2446            break;
2447        }
2448        case detail::get_type_combination_id(
2449            library_data_t::real_double, library_data_t::real_double,
2450            library_data_t::real_double, library_data_t::real_double):
2451        {
2452            detail::gemm_batch_impl<double, double, double, double>(q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb,
2453                                                                    beta, c, ldc, batch_size, matrix_info);
2454            break;
2455        }
2456        case detail::get_type_combination_id(
2457            library_data_t::real_half, library_data_t::real_half,
2458            library_data_t::real_half, library_data_t::real_half):
2459        {
2460            detail::gemm_batch_impl<sycl::half, sycl::half, sycl::half, sycl::half>(
2461                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc, batch_size, matrix_info);
2462            break;
2463        }
2464#ifdef __INTEL_MKL__
2465        case detail::get_type_combination_id(
2466            library_data_t::real_bfloat16, library_data_t::real_bfloat16,
2467            library_data_t::real_bfloat16, library_data_t::real_float):
2468        {
2469            detail::gemm_batch_impl<oneapi::mkl::bfloat16, oneapi::mkl::bfloat16, oneapi::mkl::bfloat16, float>(
2470                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc, batch_size, matrix_info);
2471            break;
2472        }
2473        case detail::get_type_combination_id(
2474            library_data_t::real_bfloat16, library_data_t::real_bfloat16,
2475            library_data_t::real_float, library_data_t::real_float):
2476        {
2477            detail::gemm_batch_impl<oneapi::mkl::bfloat16, oneapi::mkl::bfloat16, float, float>(
2478                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc, batch_size, matrix_info);
2479            break;
2480        }
2481#endif
2482        case detail::get_type_combination_id(
2483            library_data_t::real_int8, library_data_t::real_int8,
2484            library_data_t::real_int32, library_data_t::real_int32):
2485        {
2486            float alpha_float =
2487                dpct::get_value(reinterpret_cast<const std::int32_t *>(alpha), q);
2488            float beta_float =
2489                dpct::get_value(reinterpret_cast<const std::int32_t *>(beta), q);
2490            detail::gemm_batch_impl<std::int8_t, std::int8_t, std::int32_t, float>(
2491                q, a_trans, b_trans, m, n, k, &alpha_float, a, lda, b, ldb, &beta_float, c, ldc, batch_size,
2492                matrix_info);
2493            break;
2494        }
2495        case detail::get_type_combination_id(
2496            library_data_t::real_int8, library_data_t::real_int8,
2497            library_data_t::real_float, library_data_t::real_float):
2498        {
2499            detail::gemm_batch_impl<std::int8_t, std::int8_t, float, float>(
2500                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc, batch_size, matrix_info);
2501            break;
2502        }
2503        case detail::get_type_combination_id(
2504            library_data_t::real_half, library_data_t::real_half,
2505            library_data_t::real_float, library_data_t::real_float):
2506        {
2507            detail::gemm_batch_impl<sycl::half, sycl::half, float, float>(
2508                q, a_trans, b_trans, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc, batch_size, matrix_info);
2509            break;
2510        }
2511        case detail::get_type_combination_id(
2512            library_data_t::real_half, library_data_t::real_half,
2513            library_data_t::real_half, library_data_t::real_float):
2514        {
2515            float alpha_value =
2516                dpct::get_value(reinterpret_cast<const float *>(alpha), q);
2517            float beta_value =
2518                dpct::get_value(reinterpret_cast<const float *>(beta), q);
2519            sycl::half alpha_half(alpha_value);
2520            sycl::half beta_half(beta_value);
2521            detail::gemm_batch_impl<sycl::half, sycl::half, sycl::half, sycl::half>(
2522                q, a_trans, b_trans, m, n, k, &alpha_half, a, lda, b, ldb, &beta_half, c, ldc, batch_size, matrix_info);
2523            break;
2524        }
2525        default:
2526            throw std::runtime_error("the combination of data type is unsupported");
2527        }
2528    }
2529
2530    /// Computes a batch of matrix-matrix product with general matrices.
2531    /// \param [in] q The queue where the routine should be executed.
2532    /// \param [in] a_trans Specifies the operation applied to A.
2533    /// \param [in] b_trans Specifies the operation applied to B.
2534    /// \param [in] m Specifies the number of rows of the matrix op(A) and of the matrix C.
2535    /// \param [in] n Specifies the number of columns of the matrix op(B) and of the matrix C.
2536    /// \param [in] k Specifies the number of columns of the matrix op(A) and the number of rows of the matrix op(B).
2537    /// \param [in] alpha Scaling factor for the matrix-matrix product.
2538    /// \param [in] a Input matrix A.
2539    /// \param [in] a_type Data type of the matrix A.
2540    /// \param [in] lda Leading dimension of A.
2541    /// \param [in] stride_a Stride between the different A matrices.
2542    /// \param [in] b Input matrix B.
2543    /// \param [in] b_type Data type of the matrix B.
2544    /// \param [in] ldb Leading dimension of B.
2545    /// \param [in] stride_b Stride between the different B matrices.
2546    /// \param [in] beta Scaling factor for matrix C.
2547    /// \param [in, out] c Input/Output matrix C.
2548    /// \param [in] c_type Data type of the matrix C.
2549    /// \param [in] ldc Leading dimension of C.
2550    /// \param [in] stride_c Stride between the different C matrices.
2551    /// \param [in] batch_size Specifies the number of matrix multiply operations to perform.
2552    /// \param [in] scaling_type Data type of the scaling factors.
2553    inline void gemm_batch(sycl::queue & q, oneapi::mkl::transpose a_trans, oneapi::mkl::transpose b_trans, int m,
2554                           int n, int k, const void * alpha, const void * a, library_data_t a_type, int lda,
2555                           long long int stride_a, const void * b, library_data_t b_type, int ldb,
2556                           long long int stride_b, const void * beta, void * c, library_data_t c_type, int ldc,
2557                           long long int stride_c, int batch_size, library_data_t scaling_type) {
2558        if (scaling_type == library_data_t::real_float &&
2559            c_type == library_data_t::complex_float)
2560        {
2561            scaling_type = library_data_t::complex_float;
2562        }
2563        else if (scaling_type == library_data_t::real_double &&
2564                 c_type == library_data_t::complex_double)
2565        {
2566            scaling_type = library_data_t::complex_double;
2567        }
2568
2569        std::uint64_t key =
2570            detail::get_type_combination_id(a_type, b_type, c_type, scaling_type);
2571        switch (key)
2572        {
2573        case detail::get_type_combination_id(
2574            library_data_t::real_float, library_data_t::real_float,
2575            library_data_t::real_float, library_data_t::real_float):
2576        {
2577            detail::gemm_batch_impl<float, float, float, float>(
2578                q, a_trans, b_trans, m, n, k, alpha, a, lda, stride_a, b, ldb, stride_b,
2579                beta, c, ldc, stride_c, batch_size);
2580            break;
2581        }
2582        case detail::get_type_combination_id(
2583            library_data_t::real_double, library_data_t::real_double,
2584            library_data_t::real_double, library_data_t::real_double):
2585        {
2586            detail::gemm_batch_impl<double, double, double, double>(
2587                q, a_trans, b_trans, m, n, k, alpha, a, lda, stride_a, b, ldb, stride_b,
2588                beta, c, ldc, stride_c, batch_size);
2589            break;
2590        }
2591        case detail::get_type_combination_id(
2592            library_data_t::complex_float, library_data_t::complex_float,
2593            library_data_t::complex_float, library_data_t::complex_float):
2594        {
2595            detail::gemm_batch_impl<std::complex<float>, std::complex<float>,
2596                                    std::complex<float>, std::complex<float>>(
2597                q, a_trans, b_trans, m, n, k, alpha, a, lda, stride_a, b, ldb, stride_b,
2598                beta, c, ldc, stride_c, batch_size);
2599            break;
2600        }
2601        case detail::get_type_combination_id(
2602            library_data_t::complex_double, library_data_t::complex_double,
2603            library_data_t::complex_double, library_data_t::complex_double):
2604        {
2605            detail::gemm_batch_impl<std::complex<double>, std::complex<double>,
2606                                    std::complex<double>, std::complex<double>>(
2607                q, a_trans, b_trans, m, n, k, alpha, a, lda, stride_a, b, ldb, stride_b,
2608                beta, c, ldc, stride_c, batch_size);
2609            break;
2610        }
2611        case detail::get_type_combination_id(
2612            library_data_t::real_half, library_data_t::real_half,
2613            library_data_t::real_half, library_data_t::real_half):
2614        {
2615            detail::gemm_batch_impl<sycl::half, sycl::half, sycl::half,
2616                                    sycl::half>(q, a_trans, b_trans, m, n, k, alpha,
2617                                                a, lda, stride_a, b, ldb, stride_b,
2618                                                beta, c, ldc, stride_c, batch_size);
2619            break;
2620        }
2621#ifdef __INTEL_MKL__
2622        case detail::get_type_combination_id(
2623            library_data_t::real_bfloat16, library_data_t::real_bfloat16,
2624            library_data_t::real_bfloat16, library_data_t::real_float):
2625        {
2626            detail::gemm_batch_impl<oneapi::mkl::bfloat16, oneapi::mkl::bfloat16, oneapi::mkl::bfloat16, float>(
2627                q, a_trans, b_trans, m, n, k, alpha, a, lda, stride_a, b, ldb, stride_b, beta, c, ldc, stride_c,
2628                batch_size);
2629            break;
2630        }
2631        case detail::get_type_combination_id(
2632            library_data_t::real_bfloat16, library_data_t::real_bfloat16,
2633            library_data_t::real_float, library_data_t::real_float):
2634        {
2635            detail::gemm_batch_impl<oneapi::mkl::bfloat16, oneapi::mkl::bfloat16, float, float>(
2636                q, a_trans, b_trans, m, n, k, alpha, a, lda, stride_a, b, ldb, stride_b, beta, c, ldc, stride_c,
2637                batch_size);
2638            break;
2639        }
2640#endif
2641        case detail::get_type_combination_id(
2642            library_data_t::real_int8, library_data_t::real_int8,
2643            library_data_t::real_int32, library_data_t::real_int32):
2644        {
2645            detail::gemm_batch_impl<std::int8_t, std::int8_t, std::int32_t,
2646                                    std::int32_t>(q, a_trans, b_trans, m, n, k, alpha,
2647                                                  a, lda, stride_a, b, ldb, stride_b,
2648                                                  beta, c, ldc, stride_c, batch_size);
2649            break;
2650        }
2651        case detail::get_type_combination_id(
2652            library_data_t::real_int8, library_data_t::real_int8,
2653            library_data_t::real_float, library_data_t::real_float):
2654        {
2655            detail::gemm_batch_impl<std::int8_t, std::int8_t, float, float>(
2656                q, a_trans, b_trans, m, n, k, alpha, a, lda, stride_a, b, ldb, stride_b,
2657                beta, c, ldc, stride_c, batch_size);
2658            break;
2659        }
2660        case detail::get_type_combination_id(
2661            library_data_t::real_half, library_data_t::real_half,
2662            library_data_t::real_float, library_data_t::real_float):
2663        {
2664            detail::gemm_batch_impl<sycl::half, sycl::half, float, float>(
2665                q, a_trans, b_trans, m, n, k, alpha, a, lda, stride_a, b, ldb, stride_b,
2666                beta, c, ldc, stride_c, batch_size);
2667            break;
2668        }
2669        case detail::get_type_combination_id(
2670            library_data_t::real_half, library_data_t::real_half,
2671            library_data_t::real_half, library_data_t::real_float):
2672        {
2673            float alpha_value =
2674                dpct::get_value(reinterpret_cast<const float *>(alpha), q);
2675            float beta_value =
2676                dpct::get_value(reinterpret_cast<const float *>(beta), q);
2677            sycl::half alpha_half(alpha_value);
2678            sycl::half beta_half(beta_value);
2679            detail::gemm_batch_impl<sycl::half, sycl::half, sycl::half, sycl::half>(
2680                q, a_trans, b_trans, m, n, k, &alpha_half, a, lda, stride_a, b, ldb, stride_b,
2681                &beta_half, c, ldc, stride_c, batch_size);
2682            break;
2683        }
2684        default:
2685            throw std::runtime_error("the combination of data type is unsupported");
2686        }
2687    }
2688
2689    static inline void
2690    async_dpct_memcpy(void *to_ptr, size_t to_pitch, const void *from_ptr,
2691                      size_t from_pitch, size_t x, size_t y,
2692                      memcpy_direction direction = automatic,
2693                      sycl::queue &q = get_default_queue())
2694    {
2695        detail::dpct_memcpy(q, to_ptr, from_ptr, to_pitch, from_pitch, x, y,
2696                            direction);
2697    }
2698
2699    using err0 = detail::generic_error_type<struct err0_tag, int>;
2700    using err1 = detail::generic_error_type<struct err1_tag, int>;
2701
2702    static inline void dpct_free(void *ptr, sycl::queue &q = get_default_queue()) {
2703        detail::dpct_free(ptr, q);
2704    }
2705
2706    /// dpct accessor used as device function parameter.
2707    template <class T, memory_region Memory, size_t Dimension> class accessor;
2708    template <class T, memory_region Memory> class accessor<T, Memory, 3> {
2709    public:
2710        using memory_t = detail::memory_traits<Memory, T>;
2711        using element_t = typename memory_t::element_t;
2712        using pointer_t = typename memory_t::pointer_t;
2713        using accessor_t = typename memory_t::template accessor_t<3>;
2714        accessor(pointer_t data, const sycl::range<3> &in_range)
2715            : _data(data), _range(in_range) {}
2716        template <memory_region M = Memory>
2717        accessor(typename std::enable_if<M != local, const accessor_t>::type &acc)
2718            : accessor(acc, acc.get_range()) {}
2719        accessor(const accessor_t &acc, const sycl::range<3> &in_range)
2720            : accessor(acc.get_pointer(), in_range) {}
2721        accessor<T, Memory, 2> operator[](size_t index) const {
2722            sycl::range<2> sub(_range.get(1), _range.get(2));
2723            return accessor<T, Memory, 2>(_data + index * sub.size(), sub);
2724        }
2725
2726        pointer_t get_ptr() const { return _data; }
2727
2728    private:
2729        pointer_t _data;
2730        sycl::range<3> _range;
2731    };
2732    template <class T, memory_region Memory> class accessor<T, Memory, 2> {
2733    public:
2734        using memory_t = detail::memory_traits<Memory, T>;
2735        using element_t = typename memory_t::element_t;
2736        using pointer_t = typename memory_t::pointer_t;
2737        using accessor_t = typename memory_t::template accessor_t<2>;
2738        accessor(pointer_t data, const sycl::range<2> &in_range)
2739            : _data(data), _range(in_range) {}
2740        template <memory_region M = Memory>
2741        accessor(typename std::enable_if<M != local, const accessor_t>::type &acc)
2742            : accessor(acc, acc.get_range()) {}
2743        accessor(const accessor_t &acc, const sycl::range<2> &in_range)
2744            : accessor(acc.get_pointer(), in_range) {}
2745
2746        pointer_t operator[](size_t index) const {
2747            return _data + _range.get(1) * index;
2748        }
2749
2750        pointer_t get_ptr() const { return _data; }
2751
2752    private:
2753        pointer_t _data;
2754        sycl::range<2> _range;
2755    };
2756
2757    namespace detail {
2758        /// Device variable with address space of shared, global or constant.
2759        template <class T, memory_region Memory, size_t Dimension> class device_memory {
2760        public:
2761            using accessor_t =
2762                typename detail::memory_traits<Memory,
2763                                            T>::template accessor_t<Dimension>;
2764            using value_t = typename detail::memory_traits<Memory, T>::value_t;
2765            using dpct_accessor_t = dpct::accessor<T, Memory, Dimension>;
2766
2767            device_memory() : device_memory(sycl::range<Dimension>(1)) {}
2768
2769            /// Constructor of 1-D array with initializer list
2770            device_memory(const sycl::range<Dimension> &in_range,
2771                        std::initializer_list<value_t> &&init_list)
2772                : device_memory(in_range) {
2773                assert(init_list.size() <= in_range.size());
2774                _host_ptr = (value_t *)std::malloc(_size);
2775                std::memset(_host_ptr, 0, _size);
2776                std::memcpy(_host_ptr, init_list.begin(), init_list.size() * sizeof(T));
2777            }
2778
2779            /// Constructor of 2-D array with initializer list
2780            template <size_t D = Dimension>
2781            device_memory(
2782                const typename std::enable_if<D == 2, sycl::range<2>>::type &in_range,
2783                std::initializer_list<std::initializer_list<value_t>> &&init_list)
2784                : device_memory(in_range) {
2785                assert(init_list.size() <= in_range[0]);
2786                _host_ptr = (value_t *)std::malloc(_size);
2787                std::memset(_host_ptr, 0, _size);
2788                auto tmp_data = _host_ptr;
2789                for (auto sub_list : init_list) {
2790                    assert(sub_list.size() <= in_range[1]);
2791                    std::memcpy(tmp_data, sub_list.begin(),
2792                                sub_list.size() * sizeof(T));
2793                    tmp_data += in_range[1];
2794                }
2795            }
2796
2797            /// Constructor with range
2798            device_memory(const sycl::range<Dimension> &range_in)
2799                : _size(range_in.size() * sizeof(T)), _range(range_in),
2800                _reference(false), _host_ptr(nullptr), _device_ptr(nullptr) {
2801                static_assert(
2802                    (Memory == global) || (Memory == constant) || (Memory == shared),
2803                    "device memory region should be global, constant or shared");
2804                // Make sure that singleton class mem_mgr and dev_mgr will destruct
2805                // later than this.
2806                detail::mem_mgr::instance();
2807                dev_mgr::instance();
2808            }
2809
2810            /// Constructor with range
2811            template <class... Args>
2812            device_memory(Args... Arguments)
2813                : device_memory(sycl::range<Dimension>(Arguments...)) {}
2814
2815            ~device_memory() {
2816                if (_device_ptr && !_reference)
2817                    dpct::dpct_free(_device_ptr);
2818                if (_host_ptr)
2819                    std::free(_host_ptr);
2820            }
2821
2822            /// Allocate memory with default queue, and init memory if has initial
2823            /// value.
2824            void init() { init(dpct::get_default_queue()); }
2825            /// Allocate memory with specified queue, and init memory if has initial
2826            /// value.
2827            void init(sycl::queue &q) {
2828                if (_device_ptr)
2829                    return;
2830                if (!_size)
2831                    return;
2832                allocate_device(q);
2833                if (_host_ptr)
2834                    detail::dpct_memcpy(q, _device_ptr, _host_ptr, _size,
2835                                        host_to_device);
2836            }
2837
2838            /// The variable is assigned to a device pointer.
2839            void assign(value_t *src, size_t size) {
2840                this->~device_memory();
2841                new (this) device_memory(src, size);
2842            }
2843
2844            /// Get memory pointer of the memory object, which is virtual pointer when
2845            /// usm is not used, and device pointer when usm is used.
2846            value_t *get_ptr() { return get_ptr(get_default_queue()); }
2847            /// Get memory pointer of the memory object, which is virtual pointer when
2848            /// usm is not used, and device pointer when usm is used.
2849            value_t *get_ptr(sycl::queue &q) {
2850                init(q);
2851                return _device_ptr;
2852            }
2853
2854            /// Get the device memory object size in bytes.
2855            size_t get_size() { return _size; }
2856
2857            template <size_t D = Dimension>
2858            typename std::enable_if<D == 1, T>::type &operator[](size_t index) {
2859                init();
2860                return _device_ptr[index];
2861            }
2862
2863            /// Get dpct::accessor with dimension info for the device memory object
2864            /// when usm is used and dimension is greater than 1.
2865            template <size_t D = Dimension>
2866            typename std::enable_if<D != 1, dpct_accessor_t>::type
2867            get_access([[maybe_unused]] sycl::handler &cgh) {
2868                return dpct_accessor_t((T *)_device_ptr, _range);
2869            }
2870
2871        private:
2872            device_memory(value_t *memory_ptr, size_t size)
2873                : _size(size), _range(size / sizeof(T)), _reference(true),
2874                _device_ptr(memory_ptr) {}
2875
2876            void allocate_device(sycl::queue &q) {
2877        #ifndef DPCT_USM_LEVEL_NONE
2878                if (Memory == shared) {
2879                    _device_ptr = (value_t *)sycl::malloc_shared(_size, q.get_device(),
2880                                                                q.get_context());
2881                    return;
2882                }
2883        #ifdef SYCL_EXT_ONEAPI_USM_DEVICE_READ_ONLY
2884                if (Memory == constant) {
2885                    _device_ptr = (value_t *)sycl::malloc_device(
2886                        _size, q.get_device(), q.get_context(),
2887                        sycl::ext::oneapi::property::usm::device_read_only());
2888                    return;
2889                }
2890        #endif
2891        #endif
2892                _device_ptr = (value_t *)detail::dpct_malloc(_size, q);
2893            }
2894
2895            size_t _size;
2896            sycl::range<Dimension> _range;
2897            bool _reference;
2898            value_t *_host_ptr;
2899            value_t *_device_ptr;
2900        };
2901        template <class T, memory_region Memory>
2902        class device_memory<T, Memory, 0> : public device_memory<T, Memory, 1> {
2903        public:
2904            using base = device_memory<T, Memory, 1>;
2905            using value_t = typename base::value_t;
2906            using accessor_t =
2907                typename detail::memory_traits<Memory, T>::template accessor_t<0>;
2908
2909            /// Constructor with initial value.
2910            device_memory(const value_t &val) : base(sycl::range<1>(1), {val}) {}
2911
2912            /// Default constructor
2913            device_memory() : base(1) {}
2914        };
2915        } // namespace detail
2916
2917    template <class T, size_t Dimension>
2918    using global_memory = detail::device_memory<T, global, Dimension>;
2919    template <class T, size_t Dimension>
2920    using constant_memory = detail::device_memory<T, constant, Dimension>;
2921    template <class T, size_t Dimension>
2922    using shared_memory = detail::device_memory<T, shared, Dimension>;
2923
2924
2925    template <typename T,
2926            sycl::access::address_space addressSpace =
2927                sycl::access::address_space::global_space,
2928            sycl::memory_order memoryOrder = sycl::memory_order::relaxed,
2929            sycl::memory_scope memoryScope = sycl::memory_scope::device>
2930    inline T atomic_fetch_add(T *addr, T operand) {
2931    auto atm =
2932        sycl::atomic_ref<T, memoryOrder, memoryScope, addressSpace>(addr[0]);
2933    return atm.fetch_add(operand);
2934    }
2935
2936    template <sycl::access::address_space addressSpace =
2937                sycl::access::address_space::global_space,
2938            sycl::memory_order memoryOrder = sycl::memory_order::relaxed,
2939            sycl::memory_scope memoryScope = sycl::memory_scope::device,
2940            typename T1, typename T2>
2941    inline T1 atomic_fetch_add(T1 *addr, T2 operand) {
2942    auto atm =
2943        sycl::atomic_ref<T1, memoryOrder, memoryScope, addressSpace>(addr[0]);
2944    return atm.fetch_add(operand);
2945    }
2946
2947    template <typename T, sycl::access::address_space addressSpace =
2948                            sycl::access::address_space::global_space>
2949    inline T atomic_fetch_add(T *addr, T operand,
2950                            sycl::memory_order memoryOrder) {
2951    switch (memoryOrder) {
2952        case sycl::memory_order::relaxed:
2953            return atomic_fetch_add<T, addressSpace, sycl::memory_order::relaxed,
2954                                    sycl::memory_scope::device>(addr, operand);
2955        case sycl::memory_order::acq_rel:
2956            return atomic_fetch_add<T, addressSpace, sycl::memory_order::acq_rel,
2957                                    sycl::memory_scope::device>(addr, operand);
2958        case sycl::memory_order::seq_cst:
2959            return atomic_fetch_add<T, addressSpace, sycl::memory_order::seq_cst,
2960                                    sycl::memory_scope::device>(addr, operand);
2961        default:
2962            assert(false && "Invalid memory_order for atomics. Valid memory_order for "
2963                            "atomics are: sycl::memory_order::relaxed, "
2964                            "sycl::memory_order::acq_rel, sycl::memory_order::seq_cst!");
2965        }
2966    }
2967
2968    template <sycl::access::address_space addressSpace =
2969                sycl::access::address_space::global_space,
2970            typename T1, typename T2>
2971    inline T1 atomic_fetch_add(T1 *addr, T2 operand,
2972                            sycl::memory_order memoryOrder) {
2973    atomic_fetch_add<T1, addressSpace>(addr, operand, memoryOrder);
2974    }
2975
2976    inline unsigned int byte_level_permute(
2977        unsigned int a, unsigned int b, unsigned int s) {
2978      unsigned int ret;
2979      ret = ((((std::uint64_t)b << 32 | a) >> (s & 0x7) * 8) & 0xff) |
2980            (((((std::uint64_t)b << 32 | a) >> ((s >> 4) & 0x7) * 8) & 0xff)
2981             << 8) |
2982            (((((std::uint64_t)b << 32 | a) >> ((s >> 8) & 0x7) * 8) & 0xff)
2983             << 16) |
2984            (((((std::uint64_t)b << 32 | a) >> ((s >> 12) & 0x7) * 8) & 0xff)
2985             << 24);
2986      return ret;
2987    }
2988
2989    inline uint32_t byte_level_permute_custom(
2990        uint32_t low32, uint32_t high32, uint32_t sel, int mode = 0) {
2991      constexpr uint16_t lookup[6][4] = {
2992          {0x3210, 0x4321, 0x5432, 0x6543},  // Forward 4-byte extract
2993          {0x5670, 0x6701, 0x7012, 0x0123},  // Backward 4-byte extract
2994          {0x0000, 0x1111, 0x2222, 0x3333},  // Replicate 8-bit values
2995          {0x3210, 0x3211, 0x3222, 0x3333},  // Edge clamp left
2996          {0x0000, 0x1110, 0x2210, 0x3210},  // Edge clamp right
2997          {0x1010, 0x3232, 0x1010, 0x3232}   // Replicate 16-bit values
2998      };
2999
3000      if (mode >= 1 && mode <= 6) {
3001        return byte_level_permute(low32, high32, lookup[mode - 1][sel & 0x3]);
3002      } else if (!mode) {
3003        return byte_level_permute(low32, high32, sel);
3004      }
3005      return 0;
3006    }
3007
3008    template <int n_nondefault_params, int n_default_params, typename T>
3009    class args_selector;
3010
3011    /// args_selector is a helper class for extracting arguments from an
3012    /// array of pointers to arguments or buffer of arguments to pass to a
3013    /// kernel function.
3014    ///
3015    /// \param R(Ts...) The type of the kernel
3016    /// \param n_nondefault_params The number of nondefault parameters of the
3017    /// kernel (excluding parameters that like sycl::nd_item, etc.) \param
3018    /// n_default_params The number of default parameters of the kernel
3019    ///
3020    /// Example usage:
3021    /// With the following kernel:
3022    ///   void foo(sycl::float2 *x, int n, sycl::nd_item<3> item_ct1, float
3023    ///   f=.1) {}
3024    /// and with the declaration:
3025    ///   args_selector<2, 1, decltype(foo)> selector(kernelParams, extra);
3026    /// we have:
3027    ///   selector.get<0>() returns a reference to sycl::float*,
3028    ///   selector.get<1>() returns a reference to int,
3029    ///   selector.get<2>() returns a reference to float
3030    template <int n_nondefault_params, int n_default_params, typename R,
3031              typename... Ts>
3032    class args_selector<n_nondefault_params, n_default_params, R(Ts...)> {
3033      private:
3034        void **kernel_params;
3035        char *args_buffer;
3036
3037        template <int i> static constexpr int account_for_default_params() {
3038            constexpr int n_total_params = sizeof...(Ts);
3039            if constexpr (i >= n_nondefault_params) {
3040                return n_total_params - n_default_params +
3041                       (i - n_nondefault_params);
3042            } else {
3043                return i;
3044            }
3045        }
3046
3047      public:
3048        /// Get the type of the ith argument of R(Ts...)
3049        /// \param [in] i Index of parameter to get
3050        /// \returns Type of ith parameter
3051        template <int i>
3052        using arg_type = std::tuple_element_t<account_for_default_params<i>(),
3053                                              std::tuple<Ts...>>;
3054        static constexpr int params_num = sizeof...(Ts);
3055
3056      private:
3057        template <int i> static constexpr int get_offset() {
3058            if constexpr (i == 0) {
3059                // we can assume args_buffer is properly aligned to the
3060                // first argument
3061                return 0;
3062            } else {
3063                constexpr int prev_off = get_offset<i - 1>();
3064                constexpr int prev_past_end =
3065                    prev_off + sizeof(arg_type<i - 1>);
3066                using T = arg_type<i>;
3067                // is the past-the-end of the i-1st element properly aligned
3068                // with the ith element's alignment?
3069                if constexpr (prev_past_end % alignof(T) == 0) {
3070                    return prev_past_end;
3071                }
3072                // otherwise bump prev_past_end to match alignment
3073                else {
3074                    return prev_past_end +
3075                           (alignof(T) - (prev_past_end % alignof(T)));
3076                }
3077            }
3078        }
3079
3080        static char *get_args_buffer(void **extra) {
3081            if (!extra)
3082                return nullptr;
3083            for (; (std::size_t)*extra != 0; ++extra) {
3084                if ((std::size_t)*extra == 1) {
3085                    return static_cast<char *>(*(extra + 1));
3086                }
3087            }
3088            return nullptr;
3089        }
3090
3091      public:
3092        /// If kernel_params is nonnull, then args_selector will
3093        /// extract arguments from kernel_params. Otherwise, it
3094        /// will extract them from extra.
3095        /// \param [in] kernel_params Array of pointers to arguments
3096        /// a or null pointer.
3097        /// \param [in] extra Array containing pointer to argument buffer.
3098        args_selector(void **kernel_params, void **extra)
3099            : kernel_params(kernel_params),
3100              args_buffer(get_args_buffer(extra)) {}
3101
3102        /// Get a reference to the ith argument extracted from kernel_params
3103        /// or extra.
3104        /// \param [in] i Index of argument to get
3105        /// \returns Reference to the ith argument
3106        template <int i> arg_type<i> &get() {
3107            if (kernel_params) {
3108                return *static_cast<arg_type<i> *>(kernel_params[i]);
3109            } else {
3110                return *reinterpret_cast<arg_type<i> *>(args_buffer +
3111                                                        get_offset<i>());
3112            }
3113        }
3114    }; // COPY from DPCT head file
3115       // /opt/intel/oneapi/dpcpp-ct/latest/include/dpct/util.hpp
3116
3117    /// Utility class for launching SYCL kernels through kernel
3118    /// function wrapper.
3119    /// For example:
3120    /// A SYCL kernel function:
3121    ///   void kernel_func(int *ptr, sycl::nd_item<3> item);
3122    /// Kernel function wrapper:
3123    ///   void kernel_func_wrapper(int *ptr) {
3124    ///     sycl::queue queue = *dpct::kernel_launcher::_que;
3125    ///     unsigned int localMemSize = dpct::kernel_launcher::_local_mem_size;
3126    ///     sycl::nd_range<3> nr = dpct::kernel_launcher::_nr;
3127    ///     queue.parallel_for(
3128    ///       nr,
3129    ///       [=](sycl::nd_item<3> item_ct1) {
3130    ///         kernel_func(ptr, item_ct1);
3131    ///       });
3132    ///   }
3133    /// Then launch the kernel through wrapper like:
3134    ///   typedef void(*fpt)(int *);
3135    ///   fpt fp = kernel_func_wrapper;
3136    ///   dpct::kernel_launcher::launch(fp, dpct::dim3(1), dpct::dim3(1), 0, 0,
3137    ///   device_ptr);
3138    /// If the origin function type is erased, then need to register it first:
3139    ///   void *fp = (void *)wrapper_register(&kernel_func_wrapper).get();
3140    ///   dpct::kernel_launcher::launch(fp, dpct::dim3(1), dpct::dim3(1), args,
3141    ///   0, 0);
3142    class kernel_launcher {
3143        template <typename FuncT, typename ArgSelector, std::size_t... Index>
3144        static void launch_helper(FuncT &&func, ArgSelector &selector,
3145                                  std::index_sequence<Index...>) {
3146            func(selector.template get<Index>()...);
3147        }
3148        static void set_execution_config(dim3 group_range, dim3 local_range,
3149                                         unsigned int local_mem_size,
3150                                         queue_ptr que) {
3151            if (que) {
3152                _que = que;
3153            } else {
3154                _que = &get_default_queue();
3155            }
3156            _nr = sycl::nd_range<3>(
3157                static_cast<sycl::range<3>>(group_range * local_range),
3158                static_cast<sycl::range<3>>(local_range));
3159            _local_mem_size = local_mem_size;
3160
3161
3162        };
3163        static inline std::mutex kernel_function_ptr_map_mutex;
3164
3165      public:
3166        /// Variables for storing execution configuration.
3167        static inline thread_local sycl::queue *_que = nullptr;
3168        static inline thread_local sycl::nd_range<3> _nr = sycl::nd_range<3>();
3169        static inline thread_local unsigned int _local_mem_size = 0;
3170        /// Map for retrieving launchable functor from a raw pointer.
3171        static inline std::map<
3172            const void *,
3173            std::function<void(dim3, dim3, void **, unsigned int, queue_ptr)>>
3174            kernel_function_ptr_map = {};
3175
3176        /// Registers a kernel function pointer with a corresponding launchable
3177        /// functor.
3178        /// \param [in] func Pointer to the kernel function.
3179        /// \param [in] launcher Functor to handle kernel invocation.
3180        static void register_kernel_ptr(
3181            const void *func,
3182            std::function<void(dim3, dim3, void **, unsigned int, queue_ptr)>
3183                launcher) {
3184            std::lock_guard<std::mutex> lock(kernel_function_ptr_map_mutex);
3185            kernel_function_ptr_map[func] = std::move(launcher);
3186        }
3187        /// Launches a kernel function with arguments provided directly through
3188        /// kernel function wrapper.
3189        /// \tparam FuncT Type of the kernel function wrapper.
3190        /// \tparam ArgsT Types of kernel arguments.
3191        /// \param [in] func Pointer to the kernel function wrapper.
3192        /// \param [in] group_range SYCL group range.
3193        /// \param [in] local_range SYCL local range.
3194        /// \param [in] local_mem_size The size of local memory required by the
3195        /// kernel function. \param [in] que SYCL queue used to execute kernel.
3196        /// \param [in] args Kernel arguments.
3197        template <typename FuncT, typename... ArgsT>
3198        static std::enable_if_t<std::is_invocable_v<FuncT *, ArgsT...>, void>
3199        launch(FuncT *func, dim3 group_range, dim3 local_range,
3200               unsigned int local_mem_size, queue_ptr que, ArgsT... args) {
3201            set_execution_config(group_range, local_range, local_mem_size, que);
3202            func(args...);
3203        }
3204        /// Launches a kernel function through registered kernel function
3205        /// wrapper. \param [in] func Pointer to the registered kernel function
3206        /// wrapper. \param [in] group_range SYCL group range. \param [in]
3207        /// local_range SYCL local range. \param [in] args Array of pointers to
3208        /// kernel arguments. \param [in] local_mem_size The size of local
3209        /// memory required by the kernel function. \param [in] que SYCL queue
3210        /// used to execute kernel.
3211        static void launch(const void *func, dim3 group_range, dim3 local_range,
3212                           void **args, unsigned int local_mem_size,
3213                           queue_ptr que) {
3214            std::lock_guard<std::mutex> lock(kernel_function_ptr_map_mutex);
3215            auto Iter = kernel_function_ptr_map.find(func);
3216            if (Iter == kernel_function_ptr_map.end()) {
3217                throw std::runtime_error("dpct::launch() : no registered "
3218                                         "kernel function wrapper found.");
3219            }
3220            (Iter->second)(group_range, local_range, args, local_mem_size, que);
3221        }
3222        /// Launches a kernel function with packed arguments through kernel
3223        /// function wrapper.
3224        /// \tparam FuncT Type of the kernel function wrapper.
3225        /// \param [in] func Pointer to the kernel function wrapper.
3226        /// \param [in] group_range SYCL group range.
3227        /// \param [in] local_range SYCL local range.
3228        /// \param [in] args Array of pointers to kernel arguments.
3229        /// \param [in] local_mem_size The size of local memory required by the
3230        /// kernel function. \param [in] que SYCL queue used to execute kernel.
3231        template <typename FuncT>
3232        static std::enable_if_t<std::is_function_v<FuncT>, void>
3233        launch(FuncT *func, dim3 group_range, dim3 local_range, void **args,
3234               unsigned int local_mem_size, queue_ptr que) {
3235            constexpr size_t p_num = args_selector<0, 0, FuncT>::params_num;
3236            set_execution_config(group_range, local_range, local_mem_size, que);
3237            args_selector<p_num, p_num, FuncT> selector(args, nullptr);
3238            launch_helper(func, selector, std::make_index_sequence<p_num>{});
3239        }
3240    }; // COPY from DPCT head file
3241       // /opt/intel/oneapi/dpcpp-ct/latest/include/dpct/kernel.hpp
3242
3243    // /opt/intel/oneapi/dpcpp-ct/latest/include/dpct/util.hpp
3244    template <typename T>
3245    T select_from_sub_group(
3246        sycl::sub_group g,
3247        T x,
3248        int remote_local_id,
3249        int logical_sub_group_size = 32) {
3250      unsigned int start_index = g.get_local_linear_id() /
3251                                 logical_sub_group_size *
3252                                 logical_sub_group_size;
3253      return sycl::select_from_group(
3254          g, x, start_index + remote_local_id % logical_sub_group_size);
3255    }
3256
3257    // /opt/intel/oneapi/dpcpp-ct/latest/include/dpct/math.hpp
3258    template <typename T>
3259    void ldmatrix(uintptr_t addr, T* m, bool trans = false, unsigned mat = 0) {
3260      auto sg = sycl::ext::oneapi::this_work_item::get_sub_group();
3261      int lane = sg.get_local_linear_id();
3262
3263      int lane_group8_row = lane / 8;
3264      int lane_group8_col = lane % 8;
3265
3266      if (!trans) {
3267        // calculate the source lane
3268        int src_lane = 2 * lane_group8_row;
3269        if (lane_group8_col >= 4)
3270          src_lane += 1;
3271
3272        // Broadcast the address from the source lane
3273        auto recv_addr_uintp =
3274            dpct::select_from_sub_group(sg, addr, mat * 8 + src_lane);
3275
3276        // Cast the received address from uintptr_t to the type of 'm'
3277        auto recv_addr = reinterpret_cast<T*>(recv_addr_uintp);
3278
3279        // Non-transposed load
3280        *m = recv_addr[lane_group8_col % 4];
3281      } else {
3282        // calculate the source lane
3283        int src_lane = (lane % 4) * 2;
3284
3285        // Broadcast the address from the source lane
3286        auto recv_addr_uintp_1 =
3287            dpct::select_from_sub_group(sg, addr, mat * 8 + src_lane);
3288        auto recv_addr_uintp_2 =
3289            dpct::select_from_sub_group(sg, addr, mat * 8 + src_lane + 1);
3290
3291        // Cast the received address from uintptr_t to 'half *'
3292        auto recv_addr_1 = reinterpret_cast<sycl::half*>(recv_addr_uintp_1);
3293        auto recv_addr_2 = reinterpret_cast<sycl::half*>(recv_addr_uintp_2);
3294
3295        // Transposed load
3296        int index = lane / 4;
3297        sycl::half val0 = recv_addr_1[index];
3298        sycl::half val1 = recv_addr_2[index];
3299
3300        // Combine the two 16-bits into one 32-bit value
3301        sycl::half2 val = sycl::half2(val0, val1);
3302        *m = *reinterpret_cast<T*>(&val);
3303      }
3304    }
3305
3306    template <typename T>
3307    void ldmatrix(uintptr_t addr, T* m1, T* m2, bool trans = false) {
3308      // Load 1st matrix
3309      ldmatrix(addr, m1, trans, 0);
3310      // Load 2nd matrix
3311      ldmatrix(addr, m2, trans, 1);
3312    }
3313
3314    template <typename T>
3315    void ldmatrix(
3316        uintptr_t addr, T* m1, T* m2, T* m3, T* m4, bool trans = false) {
3317      // Load 1st matrix
3318      ldmatrix(addr, m1, trans, 0);
3319      // Load 2nd matrix
3320      ldmatrix(addr, m2, trans, 1);
3321      // Load 3rd matrix
3322      ldmatrix(addr, m3, trans, 2);
3323      // Load 4th matrix
3324      ldmatrix(addr, m4, trans, 3);
3325    }
3326
3327    // /opt/intel/oneapi/dpcpp-ct/latest/include/dpct/math.hpp
3328
3329    /// A helper struct that defines the pack type for the input matrix
3330    /// fragments
3331    /// of mma() function based on the type of input matrix fragments.
3332    /// The MMAType struct is specialized for different types of input matrices.
3333    /// Currently, the specialization for f16, bf16 and s8 types is defined
3334    /// below. \tparam [in] T The type of the input matrix fragments
3335    template <typename T>
3336    struct MMAType {
3337      using PackType = uint32_t;
3338    };
3339
3340    /// Each work item of a sub-group (limited to size 32) calling this function
3341    /// calculates a subset fragment for the output matrix D using MAD operation
3342    /// on A, B & C matrix fragments (D = A * B + C). Current supported shapes &
3343    /// types:
3344    /// - m8n8k4 (f32.f16.f16.f32)
3345    /// - m8n8k16 (s32.s8.s8.s32)
3346    /// - m16n8k8 (f32.f16.f16.f32 & f32.bf16.bf16.f32)
3347    /// - m16n8k16 (f32.f16.f16.f32 & s32.s8.s8.s32)
3348    /// - m16n8k32 (s32.s8.s8.s32)
3349    /// Here, m, n & k define the shapes of A, B & C matrices respectively
3350    /// (A = [m x k], B = [k x n], C = [m x n]).
3351    /// \tparam [in] M The rows of A, C & D matrices
3352    /// \tparam [in] N The columns of B, C, D matrices
3353    /// \tparam [in] K The columns & rows of A & B matrices respectively
3354    /// \tparam [in] ABType The type of the input matrix (A & B) fragment
3355    /// \tparam [in] CDType The type of the output matrix (C & D) fragment
3356    /// \param [out] d_mat_frag The fragment of the output matrix D to store the
3357    /// result of A * B + C
3358    /// \param [in] a_mat_frag The fragment of the input matrix A to be
3359    /// multiplied with B matrix fragment \param [in] b_mat_frag The fragment of
3360    /// the input matrix B to be multiplied with A matrix fragment \param [in]
3361    /// c_mat_frag The fragment of the input matrix C to be added with the
3362    /// result of A * B fragments
3363    template <int M, int N, int K, typename ABType, typename CDType>
3364    void mma(
3365        volatile void** d_mat_frag,
3366        void* a_mat_frag,
3367        void* b_mat_frag,
3368        void* c_mat_frag) {
3369      auto d = reinterpret_cast<volatile CDType**>(d_mat_frag);
3370      auto a =
3371          reinterpret_cast<typename MMAType<ABType>::PackType*>(a_mat_frag);
3372      auto b =
3373          reinterpret_cast<typename MMAType<ABType>::PackType*>(b_mat_frag);
3374      auto c = reinterpret_cast<CDType*>(c_mat_frag);
3375
3376      auto sg = sycl::ext::oneapi::this_work_item::get_sub_group();
3377      int lane = sg.get_local_linear_id();
3378
3379      static_assert(
3380          (M == 8 && N == 8 && K == 4) || (M == 8 && N == 8 && K == 16) ||
3381              (M == 16 && N == 8 && K == 8) || (M == 16 && N == 8 && K == 16) ||
3382              (M == 16 && N == 8 && K == 32),
3383          "Unsupported MMA shape!");
3384
3385      short row_load_offset = 4 * (lane >> 2);
3386      short col_load_offset = 8 * (lane % 4);
3387
3388      if constexpr (M == 8 && N == 8 && K == 4) {
3389        if constexpr (std::is_floating_point_v<CDType>) {
3390          col_load_offset = row_load_offset % 16;
3391
3392          // Init D matrix with fragments of C matrix
3393          *d[0] = c[0];
3394          *d[1] = c[1];
3395          *d[2] = c[2];
3396          *d[3] = c[3];
3397          *d[4] = c[4];
3398          *d[5] = c[5];
3399          *d[6] = c[6];
3400          *d[7] = c[7];
3401
3402          // Calculate the row and col offset indices to iterate through the row
3403          // & col fragments of A & B matrices
3404          int r_ind = (lane % 2) ? 1 : 0;
3405          int c_ind = ((lane % 4) / 2) ? 2 : 0;
3406
3407          // Each sub-group is responsible for computing a fragment size of 8*8
3408          // elements of matrix D for each of 4 MMA computations.
3409          // Each work item computes 8 elements of matrix D by gathering
3410          // their corresponding col & row matrix fragments of length k (4)
3411          // from A & B matrices respectively using below mapping logic:
3412          // row0 = (i % 4) if (lane < 16) else (i % 4) + 4
3413          // col0 = (lane % 4)
3414          // As each row & col fragment of A & B matrices is distributed across
3415          // 4 work items, each iteration of below loop loads a partial fragment
3416          // of matrix A (row) and matrix B (col) using the row & col offsets.
3417          typename MMAType<ABType>::PackType recv_a[2], recv_b[2];
3418
3419          for (int i = 0; i < 4; i++) {
3420            // Load partial fragment from col0 of matrix A ({a0, a1})
3421            recv_a[0] =
3422                dpct::select_from_sub_group(sg, a[0], row_load_offset + i);
3423            // Load partial fragment from col0 of matrix A ({a2, a3})
3424            recv_a[1] =
3425                dpct::select_from_sub_group(sg, a[1], row_load_offset + i);
3426
3427            // Load partial fragment from row0 of matrix B ({b0, b1})
3428            recv_b[0] =
3429                dpct::select_from_sub_group(sg, b[0], col_load_offset + i);
3430            // Load partial fragment from row0 of matrix B ({b2, b3})
3431            recv_b[1] =
3432                dpct::select_from_sub_group(sg, b[1], col_load_offset + i);
3433
3434            auto ra = reinterpret_cast<ABType*>(recv_a);
3435            auto rb = reinterpret_cast<ABType*>(recv_b);
3436
3437            // Each work item calculates a partial product of A & B matrix
3438            // fragments and adds it to the corresponding D matrix fragment (for
3439            // even work item indices) d0 += col0{ a0 } * row0{ b0 } d1 += col0{
3440            // a0 } * row0{ b1 } d2 += col1{ a2 } * row0{ b0 } d3 += col1{ a2 }
3441            // * row0{ b1 } (for odd work item indices) d0 += col0{ a1 } * row0{
3442            // b2 } d1 += col0{ a1 } * row0{ b3 } d2 += col1{ a3 } * row0{ b2 }
3443            // d3 += col1{ a3 } * row0{ b3 }
3444            *d[0] +=
3445                static_cast<float>(ra[r_ind]) * static_cast<float>(rb[c_ind]);
3446            *d[1] += static_cast<float>(ra[r_ind]) *
3447                     static_cast<float>(rb[c_ind + 1]);
3448            *d[2] += static_cast<float>(ra[r_ind + 2]) *
3449                     static_cast<float>(rb[c_ind]);
3450            *d[3] += static_cast<float>(ra[r_ind + 2]) *
3451                     static_cast<float>(rb[c_ind + 1]);
3452
3453            // Load partial fragment from row1 of matrix B ({b0, b1})
3454            recv_b[0] =
3455                dpct::select_from_sub_group(sg, b[0], col_load_offset + i + 16);
3456            // Load partial fragment from row1 of matrix B ({b2, b3})
3457            recv_b[1] =
3458                dpct::select_from_sub_group(sg, b[1], col_load_offset + i + 16);
3459
3460            // (for even work item indices)
3461            // d0 += col0{ a0 } * row1{ b0 }
3462            // d1 += col0{ a0 } * row1{ b1 }
3463            // d2 += col1{ a2 } * row1{ b0 }
3464            // d3 += col1{ a2 } * row1{ b1 }
3465            // (for odd work item indices)
3466            // d0 += col0{ a1 } * row1{ b2 }
3467            // d1 += col0{ a1 } * row1{ b3 }
3468            // d2 += col1{ a3 } * row1{ b2 }
3469            // d3 += col1{ a3 } * row1{ b3 }
3470            *d[4] +=
3471                static_cast<float>(ra[r_ind]) * static_cast<float>(rb[c_ind]);
3472            *d[5] += static_cast<float>(ra[r_ind]) *
3473                     static_cast<float>(rb[c_ind + 1]);
3474            *d[6] += static_cast<float>(ra[r_ind + 2]) *
3475                     static_cast<float>(rb[c_ind]);
3476            *d[7] += static_cast<float>(ra[r_ind + 2]) *
3477                     static_cast<float>(rb[c_ind + 1]);
3478          }
3479        }
3480      } else if constexpr (M == 8 && N == 8 && K == 16) {
3481        if constexpr (std::is_integral_v<ABType>) {
3482          // Init D matrix with fragments of C matrix
3483          *d[0] = c[0];
3484          *d[1] = c[1];
3485
3486          // Each sub-group is responsible for computing a fragment size of 16*8
3487          // elements of matrix D.
3488          // Each work item computes 2 elements of matrix D by gathering
3489          // their corresponding row & col matrix fragments of length k (16)
3490          // from A & B matrices respectively using below mapping logic:
3491          // row0 = ((lane % 4) * 4) + i
3492          // col0 = (lane >> 2)
3493          // As each row & col fragment of A & B matrices is distributed across
3494          // 4 work items, each iteration of below loop loads a partial fragment
3495          // of matrix A (row) and matrix B (col) using the row & col offsets.
3496          for (int i = 0; i < 4; i++) {
3497            typename MMAType<ABType>::PackType recv_a, recv_b[2];
3498
3499            // Load partial fragment from row0 of matrix A ({a0, a1, a2, a3})
3500            recv_a = dpct::select_from_sub_group(sg, a[0], row_load_offset + i);
3501            // Load partial fragment from col0 of matrix B ({b0, b1, b2, b3})
3502            recv_b[0] =
3503                dpct::select_from_sub_group(sg, b[0], col_load_offset + i);
3504            // Load partial fragment from col1 of matrix B ({b0, b1, b2, b3})
3505            recv_b[1] =
3506                dpct::select_from_sub_group(sg, b[0], col_load_offset + i + 4);
3507
3508            auto a = reinterpret_cast<ABType*>(&recv_a);
3509            auto b = reinterpret_cast<ABType*>(recv_b);
3510
3511            // Each work item calculates a partial product of A & B matrix
3512            // fragments and adds it to the corresponding D matrix fragment d0
3513            // += row0{ a0, a1, a2, a3 } * col0{ b0, b1, b2, b3 } d1 += row0{
3514            // a0, a1, a2, a3 } * col1{ b0, b1, b2, b3 } d2 += row0{ a0, a1, a2,
3515            // a3 } * col0{ b0, b1, b2, b3 } d3 += row0{ a0, a1, a2, a3 } *
3516            // col1{ b0, b1, b2, b3 }
3517            for (int j = 0; j < 4; j++) {
3518              *d[0] += a[j] * b[j];
3519              *d[1] += a[j] * b[j + 4];
3520            }
3521          }
3522        }
3523      } else if constexpr (M == 16 && N == 8 && K == 8) {
3524        if constexpr (std::is_floating_point_v<CDType>) {
3525          // Init D matrix fragment with C matrix fragment
3526          *d[0] = c[0];
3527          *d[1] = c[1];
3528          *d[2] = c[2];
3529          *d[3] = c[3];
3530
3531          // Each sub-group is responsible for computing a fragment size of 16*8
3532          // elements of matrix D.
3533          // Each work item computes 4 elements of matrix D by gathering
3534          // their corresponding row & col matrix fragments of length k (8)
3535          // from A & B matrices respectively using below mapping logic:
3536          // row0 = (lane >> 2) & row1 = (lane >> 2) + 8
3537          // col0 = (lane % 4) * 2 + (i & 0x1)
3538          // As each row & col fragment of A & B matrices is distributed across
3539          // 4 work items, each iteration of below loop loads a partial fragment
3540          // of matrix A (row) and matrix B (col) using the row & col offsets.
3541          for (int i = 0; i < 4; i++) {
3542            typename MMAType<ABType>::PackType recv_a[2], recv_b[2];
3543
3544            // Load partial fragment from row0 of matrix A ({a0, a1})
3545            recv_a[0] =
3546                dpct::select_from_sub_group(sg, a[0], row_load_offset + i);
3547            // Load partial fragment from row1 of matrix A ({a2, a3})
3548            recv_a[1] =
3549                dpct::select_from_sub_group(sg, a[1], row_load_offset + i);
3550            // Load partial fragment from col0 of matrix B ({b0, b1})
3551            recv_b[0] =
3552                dpct::select_from_sub_group(sg, b[0], col_load_offset + i);
3553            // Load partial fragment from col1 of matrix B ({b0, b1})
3554            recv_b[1] =
3555                dpct::select_from_sub_group(sg, b[0], col_load_offset + i + 4);
3556
3557            auto ra = reinterpret_cast<ABType*>(recv_a);
3558            auto rb = reinterpret_cast<ABType*>(recv_b);
3559
3560            // Each work item calculates a partial product of A & B matrix
3561            // fragments and adds it to the corresponding D matrix fragment d0
3562            // += row0{ a0, a1 } * col0{ b0, b1 } d1 += row0{ a0, a1 } * col1{
3563            // b0, b1 } d2 += row1{ a2, a3 } * col0{ b0, b1 } d3 += row1{ a2, a3
3564            // } * col1{ b0, b1 }
3565            for (int j = 0; j < 2; j++) {
3566              *d[0] += static_cast<float>(ra[j]) * static_cast<float>(rb[j]);
3567              *d[1] +=
3568                  static_cast<float>(ra[j]) * static_cast<float>(rb[j + 2]);
3569              *d[2] +=
3570                  static_cast<float>(ra[j + 2]) * static_cast<float>(rb[j]);
3571              *d[3] +=
3572                  static_cast<float>(ra[j + 2]) * static_cast<float>(rb[j + 2]);
3573            }
3574          }
3575        }
3576      } else if constexpr (M == 16 && N == 8 && K == 16) {
3577        if constexpr (std::is_floating_point_v<CDType>) {
3578          // Init D matrix fragment with C matrix fragment
3579          *d[0] = c[0];
3580          *d[1] = c[1];
3581          *d[2] = c[2];
3582          *d[3] = c[3];
3583
3584          // Each sub-group is responsible for computing a fragment size of 16*8
3585          // elements of matrix D.
3586          // Each work item computes 4 elements of matrix D by gathering
3587          // their corresponding row & col matrix fragments of length k (8)
3588          // from A & B matrices respectively using below mapping logic:
3589          // row0 = (lane >> 2)    & row1 = (lane >> 2) + 8
3590          // col0 = (lane % 4) * 2 & col1 = (lane % 4) * 2 + 1
3591          // As each row & col fragment of A & B matrices is distributed across
3592          // 4 work items, each iteration of below loop loads a partial fragment
3593          // of matrix A (row) and matrix B (col) using the row & col offsets.
3594          for (int i = 0; i < 4; i++) {
3595            typename MMAType<ABType>::PackType recv_a[4], recv_b[4];
3596
3597            // Load partial fragment from row0 of matrix A ({a0, a1})
3598            recv_a[0] =
3599                dpct::select_from_sub_group(sg, a[0], row_load_offset + i);
3600            // Load partial fragment from row0 of matrix A ({a2, a3})
3601            recv_a[1] =
3602                dpct::select_from_sub_group(sg, a[2], row_load_offset + i);
3603            // Load partial fragment from row1 of matrix A ({a0, a1})
3604            recv_a[2] =
3605                dpct::select_from_sub_group(sg, a[1], row_load_offset + i);
3606            // Load partial fragment from row1 of matrix A ({a2, a3})
3607            recv_a[3] =
3608                dpct::select_from_sub_group(sg, a[3], row_load_offset + i);
3609
3610            // Load partial fragment from col0 of matrix B ({b0, b1})
3611            recv_b[0] =
3612                dpct::select_from_sub_group(sg, b[0], col_load_offset + i);
3613            // Load partial fragment from col0 of matrix B ({b2, b3})
3614            recv_b[1] =
3615                dpct::select_from_sub_group(sg, b[1], col_load_offset + i);
3616            // Load partial fragment from col1 of matrix B ({b0, b1})
3617            recv_b[2] =
3618                dpct::select_from_sub_group(sg, b[0], col_load_offset + 4 + i);
3619            // Load partial fragment from col1 of matrix B ({b2, b3})
3620            recv_b[3] =
3621                dpct::select_from_sub_group(sg, b[1], col_load_offset + 4 + i);
3622
3623            auto ra = reinterpret_cast<ABType*>(recv_a);
3624            auto rb = reinterpret_cast<ABType*>(recv_b);
3625
3626            // Each work item calculates a partial product of A & B matrix
3627            // fragments and adds it to the corresponding D matrix fragment d0
3628            // += row0{ a0, a1, a2, a3 } * col0{ b0, b1, b2, b3 } d1 += row0{
3629            // a0, a1, a2, a3 } * col1{ b0, b1, b2, b3 } d2 += row1{ a0, a1, a2,
3630            // a3 } * col0{ b0, b1, b2, b3 } d3 += row1{ a0, a1, a2, a3 } *
3631            // col1{ b0, b1, b2, b3 }
3632            for (int j = 0; j < 4; j++) {
3633              *d[0] += static_cast<CDType>(ra[j]) * static_cast<CDType>(rb[j]);
3634              *d[1] +=
3635                  static_cast<CDType>(ra[j]) * static_cast<CDType>(rb[j + 4]);
3636              *d[2] +=
3637                  static_cast<CDType>(ra[j + 4]) * static_cast<CDType>(rb[j]);
3638              *d[3] += static_cast<CDType>(ra[j + 4]) *
3639                       static_cast<CDType>(rb[j + 4]);
3640            }
3641          }
3642        } else if constexpr (std::is_integral_v<ABType>) {
3643          // Init D matrix with fragments of C matrix
3644          *d[0] = c[0];
3645          *d[1] = c[1];
3646          *d[2] = c[2];
3647          *d[3] = c[3];
3648
3649          // Each sub-group is responsible for computing a fragment size of 16*8
3650          // elements of matrix D.
3651          // Each work item computes 4 elements of matrix D by gathering
3652          // their corresponding row & col matrix fragments of length k (8)
3653          // from A & B matrices respectively using below mapping logic:
3654          // row0 = (lane >> 2)    & row1 = (lane >> 2) + 8
3655          // col0 = (lane % 4) * 2 & col1 = (lane % 4) * 2 + 1
3656          // As each row & col fragment of A & B matrices is distributed across
3657          // 4 work items, each iteration of below loop loads a partial fragment
3658          // of matrix A (row) and matrix B (col) using the row & col offsets.
3659          for (int i = 0; i < 4; i++) {
3660            typename MMAType<ABType>::PackType recv_a[2], recv_b[2];
3661
3662            // Load partial fragment from row0 of matrix A ({a0, a1, a2, a3})
3663            recv_a[0] =
3664                dpct::select_from_sub_group(sg, a[0], row_load_offset + i);
3665            // Load partial fragment from row1 of matrix A ({a4, a5, a6, a7})
3666            recv_a[1] =
3667                dpct::select_from_sub_group(sg, a[1], row_load_offset + i);
3668            // Load partial fragment from col0 of matrix B ({b0, b1, b2, b3})
3669            recv_b[0] =
3670                dpct::select_from_sub_group(sg, b[0], col_load_offset + i);
3671            // Load partial fragment from col1 of matrix B ({b4, b5, b6, b7})
3672            recv_b[1] =
3673                dpct::select_from_sub_group(sg, b[0], col_load_offset + i + 4);
3674
3675            auto ra = reinterpret_cast<ABType*>(recv_a);
3676            auto rb = reinterpret_cast<ABType*>(recv_b);
3677
3678            // Each work item calculates a partial product of A & B matrix
3679            // fragments and adds it to the corresponding D matrix fragment d0
3680            // += row0{ a0, a1, a2, a3 } * col0{ b0, b1, b2, b3 } d1 += row0{
3681            // a0, a1, a2, a3 } * col1{ b4, b5, b6, b7 } d2 += row1{ a4, a5, a6,
3682            // a7 } * col0{ b0, b1, b2, b3 } d3 += row1{ a4, a5, a6, a7 } *
3683            // col1{ b4, b5, b6, b7 }
3684            for (int i = 0; i < 4; i++) {
3685              *d[0] += ra[i] * rb[i];
3686              *d[1] += ra[i] * rb[i + 4];
3687              *d[2] += ra[i + 4] * rb[i];
3688              *d[3] += ra[i + 4] * rb[i + 4];
3689            }
3690          }
3691        }
3692      } else if constexpr (M == 16 && N == 8 && K == 32) {
3693        if constexpr (std::is_integral_v<ABType>) {
3694          // Init D matrix with fragments of C matrix
3695          *d[0] = c[0];
3696          *d[1] = c[1];
3697          *d[2] = c[2];
3698          *d[3] = c[3];
3699
3700          // Each sub-group is responsible for computing a fragment size of 16*8
3701          // elements of matrix D.
3702          // Each work item computes 4 elements of matrix D by gathering
3703          // their corresponding row & col matrix fragments of length k (32)
3704          // from A & B matrices respectively using below mapping logic:
3705          // row0 = (lane >> 2)    & row1 = (lane >> 2) + 8
3706          // col0 = ((lane % 4) * 4) + (i & 0x3) & col1 = ((lane % 4) * 4) + (i
3707          // & 0x3) As each row & col fragment of A & B matrices is distributed
3708          // across 4 work items, each iteration of below loop loads a partial
3709          // fragment of matrix A (row) and matrix B (col) using the row & col
3710          // offsets.
3711          for (int i = 0; i < 4; i++) {
3712            typename MMAType<ABType>::PackType recv_a[2], recv_b[2];
3713
3714            // Load partial fragment from row0 of matrix A ({a0, a1, a2, a3})
3715            recv_a[0] =
3716                dpct::select_from_sub_group(sg, a[0], row_load_offset + i);
3717            // Load partial fragment from row1 of matrix A ({a4, a5, a6, a7})
3718            recv_a[1] =
3719                dpct::select_from_sub_group(sg, a[1], row_load_offset + i);
3720            // Load partial fragment from col0 of matrix B ({b0, b1, b2, b3})
3721            recv_b[0] =
3722                dpct::select_from_sub_group(sg, b[0], col_load_offset + i);
3723            // Load partial fragment from col1 of matrix B ({b0, b1, b2, b3})
3724            recv_b[1] =
3725                dpct::select_from_sub_group(sg, b[0], col_load_offset + i + 4);
3726
3727            auto a = reinterpret_cast<ABType*>(recv_a);
3728            auto b = reinterpret_cast<ABType*>(recv_b);
3729
3730            // Each work item calculates a partial product of A & B matrix
3731            // fragments and adds it to the corresponding D matrix fragment d0
3732            // += row0{ a0, a1, a2, a3 } * col0{ b0, b1, b2, b3 } d1 += row0{
3733            // a0, a1, a2, a3 } * col1{ b0, b1, b2, b3 } d2 += row1{ a4, a5, a6,
3734            // a7 } * col0{ b0, b1, b2, b3 } d3 += row1{ a4, a5, a6, a7 } *
3735            // col1{ b0, b1, b2, b3 }
3736            for (int j = 0; j < 4; j++) {
3737              *d[0] += a[j] * b[j];
3738              *d[1] += a[j] * b[j + 4];
3739              *d[2] += a[j + 4] * b[j];
3740              *d[3] += a[j + 4] * b[j + 4];
3741            }
3742          }
3743
3744          for (int i = 0; i < 4; i++) {
3745            typename MMAType<ABType>::PackType recv_a[2], recv_b[2];
3746
3747            // Load partial fragment from row0 of matrix A ({a8, a9, a10, a11})
3748            recv_a[0] =
3749                dpct::select_from_sub_group(sg, a[2], row_load_offset + i);
3750            // Load partial fragment from row1 of matrix A ({a12, a13, a14,
3751            // a15})
3752            recv_a[1] =
3753                dpct::select_from_sub_group(sg, a[3], row_load_offset + i);
3754            // Load partial fragment from col0 of matrix B ({b4, b5, b6, b7})
3755            recv_b[0] =
3756                dpct::select_from_sub_group(sg, b[1], col_load_offset + i);
3757            // Load partial fragment from col1 of matrix B ({b4, b5, b6, b7})
3758            recv_b[1] =
3759                dpct::select_from_sub_group(sg, b[1], col_load_offset + i + 4);
3760
3761            auto a = reinterpret_cast<ABType*>(recv_a);
3762            auto b = reinterpret_cast<ABType*>(recv_b);
3763
3764            // Each work item calculates a partial product of A & B matrix
3765            // fragments and adds it to the corresponding D matrix fragment d0
3766            // += row0{ a8, a9, a10, a11 } * col0{ b4, b5, b6, b7 } d1 += row0{
3767            // a8, a9, a10, a11 } * col1{ b4, b5, b6, b7 } d2 += row1{ a12, a13,
3768            // a14, a15 } * col0{ b4, b5, b6, b7 } d3 += row1{ a12, a13, a14,
3769            // a15 } * col1{ b4, b5, b6, b7 }
3770            for (int j = 0; j < 4; j++) {
3771              *d[0] += a[j] * b[j];
3772              *d[1] += a[j] * b[j + 4];
3773              *d[2] += a[j + 4] * b[j];
3774              *d[3] += a[j + 4] * b[j + 4];
3775            }
3776          }
3777        }
3778      }
3779    }
3780} // COPY from DPCT head files
3781
3782#endif // GGML_SYCL_DPCT_HELPER_HPP