ib_memory.h

Go to the documentation of this file.

/*
 * Copyright 2025 Stanford University, NVIDIA Corporation
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
// IB (Intermediate Buffer) Memory implementations for Realm
 
#ifndef REALM_IB_MEMORY_H
#define REALM_IB_MEMORY_H
 
#include "realm/realm_config.h"
 
#include "realm/mem_impl.h"
 
#include <limits>
 
namespace Realm {
 
  // IBMemory::do_alloc rounds every allocation up to this alignment; planner-side
  // size math must account for it so any plan the planner accepts is allocatable.
  static constexpr size_t IB_ALLOC_ALIGNMENT = 256;
 
  // round n up to the next multiple of mult (matches the semantics of the
  // file-local roundup() in inst_impl.cc)
  inline size_t ib_align_up(size_t n, size_t mult = IB_ALLOC_ALIGNMENT)
  {
    if(mult == 0) {
      return n;
    }
    size_t leftover = n % mult;
    if(leftover == 0) {
      return n;
    }
    size_t extra = mult - leftover;
    if(n > (std::numeric_limits<size_t>::max() - extra)) {
      return std::numeric_limits<size_t>::max();
    }
    return n + extra;
  }
 
  // a simple memory used for intermediate buffers in dma system
  class REALM_INTERNAL_API_EXTERNAL_LINKAGE IBMemory : public MemoryImpl {
  public:
    IBMemory(RuntimeImpl *_runtime_impl, Memory _me, size_t _size, MemoryKind _kind,
             Memory::Kind _lowlevel_kind, void *prealloc_base, NetworkSegment *_segment);
 
    virtual ~IBMemory();
 
    // old-style allocation used by IB memories
    virtual off_t alloc_bytes_local(size_t size);
    void free_bytes_local(off_t offset, size_t size, TimeLimit work_until);
 
    virtual void *get_direct_ptr(off_t offset, size_t size);
 
    // not used by IB memories
    virtual AllocationResult allocate_storage_immediate(RegionInstanceImpl *inst,
                                                        bool need_alloc_result,
                                                        bool poisoned,
                                                        TimeLimit work_until);
 
    virtual void release_storage_immediate(RegionInstanceImpl *inst, bool poisoned,
                                           TimeLimit work_until);
 
    virtual void get_bytes(off_t offset, void *dst, size_t size);
    virtual void put_bytes(off_t offset, const void *src, size_t size);
 
    // attempts to allocate one or more IBs - either all succeed or all fail
    bool attempt_immediate_allocation(NodeID requestor, uintptr_t req_op, size_t count,
                                      const size_t *sizes, off_t *offsets);
 
    // enqueues a batch of PendingIBRequests to be satisfied as soon as possible
    void enqueue_requests(PendingIBRequests *reqs, TimeLimit work_until);
 
    void free_multiple(size_t count, const off_t *offsets, const size_t *sizes,
                       TimeLimit work_until);
 
  protected:
    // these must be called with the mutex held
    off_t do_alloc(size_t size);
    void do_free(off_t offset, size_t size);
    PendingIBRequests *satisfy_pending_reqs();
    void forward_satisfied_reqs(PendingIBRequests *reqs, TimeLimit work_until);
 
    Mutex mutex; // protection for resizing vectors
    std::map<off_t, off_t> free_blocks;
    char *base;
    PendingIBRequests *ibreq_head;
    PendingIBRequests **ibreq_tail;
  };
 
  // helper routine to free IB whether it is local or remote
  void free_intermediate_buffer(Memory mem, off_t offset, size_t size,
                                TimeLimit work_until);
 
  // active messages related to IB allocation/release
 
  struct RemoteIBAllocRequestSingle {
    Memory memory;
    size_t size;
    uintptr_t req_op;
    unsigned req_index;
    bool immediate;
 
    static void handle_message(NodeID sender, const RemoteIBAllocRequestSingle &args,
                               const void *data, size_t msglen, TimeLimit work_until);
  };
 
  struct RemoteIBAllocRequestMultiple {
    NodeID requestor;
    unsigned count, first_index, curr_index;
    uintptr_t req_op;
    bool immediate;
 
    static void handle_message(NodeID sender, const RemoteIBAllocRequestMultiple &args,
                               const void *data, size_t msglen, TimeLimit work_until);
  };
 
  struct RemoteIBAllocResponseSingle {
    uintptr_t req_op;
    unsigned req_index;
    off_t offset;
 
    static void handle_message(NodeID sender, const RemoteIBAllocResponseSingle &args,
                               const void *data, size_t msglen, TimeLimit work_until);
  };
 
  struct RemoteIBAllocResponseMultiple {
    uintptr_t req_op;
    unsigned count, first_index;
 
    static void handle_message(NodeID sender, const RemoteIBAllocResponseMultiple &args,
                               const void *data, size_t msglen, TimeLimit work_until);
  };
 
  struct RemoteIBReleaseSingle {
    Memory memory;
    size_t size;
    off_t offset;
 
    static void handle_message(NodeID sender, const RemoteIBReleaseSingle &args,
                               const void *data, size_t msglen, TimeLimit work_until);
  };
 
  struct RemoteIBReleaseMultiple {
    unsigned count;
 
    static void handle_message(NodeID sender, const RemoteIBReleaseMultiple &args,
                               const void *data, size_t msglen, TimeLimit work_until);
  };
 
}; // namespace Realm
 
#endif