Support transfers for layouts with multiple id-indexed fields

src/realm/cuda/cuda_internal.h

@@ -416,7 +416,8 @@ namespace Realm {
       void launch_batch_affine_fill_kernel(void *fill_info, size_t dim, size_t elemSize,
                                            size_t volume, GPUStream *stream);
       void launch_batch_affine_kernel(void *copy_info, size_t dim, size_t elemSize,
-                                      size_t volume, GPUStream *stream);
+                                      size_t volume, bool multified_optimized,
+                                      GPUStream *stream);
       void launch_transpose_kernel(MemcpyTransposeInfo<size_t> &copy_info,
                                    size_t elemSize, GPUStream *stream);
 
@@ -470,6 +471,8 @@ namespace Realm {
       GPUFuncInfo indirect_copy_kernels[REALM_MAX_DIM][CUDA_MEMCPY_KERNEL_MAX2_LOG2_BYTES]
                                        [CUDA_MEMCPY_KERNEL_MAX2_LOG2_BYTES];
       GPUFuncInfo batch_affine_kernels[REALM_MAX_DIM][CUDA_MEMCPY_KERNEL_MAX2_LOG2_BYTES];
+      GPUFuncInfo multi_batch_affine_kernels[REALM_MAX_DIM]
+                                            [CUDA_MEMCPY_KERNEL_MAX2_LOG2_BYTES];
       GPUFuncInfo batch_fill_affine_kernels[REALM_MAX_DIM]
                                            [CUDA_MEMCPY_KERNEL_MAX2_LOG2_BYTES];
       GPUFuncInfo fill_affine_large_kernels[REALM_MAX_DIM]
@@ -800,9 +803,28 @@ namespace Realm {
 
       bool progress_xd(GPUChannel *channel, TimeLimit work_until);
 
+      static size_t read_address_entry(AffineCopyInfo<3> &copy_infos, size_t &min_align,
+                                       MemcpyTransposeInfo<size_t> &transpose_info,
+                                       AddressListCursor &in_alc, uintptr_t in_base,
+                                       AddressListCursor &out_alc, uintptr_t out_base,
+                                       size_t bytes_left, size_t max_xfer_fields,
+                                       size_t &fields_total);
+
     private:
       std::vector<GPU *> src_gpus, dst_gpus;
       std::vector<bool> dst_is_ipc;
+
+      // Mininum amount to transfer in a single quantum before returning in order to
+      // ensure forward progress
+      // TODO: make controllable
+      static constexpr size_t min_xfer_size = 4 << 20;
+      // Maximum amount to transfer in a single quantum in order to ensure other requests
+      // have a chance to make forward progress.  This should be large enough that the
+      // overhead of splitting the copy shouldn't be noticable in terms of latency (4GiB
+      // should be good here for most purposes)
+      // TODO: make controllable
+      static constexpr size_t max_xfer_size = 4ULL * 1024ULL * 1024ULL * 1024ULL;
+      static constexpr size_t max_xfer_fields = 2000;
     };
 
     class GPUIndirectChannel;
@@ -919,9 +941,46 @@ namespace Realm {
       long submit(Request **requests, long nr);
       GPU *get_gpu() const { return src_gpu; }
 
+      virtual RemoteChannelInfo *construct_remote_info() const;
+
+      virtual bool support_idindexed_fields(Memory src_mem, Memory dst_mem) const
+      {
+        return true;
+      }
+
     private:
       GPU *src_gpu;
-      // std::deque<Request*> pending_copies;
+    };
+
+    class GPURemoteChannelInfo : public SimpleRemoteChannelInfo {
+    public:
+      GPURemoteChannelInfo(NodeID _owner, XferDesKind _kind, uintptr_t _remote_ptr,
+                           const std::vector<Channel::SupportedPath> &_paths);
+
+      virtual RemoteChannel *create_remote_channel();
+
+      template <typename S>
+      bool serialize(S &serializer) const;
+
+      template <typename S>
+      static RemoteChannelInfo *deserialize_new(S &deserializer);
+
+    protected:
+      static Serialization::PolymorphicSerdezSubclass<RemoteChannelInfo,
+                                                      GPURemoteChannelInfo>
+          serdez_subclass;
+    };
+
+    class GPURemoteChannel : public RemoteChannel {
+      friend class GPURemoteChannelInfo;
+
+      GPURemoteChannel(uintptr_t _remote_ptr);
+
+    public:
+      virtual bool support_idindexed_fields(Memory src_mem, Memory dst_mem) const
+      {
+        return true;
+      }
     };
 
     class GPUfillChannel;

src/realm/cuda/cuda_memcpy.cu

@@ -139,8 +139,73 @@ memcpy_kernel_transpose(Realm::Cuda::MemcpyTransposeInfo<Offset_t> info, T *tile
 
 template <typename T, size_t N, typename Offset_t = size_t>
 static __device__ inline void
-memcpy_affine_batch(Realm::Cuda::AffineCopyPair<N, Offset_t> *info,
-                    size_t nrects, size_t start_offset = 0)
+memcpy_multi_affine_batch(Realm::Cuda::AffineCopyPair<N, Offset_t> *info, size_t nrects,
+                          size_t start_offset = 0)
+{
+  const Offset_t blk_stride = blockDim.x;
+  const Offset_t tid_global = threadIdx.x;
+
+  /* -------- iterate over copy rectangles -------- */
+  for(size_t r = 0; r < nrects; ++r) {
+    auto &cp = info[r];
+    Offset_t v = cp.volume; // elements in one field
+    Offset_t n = max(size_t(1), max(cp.src.num_fields, cp.dst.num_fields));
+
+    const T *__restrict__ src = reinterpret_cast<const T *>(cp.src.addr);
+    T *__restrict__ dst = reinterpret_cast<T *>(cp.dst.addr);
+
+    /* -------- iterate over fields handled by this block -------- */
+    for(Offset_t f = blockIdx.x; f < n; f += gridDim.x) {
+
+      Offset_t src_field_base = cp.src.num_fields > 0
+                                    ? cp.src.fields[f] * cp.src.field_stride
+                                    : f * cp.src.field_stride;
+
+      Offset_t dst_field_base = cp.dst.num_fields > 0
+                                    ? cp.dst.fields[f] * cp.dst.field_stride
+                                    : f * cp.dst.field_stride;
+
+      Offset_t off = tid_global;
+
+      while(off < v) {
+        /* -------- issue up to MAX_UNROLL loads ---------- */
+        T buf[MAX_UNROLL];
+        unsigned loaded = 0;
+
+#pragma unroll
+        for(unsigned i = 0; i < MAX_UNROLL; ++i) {
+          Offset_t idx = off + i * blk_stride;
+          if(idx >= v) {
+            break;
+          }
+
+          Offset_t src_coords[N];
+          index_to_coords<N>(src_coords, idx, cp.extents);
+          Offset_t src_lin = coords_to_index<N>(src_coords, cp.src.strides);
+          buf[i] = src[src_field_base + src_lin];
+          ++loaded;
+        }
+
+/* -------- corresponding stores ------------------ */
+#pragma unroll
+        for(unsigned i = 0; i < loaded; ++i) {
+          Offset_t idx = off + i * blk_stride;
+          Offset_t dst_coords[N];
+          index_to_coords<N>(dst_coords, idx, cp.extents);
+          Offset_t dst_lin = coords_to_index<N>(dst_coords, cp.dst.strides);
+          dst[dst_field_base + dst_lin] = buf[i];
+        }
+
+        off += loaded * blk_stride;
+      } // while off < v
+    } // for each field handled by this block
+  } // for each rect
+}
+
+template <typename T, size_t N, typename Offset_t = size_t>
+static __device__ inline void
+memcpy_affine_batch(Realm::Cuda::AffineCopyPair<N, Offset_t> *info, size_t nrects,
+                    size_t start_offset = 0)
 {
   Offset_t offset = blockIdx.x * blockDim.x + threadIdx.x - start_offset;
   const unsigned grid_stride = gridDim.x * blockDim.x;
@@ -170,8 +235,7 @@ memcpy_affine_batch(Realm::Cuda::AffineCopyPair<N, Offset_t> *info,
       for(unsigned j = 0; j < i; j++) {
         Offset_t dst_coords[N];
 
-        index_to_coords<N, Offset_t>(dst_coords,
-                                     (offset + j * grid_stride),
+        index_to_coords<N, Offset_t>(dst_coords, (offset + j * grid_stride),
                                      current_info.extents);
 
         const size_t dst_idx =
@@ -247,7 +311,7 @@ memcpy_indirect_points(Realm::Cuda::MemcpyIndirectInfo<3, Offset_t> info)
 
 template <int N, typename T, typename Offset_t = size_t>
 static __device__ inline void
-memfill_affine_batch(const Realm::Cuda::AffineFillInfo<N, Offset_t>& info)
+memfill_affine_batch(const Realm::Cuda::AffineFillInfo<N, Offset_t> &info)
 {
   Offset_t offset = blockIdx.x * blockDim.x + threadIdx.x;
   const unsigned grid_stride = gridDim.x * blockDim.x;
@@ -277,15 +341,24 @@ memfill_affine_batch(const Realm::Cuda::AffineFillInfo<N, Offset_t>& info)
   }
 }
 
-#define MEMCPY_TEMPLATE_INST(type, dim, offt, name)                            \
-  extern "C" __global__ __launch_bounds__(256, 4) void                         \
-      memcpy_affine_batch##name(Realm::Cuda::AffineCopyInfo<dim, offt> info) { \
-    memcpy_affine_batch<type, dim, offt>(info.subrects, info.num_rects);       \
+#define MEMCPY_MULTI_TEMPLATE_INST(type, dim, offt, name)                                \
+  extern "C" __global__ __launch_bounds__(256, 4) void multi_affine_batch##name(         \
+      Realm::Cuda::AffineCopyInfo<dim, offt> info)                                       \
+  {                                                                                      \
+    memcpy_multi_affine_batch<type, dim, offt>(info.subrects, info.num_rects);           \
+  }
+
+#define MEMCPY_TEMPLATE_INST(type, dim, offt, name)                                      \
+  extern "C" __global__ __launch_bounds__(256, 4) void memcpy_affine_batch##name(        \
+      Realm::Cuda::AffineCopyInfo<dim, offt> info)                                       \
+  {                                                                                      \
+    memcpy_affine_batch<type, dim, offt>(info.subrects, info.num_rects);                 \
   }
 
 #define FILL_TEMPLATE_INST(type, dim, offt, name)                                        \
   extern "C" __global__ void fill_affine_batch##name(                                    \
-      Realm::Cuda::AffineFillInfo<dim, offt> info) {                                     \
+      Realm::Cuda::AffineFillInfo<dim, offt> info)                                       \
+  {                                                                                      \
     memfill_affine_batch<dim, type, offt>(info);                                         \
   }
 
@@ -311,16 +384,17 @@ memfill_affine_batch(const Realm::Cuda::AffineFillInfo<N, Offset_t>& info)
 
 #define INST_TEMPLATES(type, sz, dim, off)                                               \
   MEMCPY_TEMPLATE_INST(type, dim, off, dim##D_##sz)                                      \
+  MEMCPY_MULTI_TEMPLATE_INST(type, dim, off, dim##D_##sz)                                \
   FILL_TEMPLATE_INST(type, dim, off, dim##D_##sz)                                        \
   FILL_LARGE_TEMPLATE_INST(type, dim, off, dim##D_##sz)                                  \
   MEMCPY_INDIRECT_TEMPLATE_INST(int, type, dim, off, dim##D_##sz##32)                    \
   MEMCPY_INDIRECT_TEMPLATE_INST(long long, type, dim, off, dim##D_##sz##64)
 
-#define INST_TEMPLATES_FOR_TYPES(dim, off)                                     \
-  INST_TEMPLATES(unsigned char, 8, dim, off)                                   \
-  INST_TEMPLATES(unsigned short, 16, dim, off)                                 \
-  INST_TEMPLATES(unsigned int, 32, dim, off)                                   \
-  INST_TEMPLATES(unsigned long long, 64, dim, off)                             \
+#define INST_TEMPLATES_FOR_TYPES(dim, off)                                               \
+  INST_TEMPLATES(unsigned char, 8, dim, off)                                             \
+  INST_TEMPLATES(unsigned short, 16, dim, off)                                           \
+  INST_TEMPLATES(unsigned int, 32, dim, off)                                             \
+  INST_TEMPLATES(unsigned long long, 64, dim, off)                                       \
   INST_TEMPLATES(uint4, 128, dim, off)
 
 #define INST_TEMPLATES_FOR_DIMS()                                                        \

src/realm/cuda/cuda_memcpy.h

@@ -31,10 +31,15 @@ namespace Realm {
 
     template <size_t N, typename Offset_t = size_t>
     struct alignas(8) AffineSubRect {
+      using FieldID = int;
       // Extent of the ND array
       Offset_t strides[N - 1];
       // Address of the ND array
       uintptr_t addr;
+
+      const FieldID *fields;
+      size_t num_fields;
+      Offset_t field_stride;
     };
 
     template <size_t N, typename Offset_t = size_t>

src/realm/cuda/cuda_module.cc

@@ -1142,7 +1142,8 @@ namespace Realm {
     }
 
     void GPU::launch_batch_affine_kernel(void *copy_info, size_t dim, size_t elem_size,
-                                         size_t volume, GPUStream *stream)
+                                         size_t volume, bool mutlfield_optimized,
+                                         GPUStream *stream)
     {
       size_t log_elem_size = std::min(static_cast<size_t>(ctz(elem_size)),
                                       CUDA_MEMCPY_KERNEL_MAX2_LOG2_BYTES - 1);
@@ -1151,10 +1152,13 @@ namespace Realm {
       assert(dim <= REALM_MAX_DIM);
       assert(dim >= 1);
 
-      // TODO: probably replace this
-      // with a better data-structure
-      GPUFuncInfo &func_info = batch_affine_kernels[dim - 1][log_elem_size];
-      launch_kernel(func_info, copy_info, volume, stream);
+      if(!mutlfield_optimized) {
+        GPUFuncInfo &func_info = batch_affine_kernels[dim - 1][log_elem_size];
+        launch_kernel(func_info, copy_info, volume, stream);
+      } else {
+        GPUFuncInfo &func_info = multi_batch_affine_kernels[dim - 1][log_elem_size];
+        launch_kernel(func_info, copy_info, volume, stream);
+      }
     }
 
     const GPU::CudaIpcMapping *GPU::find_ipc_mapping(Memory mem) const
@@ -2095,6 +2099,15 @@ namespace Realm {
               0));
           batch_affine_kernels[d - 1][log_bit_sz] = func_info;
 
+          std::snprintf(name, sizeof(name), "multi_affine_batch%uD_%u", d, bit_sz);
+          CHECK_CU(CUDA_DRIVER_FNPTR(cuModuleGetFunction)(&func_info.func, device_module,
+                                                          name));
+
+          CHECK_CU(CUDA_DRIVER_FNPTR(cuOccupancyMaxPotentialBlockSize)(
+              &func_info.occ_num_blocks, &func_info.occ_num_threads, func_info.func, 0, 0,
+              0));
+          multi_batch_affine_kernels[d - 1][log_bit_sz] = func_info;
+
           std::snprintf(name, sizeof(name), "fill_affine_large%uD_%u", d, bit_sz);
           CHECK_CU(CUDA_DRIVER_FNPTR(cuModuleGetFunction)(&func_info.func, device_module,
                                                           name));

src/realm/inst_layout.h

@@ -227,6 +227,7 @@ namespace Realm {
 
     using FieldMap = std::map<FieldID, FieldLayout>;
     FieldMap fields;
+    bool idindexed_fields{false};
   };
 
   REALM_PUBLIC_API
@@ -421,6 +422,9 @@ namespace Realm {
     IndexSpace<N, T> space;
     std::vector<InstancePieceList<N, T>> piece_lists;
 
+    // Pre-computed dimension ordering for idindexed_fields
+    std::vector<int> preferred_dim_order;
+
     static Serialization::PolymorphicSerdezSubclass<InstanceLayoutGeneric,
                                                     InstanceLayout<N, T>>
         serdez_subclass;