[Lazy evaluation] Pytato Array Context with transformations

[kaushikcfd]

Draft because:

• Implement a global reduction parallelization strategy
• Tighter write race checking loopy#224
• Implements a utility to get the index of ReturnFromKernel at a point in the linearized kernel loopy#476
• Defines PytatoPyOpenCLArrayContext.transform_dag arraycontext#68
• Show a DAG manipulating example: materialization pytato#127

Related/linked PRs:

• [Lazy evaluation] Pytato array context transforms pytato#169
• [Lazy evaluation] Pytato array context transforms arraycontext#104

[kaushikcfd]

@inducer said in a personal meeting that this logic should also handle reductions kernels.

[inducer]

Thanks! A few quick drive-by thoughts below.

[inducer]

Doesn't this assume that each iname is only used by exactly one statement?

[inducer]

• Which type of kernels are those?
• I would like it if this were more guided by metadata, along the lines of what the pyopencl actx does.

[kaushikcfd]

We could borrow the unifier from FusionContractorArrayContext, but that would add a dependency on #284 and inducer/pytato#224. If those go in before this PR, I will put that here.

[inducer]

My recollection of our discussion was that we'd do this without aliasing... am I remembering wrong? If not, what made you change your mind?

[kaushikcfd]

I thought we wouldn't alias at the pytato's CodeGenMapper stage, but post-linearization grabbing hold of dead-temporaries would be trivial. I.e. I had thought we would alias the global temporaries not at the pytato-level but downstream as loopy transformation.

[kaushikcfd]

x-ref: inducer/pytato#122

[matthiasdiener]

When trying this on the mirgecom examples, I'm getting this type of error:

  File "/Users/mdiener/Work/emirge/grudge/grudge/discretization.py", line 623, in nodes
    return self.discr_from_dd(dd).nodes()
  File "/Users/mdiener/Work/emirge/meshmode/meshmode/discretization/__init__.py", line 610, in nodes
    result = make_obj_array([
  File "/Users/mdiener/Work/emirge/meshmode/meshmode/discretization/__init__.py", line 611, in <listcomp>
    _DOFArray(None, tuple([
  File "/Users/mdiener/Work/emirge/meshmode/meshmode/discretization/__init__.py", line 612, in <listcomp>
    actx.freeze(resample_mesh_nodes(grp, iaxis)) for grp in self.groups
  File "/Users/mdiener/Work/emirge/arraycontext/arraycontext/impl/pytato/__init__.py", line 133, in freeze
    pt_prg = pt_prg.with_transformed_program(self.transform_loopy_program)
  File "/Users/mdiener/Work/emirge/pytato/pytato/target/loopy/__init__.py", line 124, in with_transformed_program
    return self.copy(program=f(self.program))
  File "/Users/mdiener/Work/emirge/meshmode/meshmode/array_context.py", line 482, in transform_loopy_program
    t_unit = lp.linearize(lp.preprocess_kernel(t_unit))
  File "/Users/mdiener/Work/emirge/loopy/loopy/schedule/__init__.py", line 2201, in linearize
    knl = get_one_linearized_kernel(knl,
  File "/Users/mdiener/Work/emirge/loopy/loopy/schedule/__init__.py", line 2178, in get_one_linearized_kernel
    result = _get_one_scheduled_kernel_inner(kernel,
  File "/Users/mdiener/Work/emirge/loopy/loopy/schedule/__init__.py", line 2146, in _get_one_scheduled_kernel_inner
    return next(iter(generate_loop_schedules(kernel, callables_table)))
  File "/Users/mdiener/Work/emirge/loopy/loopy/schedule/__init__.py", line 1953, in generate_loop_schedules
    yield from generate_loop_schedules_inner(kernel,
  File "/Users/mdiener/Work/emirge/loopy/loopy/schedule/__init__.py", line 2084, in generate_loop_schedules_inner
    gen_sched = insert_barriers(kernel, gen_sched,
  File "/Users/mdiener/Work/emirge/loopy/loopy/schedule/__init__.py", line 1919, in insert_barriers
    result = insert_barriers_at_outer_level(result)
  File "/Users/mdiener/Work/emirge/loopy/loopy/schedule/__init__.py", line 1838, in insert_barriers_at_outer_level
    append_barrier_or_raise_error(
  File "/Users/mdiener/Work/emirge/loopy/loopy/schedule/__init__.py", line 1760, in append_barrier_or_raise_error
    raise MissingBarrierError(
loopy.diagnostic.MissingBarrierError: _pt_kernel: Dependency '_pt_out_store depends on call_lp_nodes' (for variable '_pt_temp') requires synchronization by a global barrier (add a 'no_sync_with' instruction option to state that no synchronization is needed)

[kaushikcfd]

@matthiasdiener: Sorry, the transformations in this branch depend on some PRs on the pytato/loopy/arraycontext end. I've edited the description to record that.

[kaushikcfd]

Closing this as inducer/arraycontext#216 provides a cleaner implementation.