Custom Memory Resource and Extensible Memory Management

[sungwoo-XCENA]

1. Motivation

Velox's current memory management assumes that CPU DRAM is the only resource to be managed. With heterogeneous compute extensions like the cuDF integration, this assumption no longer holds: GPUs introduce their own memory, and emerging tiered memory technologies like CXL introduce additional resources with distinct latency and bandwidth characteristics. Today the cuDF extension's GPU memory allocations are unaccounted and invisible to the rest of the engine.

One could make Velox transparently spill from a fast resource to a slow one (e.g., GPU → host DRAM). However, transparent spilling across tiers introduces hard-to-diagnose performance regressions, because the access-latency cliff between tiers is large and the spill point is invisible to the operator that triggered it. This proposal takes the position that which resource a given allocation lands on is a first-class scheduling decision, not something the system should make implicitly.

2. Goal

Allow MemoryManager to host multiple memory resources side-by-side, where each resource carries its own allocator and its own arbitrator. An extension should be able to register a new resource by supplying a single object that bundles three things: a tag, an allocator, and an arbitrator. A single registration entry feels right because the three are logically coupled — an arbitrator manages capacity over a specific allocator's address space, and the tag is what pools use to opt into that resource.

The current code path must continue to work unchanged when no extension registers a custom resource.

Out of scope for this proposal:

• Cross-resource arbitration (e.g., evict from GPU to host DRAM under pressure). Per-resource arbitration only.
• Automatic placement decisions (e.g., place this in secondary tier memory automatically).

3. Design Sketch

The shape we have in mind follows DriverAdapter. The idea is to extend pluggability to the allocator side and bundle allocator + arbitrator under one registration entry.

CustomMemoryResource

struct CustomMemoryResource {
    /// Identifier used by pools to opt into this resource. Must be unique
    /// across registered resources.
    std::string tag;

    /// Max capacity for the root MemoryPool
    int64_t maxCapacity;

    /// Allocator that backs this resource. Constructed by the extension.
    std::shared_ptr<MemoryAllocator> allocator;

    /// Arbitrator that owns capacity decisions for pools on this resource.
    /// May be null if the resource does not require arbitration.
    std::shared_ptr<MemoryArbitrator> arbitrator;

    /// Factory invoked by QueryCtx::Builder to attach a per-query reclaimer
    /// to each root pool created on this resource. May be null for resources
    /// that do not support reclaim.
    std::function<std::unique_ptr<MemoryReclaimer>(QueryCtx*)> reclaimerFactory;
};

The extension would construct the allocator and arbitrator with whatever configuration it needs, then hand the bundle to MemoryManager:

auto allocator = std::make_shared<GpuMemoryAllocator>(rmm);
auto arbitrator = std::make_shared<GpuArbitrator>(gpuCapacity);
auto reclaimerFactory = [](core::QueryCtx* ctx) { return GpuQueryReclaimer::create(ctx); };
memoryManager()->registerCustomResource(
    {"gpu", std::move(allocator), std::move(arbitrator), std::move(reclaimerFactory)});

Example end-to-end usage

// 1. Once at extension startup. We are assuming that memoryManager is instantiated.
void registerGpuMemoryResource(rmm::mr::device_memory_resource* rmm) {
    auto allocator = std::make_shared<GpuMemoryAllocator>(rmm);
    auto arbitrator = std::make_shared<GpuArbitrator>(gpuCapacity);
    auto reclaimerFactory = [](core::QueryCtx* ctx) { return GpuQueryReclaimer::create(ctx); };
    memoryManager()->registerCustomResource(
        {"gpu", std::move(allocator), std::move(arbitrator), std::move(reclaimerFactory)});
}

// 2. Per-query: a cuDF-aware caller declares the GPU root via the builder.
// Inside QueryCtx::Builder::build,
for (auto& mr : memoryManager()->customResources()) {
    auto pool = memoryManager()->addRootPool(
            fmt::format("query.{}.{}", queryId_, mr.tag),
            mr.maxCapacity,
            mr.reclaimerFactory(queryCtx.get()),   // factory now has queryCtx
            std::nullopt,
            mr.tag);

    queryCtx->addCustomPool(std::move(pool));
}

// 3. MemoryPool (or MemoryPoolImpl) constructor searches for the custom resource if optional resourceTag is passed.              
MemoryPoolImpl::MemoryPoolImpl(..., const Options& options)            
    : ...,                                                             
      manager_{memoryManager},      
      allocator_{options.resource ? options.resource->allocator.get()  
                                  : manager_->allocator()},            
      arbitrator_{options.resource ? options.resource->arbitrator.get()
                                   : manager_->arbitrator()},          
      ...
                                                              
// 4. cuDF operators allocate through the GPU root's children. The leaf pool
//    routes to GpuMemoryAllocator because that is what its root resolved at
//    construction time.
auto opPool = queryCtx->customPool("gpu")->addAggregateChild("HashJoinBuild");
auto bufPool = opPool->addLeafChild("buffers");
void* buf = bufPool->allocate(numBytes);  // -> GpuMemoryAllocator::allocateBytes

4. Open Questions

4.1 How should we handle query failure?

Maybe I should look more into this but as QueryCtx now knows custom memory pools and the memory reclaimers, we can iterate over them to reclaim outstanding memory. I'm wondering if there would be any order preservation needed when we reclaim memory resources.

4.2 Where should we draw the scope of this project if we proceed?

I think I could plumb through the frameworks with unit tests and the mock custom memory resources. Then, I guess it's the extension developer's turns to implement the custom memory resource for their extensions.

Please review the proposal and tell me if there are things I need to consider before start working.

Best regards,
Sungwoo

[srsuryadev]

cc: @kgpai @xiaoxmeng
cc: @mbasmanova @pedroerp

[devavret]

I am heavily in favour of adding tiered tracking and arbitration to Velox. I have some thoughts:

Out of scope for this proposal:

* Cross-resource arbitration (e.g., evict from GPU to host DRAM under pressure). Per-resource arbitration only.

I don't think we can escape this. Arbitration currently consists of reservation management AND spilling. If GPU memory is full, it must spill. And it's best to spill to host memory. This means GPU memory arbitration triggers allocation from CPU memory pools. And this has the potential to cause some locking issue somewhere.

Additionally, all CPU memory isn't the same as far as transfers from GPU is concerned. Pinned memory (allocated from cudaHostAlloc) is a better (faster) target than pageable memory. I'd like to request that we get the ability to overwrite the CPU memory resources so that GPU execution mode uses a custom CPU allocator.

[sungwoo-XCENA]

Thank you @devavret for the feedbacks!

For your first feedback on cross-resource arbitration, if we are going to include that in this scope, I think it will be the extension's responsibility to orchestrate cross-resource arbitration. Extensions will have all they need for orchestration inside QueryCtx, and also I don't think that the memory manager or shared arbitrator can orchestrate that automatically for you.

For your second feedback, why not use CustomMemoryResource for CPU Pinned Memory? Rather than overwriting the existing CPU memory management, which I think needs to be preserved as much as possible to guarantee backward compatibility, it would be more flexible and extensible to use CustomMemoryResource for that purpose.

[mbasmanova]

CC: @bikramSingh91 @tanjialiang @xiaoxmeng

[pedroerp]

@sungwoo-XCENA thanks for starting this proposal. Having memory management extensible with a user-provided adapter makes sense to me, so that users can provide their own extensions of memory management.

Is the idea here that we mostly want to do this from a tracking purpose, so that, say, Vectors and Operators can carry their own memory pools built on CXL (or GPU) memory, or that we actually want operators to be able to spill across the memory hierarchy?

My only concern with the latter is that the currently spill algorithms we support are likely tied to the memory<->SSD/disk concept, so I'm not sure they will translate to other memory hierarchies. Meaning it would likely require different spilling strategies to begin with, and these are not very self-contained today. But it would be interesting to understand what these are.

[sungwoo-XCENA]

Hi @pedroerp, I'll reply to your questions/opinions first and then suggest a solution for cross-resource arbitration.

Is the idea here that we mostly want to do this from a tracking purpose, so that, say, Vectors and Operators can carry their own memory pools built on CXL (or GPU) memory, or that we actually want operators to be able to spill across the memory hierarchy?

I was going to start with just tracking memory usage and arbitration within a single memory resource, but as @devavret pointed out, we'll need cross-resource arbitration eventually, so we'd better have a framework for that from the get-go.

My only concern with the latter is that the currently spill algorithms we support are likely tied to the memory<->SSD/disk concept, so I'm not sure they will translate to other memory hierarchies. Meaning it would likely require different spilling strategies to begin with, and these are not very self-contained today. But it would be interesting to understand what these are.

Good point. So we'll need a separate injection point for custom memory resources to provide their own spilling strategies. I have an idea about this. It's not solid and probably need some more investigation but I will present in the section below.

Suggested Solution

Ok, so I've given some thought to how to do cross-resource arbitration. CustomMemoryResource has a MemoryReclaimer factory that generates a MemoryReclaimer for your custom resource. The derived MemoryReclaimer could hold a reference to a root memory pool that belongs to another memory resource used as the spill target. The derived MemoryReclaimer could then allocate space in the spill target before emptying its own memory. Now, the spill target may or may not have enough capacity to serve the allocation request. If it doesn't, the spill target's memory pool will automatically trigger the memory arbitration process, and we'll end up with enough space.

So for example, if we were to implement GpuMemoryResource, the MemoryReclaimer factory would look like this:

class GpuMemoryReclaimer : public memory::MemoryReclaimer {
 public:
  static std::unique_ptr<MemoryReclaimer> create(core::QueryCtx* queryCtx) {
    // Capture the Velox core CPU root pool as the spill target.
  }

  uint64_t reclaim(MemoryPool* gpuPool, uint64_t targetBytes, ...) override {
    // 1. Ask gpuPool for actually-spillable bytes; cap targetBytes.
    // 2. allocate(targetBytes) on a leaf child of the core CPU root.
    //    -> this implicitly triggers Velox's arbitrator if it's tight.
    // 3. cudaMemcpyAsync device -> host.
    // 4. Free the GPU buffers; return bytes actually reclaimed.
  }

 private:
  std::shared_ptr<MemoryPool> spillRoot_;  // Velox core CPU root
};

That said, I think this can be risky because it affects Velox's current memory management scheme, and I'm also not sure what kind of chain reaction this could cause in Velox's arbitration process. So for safety reasons, what I'd suggest is that custom memory resources spill to either another custom memory resource or disk. For instance, we could create a PinnedCpuMemoryResource, which would serve as the backing memory resource for GPU spills outside Velox's core memory resource. Then GpuMemoryResource's MemoryReclaimer can reference PinnedCpuMemoryResource instead of Velox's core memory resource, so that it can be managed independently. PinnedCpuMemoryResource's MemoryReclaimer can spill to disk, reusing the current spill library.

So in code, it would look something like this:

void registerHeterogeneousMemoryResources(...) {
  // Order matters: pinned-cpu must exist before gpu's factory runs.
  memoryManager()->registerCustomResource({
      .tag = "pinned-cpu",
      .allocator = PinnedCpuAllocator,
      .arbitrator = PinnedCpuArbitrator,
      .reclaimerFactory = [](QueryCtx* ctx) {
        // spills down to disk.
        return PinnedCpuToDiskReclaimer::create(ctx);
      },
  });

  memoryManager()->registerCustomResource({
      .tag = "gpu",
      .allocator = GpuMemoryAllocator,
      .arbitrator = GpuArbitrator,
      .reclaimerFactory = [](QueryCtx* ctx) {
        // GPU spills into the sibling pinned-cpu pool.
        return GpuToPinnedCpuReclaimer::create(ctx->customPool("pinned-cpu"));
      },
  });
}

class GpuToPinnedCpuReclaimer : public MemoryReclaimer {
 public:
  static std::unique_ptr<MemoryReclaimer> create(
      std::shared_ptr<MemoryPool> spillTarget);

  uint64_t reclaim(MemoryPool* gpuPool, uint64_t targetBytes, ...) override {
    // 1. Cap targetBytes by gpuPool's reclaimableBytes().
    // 2. allocate() on spillTarget_ (pinned-cpu).
    //    -> triggers pinned-cpu arbitration if tight,
    //       which may recursively spill pinned-cpu to disk.
    // 3. cudaMemcpyAsync device -> pinned host.
    // 4. Free GPU buffers; return bytes reclaimed.
  }

 private:
  std::shared_ptr<MemoryPool> spillTarget_;  // pinned-cpu root, not core CPU
};

class PinnedCpuToDiskReclaimer : public MemoryReclaimer {
 public:
  static std::unique_ptr<MemoryReclaimer> create(QueryCtx* queryCtx);

  uint64_t reclaim(MemoryPool* pinnedPool, uint64_t targetBytes, ...) override {
    // Reuse Velox's existing file-based spill machinery to push
    // pinned-host buffers to disk. Return bytes reclaimed.
  }
};

What do you think? Does this seem like a feasible approach? Let me know in the comments.

[sungwoo-XCENA]

I'm waiting for a consensus on this. Meanwhile I implemented the proposal in my fork. Let me know if I can make a PR.

[mbasmanova]

@sungwoo-XCENA Please, go ahead and submit a PR. Thanks!

CC: @tanjialiang @bikramSingh91 @Yuhta @xiaoxmeng

[devavret]

Yes please. I was going to ask if you already had a branch with this but didn't want to presume. I can test with it and let you know if it works for us.

[sungwoo-XCENA]

Thank you @mbasmanova @devavret. I submitted a PR #17529.