How to merge two computations that use reduction domain

[apartridge]

Hi.

I have two functions that both depend on a rdom. I have made a minimal example below, where I calculate both the sum_ and sum_squared of some region of the input data. However, by default sum_ and sum_squared will be calculated independently, in two loops over the rdom.

Here is the loop nest that is printed. How would I be able to schedule this so that sum_ and sum_squared are calculated together in the same loop? Thanks for any tips!

produce out:
  for img:
    for row:
      for col:
        produce sum_squared:
          sum_squared(...) = ...
          for rdom in [-2, 2]:
            for rdom in [-2, 2]:
              sum_squared(...) = ...
              # I also want to compute sum_ here, in this loop
        consume sum_squared:
          produce sum_:
            sum_(...) = ...
            for rdom in [-2, 2]:
              for rdom in [-2, 2]:
                sum_(...) = ...
          consume sum_:
            out(...) = ...

Code:

import halide as hl
from halide import BoundaryConditions, Buffer, Func, Range, RDom, Var, f32


def find_target() -> hl.Target:
    host_target = hl.get_host_target()
    operating_system, architecture, bitness = host_target.os, host_target.arch, host_target.bits
    target_extensions = [
        hl.TargetFeature.OpenCL
    ]
    return hl.Target(operating_system, architecture, bitness, target_extensions)


def run(xyz):
    xyz_bound = BoundaryConditions.mirror_interior(xyz, [Range(0, 10), Range(0, 10)])

    rdom = RDom([(-2, 5), (-2, 5)], name="rdom")

    col = Var("col")
    row = Var("row")
    img = Var("img")

    sum_ = Func("sum_")
    sum_[col, row, img] = f32(0.0)
    sum_[col, row, img] += xyz_bound[col+rdom.x, row+rdom.y, img]

    sum_squared = Func("sum_squared")
    sum_squared[col, row, img] = f32(0.0)
    sum_squared[col, row, img] += xyz_bound[col+rdom.x, row+rdom.y, img]*xyz_bound[col+rdom.x, row+rdom.y, img]

    out = Func("out")
    out[col, row, img] = sum_[col, row, img] / sum_squared[col, row, img]

    target = find_target()
    out.print_loop_nest()
    return out.realize([10, 10, 3], target)


if __name__ == "__main__":
    xyz_buffer = Buffer(type=hl.Float(32), sizes=[10, 10, 3], name="xyz_buffer")
    run(xyz=xyz_buffer)

[stewpend0us]

Were you able to figure this out?