Introducing MemoryTier

[sungwoo-XCENA]

Motivation

In PR #17529, we introduced CustomMemoryResource where an extension can register and use custom memory allocator and arbitrator. In PR #17960, we implemented CxlMemoryResource using the CustomMemoryResource, and in PR #18000, we implemented CXL-aware Hash Aggregation using CxlMemoryResource. The PR #17529 and #17960 are merged, and PR #18000 is being reviewed.

So, that's the history of development for those who weren't tracking. When I was developing PR #18000, I felt inconvenient of using CustomMemoryResource because after registering a memory resource, you need to keep the resource tag somewhere to use it in the operator. Either you have to sink that tag through QueryCtx or you have to statically define a field in QueryConfig for operator to be able to retrieve that tag. This is not a great idea because, first extensions that wants to add custom memory resource will have to add tag fields all the time, and second it strongly couples the custom memory resource and the extension that wants to use the memory resource.

What we need is a hierarchy of memory resources, based on the memory latency from the current compute. We can picture it like CPU cache hierarchy, where we have L1/L2/L3 and the compute uses L1 first, then L2, then L3 as each layer fills up. This hierarchy should also support operator-level scope because extensions that use DriverAdapter would want different hierarchy than the normal paths.

Design

That's why I'm suggesting two things:

1. MemoryTier. It's a ScopedRegistry-based registry that keeps resource tags in each memory tier. It's an array of arrays, with each slot holds list of memory resources that are considered same tier. For example, for cuDF extension, L1 would be GPU memory managed by RMM, L2 would be pinned host memory, L3 would be the default Velox memory or CXL memory or both in the same tier.
2. Operator-level Scoped Registry. Today, Velox only has ScopedRegistry globally and in QueryCtx. However, with extensions replacing operators with DriverAdapter, each operator may need different MemoryTiers. I think OperatorCtx is a good place to introduce another ScopedRegistry, so that we can override query-level configuration.

Implementation

1. MemoryTier

class MemoryTier {
  vector<vector<string>> tiers_;              // tiers_[level] = tags at that tier

  size_t numTiers() const;
  const vector<string>& tierAt(size_t level) const;   // just the tags
};

constexpr string_view kMemoryTierKey = "memoryTier";

2. Operator-level registry

class OperatorCtx {
  // same shape as QueryCtx::registries_ : key -> {shared_ptr<void>, type_index}
  Synchronized<F14FastMap<string, RegistryEntry>> registries_;

  template <typename T>
  bool setRegistry(string_view key, shared_ptr<T> reg, bool overwrite=false);  // local override

  template <typename T>
  shared_ptr<T> registry(string_view key) const {
    if (auto local = lookupLocal<T>(key)) return local;          // operator scope
    return driverCtx_->task->queryCtx()->registry<T>(key);       // fall through to query (-> global)
  }
};

3. DriverAdapter Integration

// DriverAdapter wants a different hierarchy for just this operator:
opCtx->setRegistry<MemoryTier>(kMemoryTierKey,
    make_shared<MemoryTier>({{"rmm"}, {"pinnedHost"}, {"cxl","dram"}}));

The operator fetches its effective MemoryTier and picks the next tier below the one it is filling.

MemoryPool* relocationTarget(size_t currentLevel) {
  auto tiers = operatorCtx_->registry<MemoryTier>(kMemoryTierKey);   // tiered lookup
  if (!tiers || currentLevel + 1 >= tiers->numTiers()) return nullptr;
  for (tag : tiers->tagsAt(currentLevel + 1))
    if (auto* p = customPool(tag)) return p;    // reuse existing tag -> leaf-pool wiring
  return nullptr;
}

What do you think? I appreciate all the feedbacks!

[bdice]

Thanks for starting this discussion! There is a library being developed by some colleagues at NVIDIA called cuCascade (https://github.com/NVIDIA/cuCascade). It is also designed for memory tiers, though it incorporates some additional features like topology discovery that helps with defining the hierarchy and making performance-aware choices. Maybe there are ideas worth borrowing here. We have considered similar approaches in Velox but no implementation to my knowledge. @devavret may have more to say on this topic.

In general I am cautious about global registries, it is easy to lose flexibility in atypical use cases. I’m glad to hear there will be some customization points (per-operator, etc.).

[sungwoo-XCENA]

Thanks for the feedback. cuCascade would be a good reference point for this project. I will see what I can borrow.

[mbasmanova]

CC: @xiaoxmeng @tanjialiang

[sungwoo-XCENA]

I created two PRs that implemented this proposal.

PR #18121 creates Operator-scoped registry.
PR #18122 creates MemoryTier.

Wiring MemoryTier with Operator-scoped registry could be a third PR or in-place with PR #18122.

I would appreciate any feedback. Thank you!

[sungwoo-XCENA]

Update.

PR #18121 is closed now. We reached an agreement that Operator-level Memory Tier isn't necessary as the compute backend will be switched at plan fragment level.