Pygion projection functor

.gitlab-ci.yml

@@ -666,14 +666,14 @@ gcc49_cxx98_debug_hdf5_mpi_regent:
   variables:
     <<: [*gcc49, *terra38, *debug, *cxx98_normal, *hdf5, *mpi, *regent]
 # Multi-node Legion Spy
-gcc49_cxx98_debug_spy_gasnet_regent:
-  <<: [*linux_compute, *image, *tests]
-  variables:
-    <<: [*gcc49, *terra38, *cxx98_normal, *spy, *gasnet, *regent]
-gcc49_cxx98_debug_spy_mpi_regent:
-  <<: [*linux_compute, *image, *tests]
-  variables:
-    <<: [*gcc49, *terra38, *cxx98_normal, *spy, *mpi, *regent]
+#gcc49_cxx98_debug_spy_gasnet_regent:
+#  <<: [*linux_compute, *image, *tests]
+#  variables:
+#    <<: [*gcc49, *terra38, *cxx98_normal, *spy, *gasnet, *regent]
+#gcc49_cxx98_debug_spy_mpi_regent:
+#  <<: [*linux_compute, *image, *tests]
+#  variables:
+#    <<: [*gcc49, *terra38, *cxx98_normal, *spy, *mpi, *regent]
 # * Different architectures
 gcc49_cxx98_32bit_debug_legion:
   <<: [*linux, *image, *tests]

bindings/python/legion_top.py

@@ -263,7 +263,39 @@ def import_global(module, check_depth=True, block=True):
 # not safe to use with control replication so this will give them a way
 # to check whether they are running in a safe context or not
 def is_control_replicated():
-    return False
+    try:
+        # We should only be doing something for this if we're the top-level task
+        return c.legion_context_get_num_shards(top_level.runtime[0],
+                top_level.context[0], True) > 1
+    except AttributeError:
+        raise RuntimeError('"is_control_replicated" must be called in a legion_python task')
+
+
+# Helper class for deduplicating output streams with control replication
+class LegionOutputStream(object):
+    def __init__(self, shard_id, stream):
+        self.shard_id = shard_id
+        # This is the original stream
+        self.stream = stream
+
+    def close(self):
+        self.stream.close()
+
+    def flush(self):
+        self.stream.flush()
+
+    def write(self, string):
+        # Only do the write if we are shard 0
+        if self.shard_id == 0:
+            self.stream.write(string)
+
+    def writelines(self, sequence):
+        # Only do the write if we are shard 0
+        if self.shard_id == 0:
+            self.stream.writelines(sequence)
+
+    def isatty(self):
+        return self.stream.isatty()
 
 
 def legion_python_main(raw_args, user_data, proc):
@@ -282,6 +314,11 @@ def legion_python_main(raw_args, user_data, proc):
 
     top_level.runtime, top_level.context, top_level.task = runtime, context, task
 
+    # If we're control replicated, deduplicate stdout
+    if is_control_replicated():
+        shard_id = c.legion_context_get_shard_id(runtime[0], context[0], True)
+        sys.stdout = LegionOutputStream(shard_id, sys.stdout)
+
     # Run user's script.
     args = input_args(True)
     start = 1

bindings/python/main.cc

@@ -27,10 +27,21 @@ static bool control_replicate = true;
 static const char * const unique_name = "legion_python";
 static const VariantID vid = 1;
 
+class LegionPyShardingFunctor : public ShardingFunctor {
+public:
+  LegionPyShardingFunctor(void) { }
+  virtual ~LegionPyShardingFunctor(void) { }
+public:
+  virtual ShardID shard(const DomainPoint &point,
+                        const Domain &full_space,
+                        const size_t total_shards);
+};
+
 // Special mapper just for mapping the top-level Python tasks
 class LegionPyMapper : public Legion::Mapping::NullMapper {
 public:
-  LegionPyMapper(MapperRuntime *runtime, Machine machine, TaskID top_task_id);
+  LegionPyMapper(MapperRuntime *runtime, Machine machine, 
+                 TaskID top_task_id, ShardingID sharding_id);
   virtual ~LegionPyMapper(void);
 public:
   static AddressSpaceID get_local_node(void);
@@ -57,6 +68,11 @@ class LegionPyMapper : public Legion::Mapping::NullMapper {
                                   const MapTaskInput&      input,
                                   const MapTaskOutput&     default_output,
                                   MapReplicateTaskOutput&  output);
+  virtual void select_sharding_functor(
+                             const MapperContext                ctx,
+                             const Task&                        task,
+                             const SelectShardingFunctorInput&  input,
+                                   SelectShardingFunctorOutput& output);
   virtual void select_tunable_value(const MapperContext         ctx,
                                     const Task&                 task,
                                     const SelectTunableInput&   input,
@@ -81,6 +97,7 @@ class LegionPyMapper : public Legion::Mapping::NullMapper {
   const size_t total_nodes;
   const char *const mapper_name;
   const TaskID top_task_id;
+  const ShardingID sharding_id;
 protected:
   std::vector<Processor> local_pys; // Python processors
 };
@@ -125,11 +142,14 @@ static void python_main_callback(Machine machine, Runtime *runtime,
     runtime->register_task_variant(registrar, code_desc, NULL, 0, task_name, vid);
     runtime->attach_name(top_task_id+2, task_name, false/*mutable*/, true/*local only*/);
   }
+  // Register our sharding function for any global import tasks
+  const ShardingID sharding_id = runtime->generate_library_sharding_ids(unique_name, 1);
+  runtime->register_sharding_functor(sharding_id, new LegionPyShardingFunctor());
   // Register our mapper for the top-level task
   const MapperID top_mapper_id = runtime->generate_library_mapper_ids(unique_name, 1);
   runtime->set_top_level_task_mapper_id(top_mapper_id);
   runtime->add_mapper(top_mapper_id,
-      new LegionPyMapper(runtime->get_mapper_runtime(), machine, top_task_id));
+      new LegionPyMapper(runtime->get_mapper_runtime(), machine, top_task_id, sharding_id));
 }
 
 int main(int argc, char **argv)
@@ -197,10 +217,10 @@ int main(int argc, char **argv)
   return Runtime::start(argc, argv);
 }
 
-LegionPyMapper::LegionPyMapper(MapperRuntime *rt, Machine m, TaskID top_id) 
+LegionPyMapper::LegionPyMapper(MapperRuntime *rt, Machine m, TaskID top_id, ShardingID sid) 
   : NullMapper(rt, m), local_node(get_local_node()), 
     total_nodes(get_total_nodes(m)), mapper_name(create_name(local_node)),
-    top_task_id(top_id)
+    top_task_id(top_id), sharding_id(sid)
 {
   Machine::ProcessorQuery py_procs(machine);
   py_procs.local_address_space();
@@ -378,6 +398,28 @@ void LegionPyMapper::map_top_level_task(const MapperContext ctx,
   output.chosen_variant = vid;
 }
 
+ShardID LegionPyShardingFunctor::shard(const DomainPoint &point,
+                                       const Domain &full_domain,
+                                       size_t total_shards)
+{
+  Point<1> p = point;
+  Rect<1> bounds = full_domain;
+  const size_t volume = bounds.volume();
+  assert((volume % total_shards) == 0);
+  const size_t pernode = volume / total_shards;
+  return (p[0] / pernode);
+}
+
+void LegionPyMapper::select_sharding_functor(
+                             const MapperContext                ctx,
+                             const Task&                        task,
+                             const SelectShardingFunctorInput&  input,
+                                   SelectShardingFunctorOutput& output)
+{
+  assert(task.task_id == (top_task_id+1));
+  output.chosen_functor = sharding_id;
+}
+
 void LegionPyMapper::select_tunable_value(const MapperContext         ctx,
                                           const Task&                 task,
                                           const SelectTunableInput&   input,

bindings/python/projection_functor_example.c

@@ -0,0 +1,9 @@
+#include "legion.h"
+
+void proj_functor(legion_runtime_t runtime,
+                  legion_logical_partition_t parent,
+                  legion_domain_point_t point,
+                  legion_domain_t launch_domain)
+{
+  legion_projection_functor_logical_partition_print_arguments_2(parent);
+}

bindings/python/projection_functor_example.py

@@ -0,0 +1,67 @@
+from __future__ import print_function
+
+import pygion
+from pygion import (
+    index_launch,
+    task,
+    Domain,
+    ID,
+    IndexLaunch,
+    R,
+    Region,
+    Partition,
+    ProjectionFunctor,
+)
+
+from typing import cast, Callable
+
+import subprocess
+import petra as pt
+
+f = ProjectionFunctor.create(1 + ID)
+
+
+@task(privileges=[R])
+def hello(R, i, num):
+    print("hello from point %s (region %s)" % (i, R.ispace.bounds))
+    assert int(R.ispace.bounds[0, 0]) == int(i + num)
+
+
+@task
+def main():
+    R = Region([4], {"x": pygion.float64})
+    P = Partition.equal(R, [4])
+    for i in range(4):
+        print(
+            "python region %s is %s %s %s"
+            % (
+                i,
+                P[i].handle[0].tree_id,
+                P[i].handle[0].index_space.tid,
+                P[i].handle[0].index_space.id,
+            )
+        )
+    pygion.fill(R, "x", 0)
+
+    for i in IndexLaunch([3]):
+        hello(P[f(i)], i, 1)
+
+    for i in IndexLaunch([3]):
+        hello(P[i], i, 0)
+
+    for i in IndexLaunch([2]):
+        hello(P[i + 2], i, 2)
+
+    for i in IndexLaunch([2]):
+        hello(P[i + 2], i, 2)
+
+    index_launch([3], hello, P[ID], ID, 0)
+
+    # This Seg Fault when running all tests but not when it's the only test:
+    index_launch([3], hello, P[f(ID)], ID, 1)
+
+    index_launch([2], hello, P[ID + 2], ID, 2)
+
+
+if __name__ == "__main__":
+    main()