How to validate fairly involved transformations in loopy?

[kaushikcfd]

I have the following script:

import loopy as lp
import numpy as np
import math


def transform(t_unit: lp.TranslationUnit, ndim: int, ndof: int,
              n_e_per_wg: int, nwork_items_per_e: int,
              i_tiles: int, j_tiles: int) -> lp.TranslationUnit:
    kernel_name = t_unit.default_entrypoint.name

    knl = t_unit[kernel_name]
    insn_id, = [insn.id for insn in knl.instructions]
    del knl

    # {{{ get corresponding variables in t_unit

    vng = t_unit[kernel_name].get_var_name_generator()
    ing = t_unit[kernel_name].get_var_name_generator()

    i = "i"
    j = "j"
    J = "J"
    e = "e"
    D = "D"
    x = "x"
    r = "r"

    j_inner, j_tile = vng(f"{j}_inner"), vng(f"{j}_tile")
    e_inner, e_outer = vng(f"{e}_inner"), vng(f"{e}_outer")
    i_inner, i_tile = vng(f"{i}_inner"), vng(f"{i}_tile")
    i_inner_inner, i_inner_outer = (vng(f"{i_inner}_inner"),
                                    vng(f"{i_inner}_outer"))
    rprftchD, iprftchD, jprftchD = (vng(f"{r}prftchD"),
                                    vng(f"{i}prftchD"),
                                    vng(f"{j}prftchD"))
    D_fetch = vng(f"{D}_fetch")
    prcmpt_x_redn = ing(f"prcmpt_{x}_redn")
    e_prcmpt_subst, r_prcmpt_subst, j_prcmpt_subst = (vng(f"{e}prcmpt_subst"),
                                                      vng(f"{r}prcmpt_subst"),
                                                      vng(f"{j}prcmpt_subst"))

    j_prcmpt_subst_inner, j_prcmpt_subst_outer = (vng(f"{j_prcmpt_subst}_inner"),
                                                  vng(f"{j_prcmpt_subst}_outer"))

    # }}}

    # {{{ term hoisting to match the flop count of opt_einsum

    from pymbolic import variables
    from loopy.symbolic import get_dependencies

    knl = t_unit[kernel_name]

    knl = lp.split_reduction_inward(knl, x)
    knl = lp.hoist_invariant_multiplicative_terms_in_sum_reduction(knl, x)
    knl = lp.extract_multiplicative_terms_in_sum_reduction_as_subst(
        knl,
        subst_name="subst",
        within=None,
        arguments=variables(f"{e} {j} {r}"),
        terms_filter=lambda x: ((get_dependencies(x) & knl.all_inames())
                                <= {r, e, j})
    )

    knl = lp.split_iname(knl, j, math.ceil(ndof/j_tiles),
                         outer_iname=j_tile, inner_iname=j_inner)

    t_unit = t_unit.with_kernel(knl)

    # }}}

    t_unit = lp.split_iname(t_unit, e, n_e_per_wg,
                            inner_iname=e_inner, outer_iname=e_outer,
                            inner_tag="l.1", outer_tag="g.0")

    t_unit = lp.precompute(t_unit, J,
                           sweep_inames=[x, r],
                           precompute_outer_inames=frozenset({e_outer, e_inner,
                                                              i_inner_inner}),
                           temporary_address_space=lp.AddressSpace.PRIVATE,
                           default_tag=None,
                           precompute_inames=[vng("J_prftch_x"), vng("J_prftch_r")])

    # {{{ tile and prefetch D

    t_unit = lp.split_iname(t_unit, i, math.ceil(ndof/i_tiles),
                            inner_iname=i_inner, outer_iname=i_tile,
                            )
    t_unit = lp.precompute(t_unit, D, [i_inner, r, j_inner],
                           precompute_outer_inames=frozenset([e_outer,
                                                              i_tile,
                                                              j_tile]),
                           precompute_inames=[rprftchD, iprftchD, jprftchD],
                           temporary_address_space=lp.AddressSpace.LOCAL,
                           temporary_name=D_fetch,
                           default_tag=None)
    t_unit = lp.split_iname(t_unit, iprftchD, n_e_per_wg, inner_tag="l.1")
    t_unit = lp.split_iname(t_unit, jprftchD, nwork_items_per_e, inner_tag="l.0")

    # }}}

    # {{{ precompute 'subst'

    t_unit = lp.precompute(t_unit, "subst",
                           sweep_inames=[r, j_inner, e_inner],
                           precompute_inames=[e_prcmpt_subst,
                                              j_prcmpt_subst,
                                              r_prcmpt_subst],
                           precompute_outer_inames=frozenset({e_outer, j_tile}),
                           default_tag=None,
                           compute_insn_id=prcmpt_x_redn,
                           temporary_address_space=lp.AddressSpace.LOCAL)
    t_unit = lp.tag_inames(t_unit, {e_prcmpt_subst: "l.1"})

    t_unit = lp.split_iname(t_unit, j_prcmpt_subst, nwork_items_per_e,
                            inner_iname=j_prcmpt_subst_inner,
                            outer_iname=j_prcmpt_subst_outer,
                            inner_tag="l.0",
                            )

    # }}}

    t_unit = lp.split_iname(t_unit, i_inner, nwork_items_per_e,
                            inner_iname=i_inner_inner,
                            outer_iname=i_inner_outer,
                            inner_tag="l.0",
                            )

    # {{{ TODO: remove once github.com/inducer/loopy/issues/666 is resolved.

    t_unit = lp.realize_reduction(t_unit, insn_id_filter=insn_id)
    inames_to_duplicate = (frozenset({i_tile, i_inner_outer})
                           & t_unit[kernel_name].all_inames())
    acc_name = f"acc_{r}_{j_tile}_{j_inner}"
    t_unit = lp.privatize_temporaries_with_inames(t_unit,
                                                  inames_to_duplicate,
                                                  only_var_names={acc_name})

    t_unit = lp.duplicate_inames(
        t_unit,
        inames_to_duplicate,
        within=f"writes:{acc_name} and not reads:{acc_name}")
    t_unit = lp.duplicate_inames(
        t_unit,
        inames_to_duplicate,
        within=f"reads:{acc_name} and not writes:{acc_name}")

    # }}}

    return t_unit


if __name__ == "__main__":
    import pyopencl as cl
    from functools import partial

    Ndim = 3
    Ndof = 35

    bound_transform = partial(transform,
                              ndim=3,
                              ndof=35,
                              n_e_per_wg=10,
                              nwork_items_per_e=12,
                              i_tiles=4, j_tiles=4)

    cl_ctx = cl.create_some_context()

    t_unit = lp.make_kernel(
        "{[x,r,e,i,j]: 0<=x,r<3 and 0<=i,j<35 and 0<=e<Nel}",
        """
        u(_0, _1, _2) := u_ary[_0, _1, _2]
        J(_0, _1, _2) := J_ary[_0, _1, _2]
        D(_0, _1, _2) := D_ary[_0, _1, _2]

        out[e,i] = sum([x,r,j], J(x,r,e) * D(r, i, j) * u(x,e,j))
        """,
        [lp.GlobalArg("u_ary,J_ary,D_ary", dtype=np.float64, shape=lp.auto),
         ...],
        lang_version=(2018, 2))

    lp.auto_test_vs_ref(t_unit, cl_ctx, bound_transform(t_unit),
                        parameters={"Nel": 10})

# vim: fdm=marker

which on running as python file.py raises:

loopy.diagnostic.AutomaticTestFailure: results do not match -- (rel) l_2 err: 0.0462598, l_inf err: 1

This error can be attributed to incorrect loop bounds generated by loopy on Line 76 of the generated code (cf https://gist.github.com/kaushikcfd/44bc70d3431fcdc7b391505ee4026cfd#file-complicated_transfom_result-c-L76). However, I am unsure if this is due to a dependency violation in the transform script or Loopy's bounds generation is wrong. Any suggestions on how to debug this?

[kaushikcfd]

I am unsure if this is due to a dependency violation in the transform script or Loopy's bounds generation is wrong. Any suggestions on how to debug this?

FWIW, the transformed kernel has the following Basic Set as its domain:

---------------------------------------------------------------------------                                                                                    
DOMAINS:                                                                                                                                                       
[Nel] -> { [x, r, j_tile, j_inner, e_outer, e_inner, J_prftch_x, J_prftch_r, i_tile, rprftchD, iprftchD_outer, iprftchD_inner, jprftchD_outer, jprftchD_inner, 
eprcmpt_subst, rprcmpt_subst, jprcmpt_subst_outer, jprcmpt_subst_inner, i_inner_outer, i_inner_inner, i_inner_outer_0, i_tile_0, i_inner_outer_1, i_tile_1] :
          0 <= x <= 2
          and 0 <= r <= 2
          and j_tile >= 0
          and 0 <= j_inner <= 34 - 9j_tile
          and j_inner <= 8
          and e_inner >= 0
          and -10e_outer <= e_inner <= 9
          and e_inner < Nel - 10e_outer
          and 0 <= J_prftch_x <= 2
          and 0 <= J_prftch_r <= 2
          and i_tile >= 0
          and 0 <= rprftchD <= 2
          and iprftchD_inner >= 0
          and -10iprftchD_outer <= iprftchD_inner <= 34 - 9i_tile - 10iprftchD_outer
          and iprftchD_inner <= 8 - 10iprftchD_outer
          and iprftchD_inner <= 9
          and jprftchD_inner >= 0
          and -12jprftchD_outer <= jprftchD_inner <= 34 - 9j_tile - 12jprftchD_outer
          and jprftchD_inner <= 8 - 12jprftchD_outer
          and jprftchD_inner <= 11
          and eprcmpt_subst >= 0
          and -10e_outer <= eprcmpt_subst <= 9
          and eprcmpt_subst < Nel - 10e_outer
          and 0 <= rprcmpt_subst <= 2
          and jprcmpt_subst_inner >= 0
          and -12jprcmpt_subst_outer <= jprcmpt_subst_inner <= 34 - 9j_tile - 12jprcmpt_subst_outer
          and jprcmpt_subst_inner <= 8 - 12jprcmpt_subst_outer
          and jprcmpt_subst_inner <= 11
          and i_inner_inner >= 0
          and -12i_inner_outer <= i_inner_inner <= 34 - 9i_tile - 12i_inner_outer
          and i_inner_inner <= 8 - 12i_inner_outer
          and i_inner_inner <= 11
          and -i_inner_inner <= 12i_inner_outer_0 <= 34 - 9i_tile - i_inner_inner
          and 12i_inner_outer_0 <= 8 - i_inner_inner
          and i_tile_0 >= 0
          and 9i_tile_0 <= 34 - 10iprftchD_outer - iprftchD_inner
          and 9i_tile_0 <= 34 - 12i_inner_outer - i_inner_inner
          and 9i_tile_0 <= 34 - i_inner_inner - 12i_inner_outer_0
          and -i_inner_inner <= 12i_inner_outer_1 <= 34 - 9i_tile - i_inner_inner
          and 12i_inner_outer_1 <= 34 - i_inner_inner - 9i_tile_0
          and 12i_inner_outer_1 <= 8 - i_inner_inner
          and i_tile_1 >= 0
          and 9i_tile_1 <= 34 - 10iprftchD_outer - iprftchD_inner
          and 9i_tile_1 <= 34 - 12i_inner_outer - i_inner_inner
          and 9i_tile_1 <= 34 - i_inner_inner - 12i_inner_outer_0
          and 9i_tile_1 <= 34 - i_inner_inner - 12i_inner_outer_1 }                                                                                                                      
--------------------------------------------------------------------------- 

which is making it difficult to answer this question.

[kaushikcfd]

After spending some time here, the incorrect loop bound is coming because of multiple inames tagged with "l.0" being part of the same basic set. However, I do not think anything in the above transform script is a non-pattern.

[kaushikcfd]

the incorrect loop bound is coming because of multiple inames tagged with "l.0" being part of the same basic set

In this particular example, this could decoupled by restructuring the transformation as:

97,98c97,98
<     t_unit = lp.split_iname(t_unit, iprftchD, n_e_per_wg, inner_tag="l.1")
<     t_unit = lp.split_iname(t_unit, jprftchD, nwork_items_per_e, inner_tag="l.0")
---
>     t_unit = lp.split_iname(t_unit, iprftchD, nwork_items_per_e, inner_tag="l.0")
>     t_unit = lp.split_iname(t_unit, jprftchD, n_e_per_wg, inner_tag="l.1")

This emits the desired loopy kernel and passes the test.

[inducer]

Can you make out which iname has its bound miscomputed?

[kaushikcfd]

Yep, i_tile_1. The dependence of i_tile_1 on lid(1) is spurious.

[inducer]

Is this connected to #724? Or is it a separate issue?

[kaushikcfd]

The reason for incorrect code-generation for the example here is the same as #724, but, this discussion tries to ask how can debuggability be improved in Loopy and #724 is tracking the incorrect code-gen for the specific example used here.