_DistributedLazilyCompilingFunctionCaller._dag_to_compiled_func loses metadata.

[kaushikcfd]

We preprocess each partitioned DAG and doing so loses some information on the way.

Consider this DAG:

Notice how the knowledge of "x1" being a DOFArray is unavailable during the preprocessing of "Partition 2" and this leads to some inference failures.

/cc @matthiasdiener @inducer

[matthiasdiener]

This issue manifests itself like this:

(ie, partitioning placeholders with empty axes tags)

The runtime error is:

  File "/Users/mdiener/Work/efuse/mirgecom/mirgecom/integrators/explicit_rk.py", line 33, in rk4_step
    k1 = rhs(t, state)
  File "/Users/mdiener/Work/efuse/arraycontext/arraycontext/impl/pytato/compile.py", line 298, in __call__
    compiled_func = self._dag_to_compiled_func(
  File "/Users/mdiener/Work/efuse/grudge/grudge/array_context.py", line 173, in _dag_to_compiled_func
    part_id_to_prg[part.pid], _, _ = self._dag_to_transformed_loopy_prg(d)
  File "/Users/mdiener/Work/efuse/arraycontext/arraycontext/impl/pytato/compile.py", line 373, in _dag_to_transformed_loopy_prg
    return super()._dag_to_transformed_pytato_prg(dict_of_named_arrays)
  File "/Users/mdiener/Work/efuse/arraycontext/arraycontext/impl/pytato/compile.py", line 324, in _dag_to_transformed_pytato_prg
    pt_dict_of_named_arrays = self.actx.transform_dag(dict_of_named_arrays)
  File "/Users/mdiener/Work/efuse/meshmode/meshmode/array_context.py", line 1291, in transform_dag
    dag = unify_discretization_entity_tags(dag)
  File "/Users/mdiener/Work/efuse/meshmode/meshmode/pytato_utils.py", line 466, in unify_discretization_entity_tags
    raise RuntimeError("Unification failed.")
RuntimeError: Unification failed.