Legion Runtime Class #14: Region Tree State and Contexts
TweetOverview
In this class we'll talk about the region trees, contexts, and some utility data structures.
About Legion Runtime Class
These notes are closely based on the set of Legion Runtime Class videos produced by the Legion developers. They are my own notes and code walks, and any errors or things that are just plain wrong represent my own mistakes.
Today's notes are based on the following video:
Region Tree Creation
A logical region tree is created by taking the cross product of an index space and a field space. The logical region tree is identified by a 3-tuple consisting of the index space, field space, and a unique region tree ID. This is where a logical region is created:
LogicalRegion create_logical_region(Context ctx, IndexSpace index,
FieldSpace fields);
The you can see that the logical region structure is a light-weight handle. Below is the 3-tuple:
class LogicalRegion {
...
private:
// These are private so the user can't just arbitrarily change them
RegionTreeID tree_id;
IndexSpace index_space;
FieldSpace field_space;
};
When a logical region is created first a unique ID is generated, and then generate the handle. After we create the handle we ask the region tree forest to create the region for us.
LogicalRegion Runtime::create_logical_region(Context ctx,
IndexSpace index_space, FieldSpace field_space)
{
RegionTreeID tid = get_unique_tree_id();
LogicalRegion region(tid, index_space, field_space);
forest->create_logical_region(region);
// Register the creation of a top-level region with the context
ctx->register_region_creation(region);
return region;
}
Similar to what we've seen before, we use the overloaded create_node method. Here we are creating a new tree so there is no parent because this node will be the root of the tree.
void RegionTreeForest::create_logical_region(LogicalRegion handle)
{
// We do need the autolock here since we're going to be making
// region tree nodes
AutoLock f_lock(forest_lock,1,false/*exclusive*/);
create_node(handle, NULL/*parent*/);
}
In create_node a new PartitionNode is created. The index space and field space are referred to as the row and column sources, respectively. With references to them we create the new partition node.
Next we lookup the partition node in the part_nodes data structure. If it's found then we are done. Like previous instances of create_node we've seen, the creation can race. If we don't find it, then we add it to part_nodes, and call reserve_contexts, which we'll discuss later.
PartitionNode* RegionTreeForest::create_node(LogicalPartition p,
RegionNode *parent)
//--------------------------------------------------------------------------
{
IndexPartNode *row_src = get_node(p.index_partition);
FieldSpaceNode *col_src = get_node(p.field_space);
PartitionNode *result = new PartitionNode(p, parent, row_src,
col_src, this);
// Special case here in case multiple clients attempt
// to make the node at the same time
{
// Hole the lookup lock when modifying the lookup table
AutoLock l_lock(lookup_lock);
std::map<LogicalPartition,PartitionNode*>::const_iterator it =
part_nodes.find(p);
if (it != part_nodes.end())
{
// It already exists, delete our copy and
// return the one that has already been made
delete result;
return it->second;
}
// Now make sure that the node has the right number of contexts
// before we put it on the map and release the look-up lock.
// Note that this is only safe because the runtime ups the
// total_context count in allocate_context before calling the
// resize_node_contexts funciton.
result->reserve_contexts(runtime->get_context_count());
// Now we can put the node in the map
part_nodes[p] = result;
}
// Now we can make the other ways of accessing the node available
row_src->add_instance(result);
parent->add_child(result);
return result;
}
Looking at the high-level interface to Legion we see calls that can partition index spaces, but there are no calls to partition a logical partition. The reason for this is that when an index space is partitioned, the logical regions created from that index space automatically have those partitions created.
Here's how we lookup a logical partition.
LogicalPartition Runtime::get_logical_partition(Context ctx,
LogicalRegion parent, IndexPartition handle)
{
return forest->get_logical_partition(parent, handle);
}
In some cases, there is no lookup. We can just construct the handle and return it.
LogicalPartition RegionTreeForest::get_logical_partition(
LogicalRegion parent, IndexPartition handle)
//--------------------------------------------------------------------------
{
// No lock needed for this one
return LogicalPartition(parent.tree_id, handle, parent.field_space);
}
In other cases we may need to do some additional work. Here we want to lookup by color, so we need to first examine the forest to find elements with the specified color, but then can create the handle from the results.
LogicalPartition RegionTreeForest::get_logical_partition_by_color(
LogicalRegion parent, Color c)
//--------------------------------------------------------------------------
{
AutoLock f_lock(forest_lock,1,false/*exclusive*/);
RegionNode *parent_node = get_node(parent);
IndexPartNode *index_node = parent_node->row_source->get_child(c);
LogicalPartition result(parent.tree_id, index_node->handle,
parent.field_space);
return result;
}
Recall that index and field spaces were created eagerly. Region trees are created lazily. This is obvious from the get_logical_partition call for instance where we construct a handle without even ensuring that the tree has the relevant elements.
The RegionNode and PartitionNode that make up the elements in the region tree are both subclasses of RegionTreeNode. A RegionNode contains a lot of methods (not shown), as well as pointers to other parts of the tree for its index space and field space. The color maps track information related to children, which we'll discuss in a later talk.
class RegionNode : public RegionTreeNode {
public:
RegionNode(LogicalRegion r, PartitionNode *par, IndexSpaceNode *row_src,
FieldSpaceNode *col_src, RegionTreeForest *ctx);
RegionNode(const RegionNode &rhs);
...
public:
const LogicalRegion handle;
PartitionNode *const parent;
IndexSpaceNode *const row_source;
protected:
std::map<Color,PartitionNode*> color_map;
std::map<Color,PartitionNode*> valid_map;
The PartitionNode is similar, and includes extra information like if this partition is disjoint.
class PartitionNode : public RegionTreeNode {
public:
PartitionNode(LogicalPartition p, RegionNode *par,
IndexPartNode *row_src, FieldSpaceNode *col_src,
RegionTreeForest *ctx);
PartitionNode(const PartitionNode &rhs);
public:
const LogicalPartition handle;
RegionNode *const parent;
IndexPartNode *const row_source;
const bool disjoint;
protected:
std::map<Color,RegionNode*> color_map;
std::map<Color,RegionNode*> valid_map;
The RegionTreeNode contains state that is common to both types of region tree nodes. The context shown below isn't the same type of context we've seen in other areas, its just an overloaded term here. The NodeMask tracks which nodes have been created or destroyed in the cluster. The reservation structure restricts concurrent access to these data structures.
The logical and physical state structures are the most important. These store all of the state information needed for performing logical and physical analysis. The legion stack structure is a special structure used that allow expansion without moving already allocated data.
class RegionTreeNode {
public:
RegionTreeNode(RegionTreeForest *ctx, FieldSpaceNode *column);
virtual ~RegionTreeNode(void);
...
public:
RegionTreeForest *const context;
FieldSpaceNode *const column_source;
public:
NodeMask creation_set;
NodeMask destruction_set;
protected:
Reservation node_lock;
LegionStack<LogicalState,MAX_CONTEXTS,DEFAULT_CONTEXTS> logical_states;
LegionStack<PhysicalState,MAX_CONTEXTS,DEFAULT_CONTEXTS> physical_states;
protected:
LegionMap<SemanticTag,SemanticInfo>::aligned semantic_info;
};
So what exactly are these contexts that we keep talking about. They are RegionTreeContext. Each time we start a task that is not a leaf we create a new context object. It's a handle to index into those logical and physical states. We can see below that a region tree context is really a wrapper around an integer.
class RegionTreeContext {
public:
RegionTreeContext(void)
: ctx(-1) { }
RegionTreeContext(ContextID c)
: ctx(c) { }
public:
inline bool exists(void) const { return (ctx >= 0); }
inline ContextID get_id(void) const
{
return ContextID(ctx);
}
private:
int ctx;
};
The runtime maintains a pool of region tree contexts which are allocated via allocate_context when a task starts. If there aren't enough, then we'll allocate more. Several are allocated at once because it can be an expensive operation. Recall that the region nodes are indexed by the context. When we allocate a brand new context, we also iterate over all of the nodes in the every region tree and guarantee that there is space allocated for the new context. That's why it is expensive. The expensive call occurs at the end of the method by calling forest->resize_node_contexts(current_contexts);.
void Runtime::allocate_context(SingleTask *task)
{
// Try getting something off the list of available contexts
AutoLock avail_lock(available_lock);
if (!available_contexts.empty())
{
task->assign_context(available_contexts.front());
available_contexts.pop_front();
return;
}
// If we failed to get a context, double the number of total
// contexts and then update the forest nodes to have the right
// number of contexts available
task->assign_context(RegionTreeContext(total_contexts));
for (unsigned idx = 1; idx < total_contexts; idx++)
{
available_contexts.push_back(RegionTreeContext(total_contexts+idx));
}
// Mark that we doubled the total number of contexts
// Very important that we do this before calling the
// RegionTreeForest's resize method!
unsigned current_contexts = total_contexts;
__sync_fetch_and_add(&total_contexts,current_contexts);
if (total_contexts > MAX_CONTEXTS)
{
log_run(LEVEL_ERROR,"ERROR: Maximum number of allowed contexts %d "
"exceeded when initializing task %s (UID %lld). "
"Please change 'MAX_CONTEXTS' at top "
"of legion_config.h and recompile. It is also "
"possible to reduce context usage by annotating "
"task variants as leaf tasks since leaf tasks do "
"not require context allocation.",
MAX_CONTEXTS, task->variants->name,
task->get_unique_task_id());
exit(ERROR_EXCEEDED_MAX_CONTEXTS);
}
// Tell the forest to resize the number of available contexts
// on all the nodes
forest->resize_node_contexts(current_contexts);
}
First take a lock and then iterate over all of the nodes.
void RegionTreeForest::resize_node_contexts(unsigned total_contexts)
{
// We're only reading the maps of nodes, so we only need read permissions
AutoLock l_lock(lookup_lock,1,false/*exclusive*/);
for (std::map<LogicalRegion,RegionNode*>::const_iterator it =
region_nodes.begin(); it != region_nodes.end(); it++)
{
it->second->reserve_contexts(total_contexts);
}
for (std::map<LogicalPartition,PartitionNode*>::const_iterator it =
part_nodes.begin(); it != part_nodes.end(); it++)
{
it->second->reserve_contexts(total_contexts);
}
}
For each node reserve_contexts does the resize operation.
void RegionTreeNode::reserve_contexts(unsigned num_contexts)
{
// Hold the lock to prevent races on multiple people
// trying to update the reserve size.
// Also since deques don't copy objects when
// appending new ones, we can add states without affecting the
// already existing ones.
AutoLock n_lock(node_lock);
logical_states.append(num_contexts);
physical_states.append(num_contexts);
for (unsigned idx = physical_states.size()-num_contexts;
idx < physical_states.size(); idx++)
physical_states[idx] = PhysicalState(idx);
}
The allocation of the context occurs in launch_task:
void SingleTask::launch_task(void)
{
...
// If we're a leaf task and we have virtual mappings
// then it's possible for the application to do inline
// mappings which require a physical context
if (!chosen_variant.leaf || (num_virtual_mappings > 0))
{
// Request a context from the runtime
runtime->allocate_context(this);
Now we'll look back at the logical and physical state objects that are tracked by the nodes in the region tree. These are LogicalState and PhysicalState. In the logical state structure the field state list tracks the children that are open below the corresponding node that the state is stored in. The field mask tracks where simultaneous coherence has been used in the tree. There are other things here we'll cover in later classes.
struct LogicalState {
public:
LogicalState(void);
LogicalState(const LogicalState &state);
~LogicalState(void);
public:
LogicalState& operator=(const LogicalState &rhs);
void* operator new(size_t count);
void* operator new[](size_t count);
void operator delete(void *ptr);
void operator delete[](void *ptr);
public:
void reset(void);
public:
LegionList<FieldState,
LOGICAL_FIELD_STATE_ALLOC>::track_aligned field_states;
LegionList<LogicalUser,CURR_LOGICAL_ALLOC>::track_aligned
curr_epoch_users;
LegionList<LogicalUser,PREV_LOGICAL_ALLOC>::track_aligned
prev_epoch_users;
// Fields on which the user has
// asked for explicit coherence
FieldMask user_level_coherence;
};
The physical state also tracks a lot of information like which data is dirty.
struct PhysicalState {
public:
PhysicalState(void);
PhysicalState(ContextID ctx);
#ifdef DEBUG_HIGH_LEVEL
PhysicalState(ContextID ctx, RegionTreeNode *node);
#endif
PhysicalState(const PhysicalState &rhs);
public:
PhysicalState& operator=(const PhysicalState &rhs);
void* operator new(size_t count);
void* operator new[](size_t count);
void operator delete(void *ptr);
void operator delete[](void *ptr);
public:
FieldMask dirty_mask;
FieldMask reduction_mask;
FieldMask remote_mask; // which fields are valid remote copies
ChildState children;
LegionMap<InstanceView*, FieldMask,
VALID_VIEW_ALLOC>::track_aligned valid_views;
LegionMap<ReductionView*, FieldMask,
VALID_REDUCTION_ALLOC>::track_aligned reduction_views;
LegionMap<MaterializedView*, LegionMap<Event,FieldMask>::aligned,
PENDING_UPDATES_ALLOC>::tracked pending_updates;
public:
// These are used for managing access to the physical state
unsigned acquired_count;
bool exclusive;
std::deque<std::pair<UserEvent,bool/*exclusive*/> > requests;
public:
ContextID ctx;
#ifdef DEBUG_HIGH_LEVEL
RegionTreeNode *node;
#endif
};
Notice the LegionStack and LegionList structures. These are wrappers around stl structures from the standard library that have special features like tracking memory usage. All of these are defined in legion_allocation.h.