GPU: make the kernel entry-point signature work on Metal

MSL has no by-value kernel parameters: every one must carry an attribute naming
a buffer or a builtin. It also has no ambient work-item builtins, so the grid
dimensions arrive as attributes on the entry point too.

GPUCA_KRNLGPU_DEF therefore gets two hooks, GPUCA_KRNL_SECTOR_ARG and
GPUCA_KRNL_GRID_ARGS, which the Metal source fills in with a buffer and the four
grid attributes. Both default to what the signature had, so CUDA, HIP and OpenCL
generate exactly the same entry point as before.

The attributes are named after the nBlocks, nThreads, iBlock and iThread that
Thread() already takes, so the get_*() helpers resolve at the entry point and in
everything it calls. Metal still needs its own definitions of them because the
host ones assume one thread per block: get_local_id() is 0 and get_local_size()
is 1 there.

The generated arguments need an attribute too, each with a distinct buffer index,
and the preprocessor cannot supply one: the kernel list splices arguments as a
flat comma-separated list, and __COUNTER__ is monotonic across the translation
unit rather than per kernel. o2_gpu_add_kernel already walks the arguments in
pairs, so it emits the index there. Declarations go through GPUPtr1(idx, type,
name) for pointers and GPUArg1(idx, type, name) for scalars, which each backend
defines as it needs. Indices start at 3, after gpu_mem, the constant memory and
the sector; the sector itself is nothing special, it just lives in the fixed part
of the macro rather than in the generated list.

Metal masks pointers as a 64-bit address exactly as OpenCL does, and for the same
reason: GPUTRDTrackerKernels takes a GPUTRDTrackerGPU*, and a pointer to a
derived class is not a valid kernel argument type there either. Binding POD
pointers directly would have worked but would not have covered that case, so both
go the same way. On the way back in, GPUPtr2 casts through device before handing
the pointer to Thread(): the kernel's own buffers are device memory, but the
Thread() entry points take the pointer unannotated, which in MSL means generic.

Generated entry points are byte-identical for CUDA, HIP and OpenCL. Kernel list
diagnostics: 408 to 0, and the translation unit 1108 to 881.

Author: ktf

GPU/Common/GPUCommonDefAPI.h

@@ -277,6 +277,16 @@
   #define get_group_id(dim) (blockIdx.x)
 #elif defined(__OPENCL__)
   // Using OpenCL defaults
+#elif defined(__METAL__)
+  // MSL has no work-item builtins. They arrive as attributes on the entry point,
+  // named there after the nBlocks / nThreads / iBlock / iThread that Thread()
+  // already takes, so these resolve both there and in every function below it.
+  #define get_global_id(dim) (iBlock * nThreads + iThread)
+  #define get_global_size(dim) (nBlocks * nThreads)
+  #define get_num_groups(dim) (nBlocks)
+  #define get_local_id(dim) (iThread)
+  #define get_local_size(dim) (nThreads)
+  #define get_group_id(dim) (iBlock)
 #else
   #define get_global_id(dim) iBlock
   #define get_global_size(dim) nBlocks

GPU/GPUTracking/Base/GPUReconstructionKernelMacros.h

@@ -63,8 +63,17 @@
 #define GPUCA_ATTRRES(...) GPUCA_M_EXPAND(GPUCA_M_CAT(GPUCA_ATTRRES_, GPUCA_M_FIRST(__VA_ARGS__)))(__VA_ARGS__)
 
 // GPU Kernel entry point
+// MSL requires every kernel parameter to carry an attribute, and supplies the
+// grid dimensions the same way, so the backend gets to shape both ends of the
+// parameter list.
+#ifndef GPUCA_KRNL_SECTOR_ARG
+#define GPUCA_KRNL_SECTOR_ARG int32_t _iSector_internal
+#endif
+#ifndef GPUCA_KRNL_GRID_ARGS
+#define GPUCA_KRNL_GRID_ARGS
+#endif
 #define GPUCA_KRNLGPU_DEF(x_class, x_attributes, x_arguments, ...) \
-  GPUg() void GPUCA_ATTRRES(GPUCA_M_STRIP(x_attributes)) GPUCA_M_CAT(krnl_, GPUCA_M_KRNL_NAME(x_class))(GPUCA_CONSMEM_PTR int32_t _iSector_internal GPUCA_M_STRIP(x_arguments))
+  GPUg() void GPUCA_ATTRRES(GPUCA_M_STRIP(x_attributes)) GPUCA_M_CAT(krnl_, GPUCA_M_KRNL_NAME(x_class))(GPUCA_CONSMEM_PTR GPUCA_KRNL_SECTOR_ARG GPUCA_M_STRIP(x_arguments) GPUCA_KRNL_GRID_ARGS)
 
 #ifdef GPUCA_KRNL_DEFONLY
 #define GPUCA_KRNLGPU(...) GPUCA_KRNLGPU_DEF(__VA_ARGS__);

GPU/GPUTracking/Base/metal/GPUReconstructionMETAL.metal

@@ -70,6 +70,16 @@ using namespace metal;
   device char* pConstantRaw   [[buffer(1)]],
 #define GPUCA_CONSMEM (*(device GPUConstantMem*)pConstantRaw)
 
+// Every kernel parameter needs an attribute, so the sector index arrives as a
+// buffer rather than by value, and the grid dimensions come in at the end, where
+// GPUCommonDefAPI.h's get_group_id() and friends pick them up.
+#define GPUCA_KRNL_SECTOR_ARG constant int32_t& _iSector_internal [[buffer(2)]]
+#define GPUCA_KRNL_GRID_ARGS \
+  , uint iBlock   [[threadgroup_position_in_grid]] \
+  , uint iThread  [[thread_position_in_threadgroup]] \
+  , uint nThreads [[threads_per_threadgroup]] \
+  , uint nBlocks  [[threadgroups_per_grid]]
+
 #include "GPUReconstructionKernelList.h"
 
 // clang-format on

GPU/GPUTracking/Definitions/GPUDef.h

@@ -21,17 +21,27 @@
 #include "GPUDefParametersWrapper.h"
 #include "GPUCommonRtypes.h"
 
-// Macros for masking ptrs in OpenCL kernel calls as uint64_t (The API only allows us to pass buffer objects)
+// Macros for kernel arguments. OpenCL can only pass buffer objects, so pointers
+// are masked as uint64_t and cast back inside the kernel. MSL needs an explicit
+// buffer index on every parameter, but can bind a pointer directly. The index is
+// emitted per argument by o2_gpu_add_kernel; 0, 1 and 2 are taken by gpu_mem,
+// the constant memory and the sector index.
 #ifdef __OPENCL__
-  #define GPUPtr1(a, b) uint64_t b
-  #ifdef __OPENCL__
-    #define GPUPtr2(a, b) ((__generic a) (a) b)
-  #else
-    #define GPUPtr2(a, b) ((__global a) (a) b)
-  #endif
+  #define GPUPtr1(idx, a, b) uint64_t b
+  #define GPUPtr2(a, b) ((__generic a) (a) b)
+  #define GPUArg1(idx, a, b) a b
+#elif defined(__METAL__)
+  // As for OpenCL, pointers travel as a 64-bit address: a pointer to a derived
+  // class is not a valid kernel argument type in MSL either.
+  #define GPUPtr1(idx, a, b) constant uint64_t& b [[buffer(idx)]]
+  // through device and then to generic: the kernel's own buffers are device
+  // memory, but the Thread() entry points take the pointer unannotated
+  #define GPUPtr2(a, b) ((a)((device a)(b)))
+  #define GPUArg1(idx, a, b) constant a& b [[buffer(idx)]]
 #else
-  #define GPUPtr1(a, b) a b
+  #define GPUPtr1(idx, a, b) a b
   #define GPUPtr2(a, b) b
+  #define GPUArg1(idx, a, b) a b
 #endif
 
 #define GPUCA_EVDUMP_FILE "event"

GPU/GPUTracking/cmake/kernel_helpers.cmake

@@ -55,11 +55,13 @@ function(o2_gpu_add_kernel kernel_name kernel_files)
     math(EXPR n "${n} - 1")
     foreach(i RANGE 3 ${n} 2)
       math(EXPR j "${i} + 1")
+      # buffer indices 0, 1 and 2 are gpu_mem, the constant memory and the sector
+      math(EXPR TMP_ARG_IDX "3 + (${i} - 3) / 2")
       if(${ARGV${i}} MATCHES "\\*$")
-        string(APPEND OPT1 ",GPUPtr1(${ARGV${i}},${ARGV${j}})")
+        string(APPEND OPT1 ",GPUPtr1(${TMP_ARG_IDX},${ARGV${i}},${ARGV${j}})")
         string(APPEND OPT2 ",GPUPtr2(${ARGV${i}},${ARGV${j}})")
       else()
-        string(APPEND OPT1 ",${ARGV${i}} ${ARGV${j}}")
+        string(APPEND OPT1 ",GPUArg1(${TMP_ARG_IDX},${ARGV${i}},${ARGV${j}})")
         string(APPEND OPT2 ",${ARGV${j}}")
       endif()
       string(APPEND OPT3 ",${ARGV${i}}")