Merge pull request #124 from inducer/refactor-pytato-compile-exec

Refactor PytatoActx.compile for usability by distributed

Author: matthiasdiener

arraycontext/impl/pytato/compile.py

@@ -187,6 +187,49 @@ class LazilyCompilingFunctionCaller:
     program_cache: Dict["PMap[Tuple[Any, ...], AbstractInputDescriptor]",
                         "CompiledFunction"] = field(default_factory=lambda: {})
 
+    def _dag_to_transformed_loopy_prg(self, dict_of_named_arrays):
+        from pytato.target.loopy import BoundPyOpenCLProgram
+
+        import loopy as lp
+
+        with ProcessLogger(logger, "transform_dag"):
+            pt_dict_of_named_arrays = self.actx.transform_dag(
+                pt.make_dict_of_named_arrays(dict_of_named_arrays))
+
+        with ProcessLogger(logger, "generate_loopy"):
+            pytato_program = pt.generate_loopy(pt_dict_of_named_arrays,
+                                               options=lp.Options(
+                                                   return_dict=True,
+                                                   no_numpy=True),
+                                               cl_device=self.actx.queue.device)
+            assert isinstance(pytato_program, BoundPyOpenCLProgram)
+
+        with ProcessLogger(logger, "transform_loopy_program"):
+
+            pytato_program = (pytato_program
+                              .with_transformed_program(
+                                  lambda x: x.with_kernel(
+                                      x.default_entrypoint
+                                      .tagged(FromArrayContextCompile()))))
+
+            pytato_program = (pytato_program
+                              .with_transformed_program(self
+                                                        .actx
+                                                        .transform_loopy_program))
+
+        return pytato_program
+
+    def _dag_to_compiled_func(self, dict_of_named_arrays,
+            input_id_to_name_in_program, output_id_to_name_in_program,
+            output_template):
+        pytato_program = self._dag_to_transformed_loopy_prg(dict_of_named_arrays)
+
+        return CompiledFunction(
+                self.actx, pytato_program,
+                input_id_to_name_in_program=input_id_to_name_in_program,
+                output_id_to_name_in_program=output_id_to_name_in_program,
+                output_template=output_template)
+
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         """
         Returns the result of :attr:`~LazilyCompilingFunctionCaller.f`'s
@@ -197,8 +240,6 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
         :attr:`~LazilyCompilingFunctionCaller.f` with *args* in a lazy-sense.
         The intermediary pytato DAG for *args* is memoized in *self*.
         """
-        from pytato.target.loopy import BoundPyOpenCLProgram
-
         arg_id_to_arg, arg_id_to_descr = _get_arg_id_to_arg_and_arg_id_to_descr(
             args, kwargs)
 
@@ -210,74 +251,70 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
             return compiled_f(arg_id_to_arg)
 
         dict_of_named_arrays = {}
-        # output_naming_map: result id to name of the named array in the
-        # generated pytato DAG.
-        output_naming_map = {}
-        # input_naming_map: argument id to placeholder name in the generated
-        # pytato DAG.
-        input_naming_map = {
+        output_id_to_name_in_program = {}
+        input_id_to_name_in_program = {
             arg_id: f"_actx_in_{_ary_container_key_stringifier(arg_id)}"
             for arg_id in arg_id_to_arg}
 
-        outputs = self.f(*[_get_f_placeholder_args(arg, iarg, input_naming_map)
-                           for iarg, arg in enumerate(args)],
-                         **{kw: _get_f_placeholder_args(arg, kw, input_naming_map)
-                            for kw, arg in kwargs.items()})
+        output_template = self.f(
+                *[_get_f_placeholder_args(arg, iarg, input_id_to_name_in_program)
+                    for iarg, arg in enumerate(args)],
+                **{kw: _get_f_placeholder_args(arg, kw, input_id_to_name_in_program)
+                    for kw, arg in kwargs.items()})
 
-        if not is_array_container_type(outputs.__class__):
+        if not is_array_container_type(output_template.__class__):
             # TODO: We could possibly just short-circuit this interface if the
             # returned type is a scalar. Not sure if it's worth it though.
             raise NotImplementedError(
                 f"Function '{self.f.__name__}' to be compiled "
                 "did not return an array container, but an instance of "
-                f"'{outputs.__class__}' instead.")
+                f"'{output_template.__class__}' instead.")
 
         def _as_dict_of_named_arrays(keys, ary):
             name = "_pt_out_" + "_".join(str(key)
                                          for key in keys)
-            output_naming_map[keys] = name
+            output_id_to_name_in_program[keys] = name
             dict_of_named_arrays[name] = ary
             return ary
 
         rec_keyed_map_array_container(_as_dict_of_named_arrays,
-                                      outputs)
+                                      output_template)
 
-        import loopy as lp
+        from pytato import DictOfNamedArrays
+        compiled_func = self._dag_to_compiled_func(
+                DictOfNamedArrays(dict_of_named_arrays),
+                input_id_to_name_in_program=input_id_to_name_in_program,
+                output_id_to_name_in_program=output_id_to_name_in_program,
+                output_template=output_template)
 
-        with ProcessLogger(logger, "transform_dag"):
-            pt_dict_of_named_arrays = self.actx.transform_dag(
-                pt.make_dict_of_named_arrays(dict_of_named_arrays))
+        self.program_cache[arg_id_to_descr] = compiled_func
+        return compiled_func(arg_id_to_arg)
 
-        with ProcessLogger(logger, "generate_loopy"):
-            pytato_program = pt.generate_loopy(pt_dict_of_named_arrays,
-                                               options=lp.Options(
-                                                   return_dict=True,
-                                                   no_numpy=True),
-                                               cl_device=self.actx.queue.device)
-            assert isinstance(pytato_program, BoundPyOpenCLProgram)
 
-        with ProcessLogger(logger, "transform_loopy_program"):
+def _args_to_cl_buffers(actx, input_id_to_name_in_program, arg_id_to_arg):
+    input_kwargs_for_loopy = {}
 
-            pytato_program = (pytato_program
-                              .with_transformed_program(
-                                  lambda x: x.with_kernel(
-                                      x.default_entrypoint
-                                      .tagged(FromArrayContextCompile()))))
-
-            pytato_program = (pytato_program
-                              .with_transformed_program(self
-                                                        .actx
-                                                        .transform_loopy_program))
+    for arg_id, arg in arg_id_to_arg.items():
+        if np.isscalar(arg):
+            arg = cla.to_device(actx.queue, np.array(arg))
+        elif isinstance(arg, pt.array.DataWrapper):
+            # got a Datwwrapper => simply gets its data
+            arg = arg.data
+        elif isinstance(arg, cla.Array):
+            # got a frozen array  => do nothing
+            pass
+        elif isinstance(arg, pt.Array):
+            # got an array expression => evaluate it
+            arg = actx.freeze(arg).with_queue(actx.queue)
+        else:
+            raise NotImplementedError(type(arg))
 
-        self.program_cache[arg_id_to_descr] = CompiledFunction(
-                                                self.actx, pytato_program,
-                                                input_naming_map, output_naming_map,
-                                                output_template=outputs)
+        input_kwargs_for_loopy[input_id_to_name_in_program[arg_id]] = arg
 
-        return self.program_cache[arg_id_to_descr](arg_id_to_arg)
+    return input_kwargs_for_loopy
 
 
-@dataclass
+@dataclass(frozen=True)
 class CompiledFunction:
     """
     A callable which captures the :class:`pytato.target.BoundProgram`  resulting
@@ -319,40 +356,18 @@ def __call__(self, arg_id_to_arg) -> ArrayContainer:
             :attr:`CompiledFunction.input_id_to_name_in_program` for input id's
             representation.
         """
-        from arraycontext.container.traversal import rec_keyed_map_array_container
-
-        input_kwargs_to_loopy = {}
-
-        # {{{ preprocess args to get arguments (CL buffers) to be fed to the
-        # loopy program
-
-        for arg_id, arg in arg_id_to_arg.items():
-            if np.isscalar(arg):
-                arg = cla.to_device(self.actx.queue, np.array(arg))
-            elif isinstance(arg, pt.array.DataWrapper):
-                # got a Datwwrapper => simply gets its data
-                arg = arg.data
-            elif isinstance(arg, cla.Array):
-                # got a frozen array  => do nothing
-                pass
-            elif isinstance(arg, pt.Array):
-                # got an array expression => evaluate it
-                arg = self.actx.freeze(arg).with_queue(self.actx.queue)
-            else:
-                raise NotImplementedError(type(arg))
-
-            input_kwargs_to_loopy[self.input_id_to_name_in_program[arg_id]] = arg
+        input_kwargs_for_loopy = _args_to_cl_buffers(
+                self.actx, self.input_id_to_name_in_program, arg_id_to_arg)
 
         evt, out_dict = self.pytato_program(queue=self.actx.queue,
                                             allocator=self.actx.allocator,
-                                            **input_kwargs_to_loopy)
+                                            **input_kwargs_for_loopy)
+
         # FIXME Kernels (for now) allocate tons of memory in temporaries. If we
         # race too far ahead with enqueuing, there is a distinct risk of
         # running out of memory. This mitigates that risk a bit, for now.
         evt.wait()
 
-        # }}}
-
         def to_output_template(keys, _):
             return self.actx.thaw(out_dict[self.output_id_to_name_in_program[keys]])