Tiling with Reduction Domain

[ossanossan]

Hi,
I tried adding tiling to the argmax operation as follows.

		RDom r(0, input.dim(3).extent(), "r");

		Func argmaxResult("argmaxResult");

		argmaxResult(x, y, z) = Tuple(cast<int32_t>(0), input(x, y, z, 0));

		Expr f32CurVal = input(x, y, z, r);
		Expr s32CurMaxIdx = argmaxResult(x, y, z)[0];
		Expr f32CurMaxVal = argmaxResult(x, y, z)[1];

		argmaxResult(x, y, z) = tuple_select(f32CurVal > f32CurMaxVal,
								Tuple(r, f32CurVal),
								Tuple(s32CurMaxIdx, f32CurMaxVal));

		s32Index(x, y, z) = argmaxResult(x, y, z)[0];

		//// Schedule
		int32_t s32TileWidth = 4;
		int32_t s32TileHeight = 4;
		int32_t s32TileChannels = 4;

		s32Index
			.split(x, x_outer, x_inner, s32TileWidth)
			.split(y, y_outer, y_inner, s32TileHeight)
			.split(z, z_outer, z_inner, s32TileChannels)
			.reorder(x_inner, y_inner, z_inner, x_outer, y_outer, z_outer);

But, because the Reduction Domain is mapped to the innermost loop, the load operation inside the innermost loop accesses outside the tile, as shown below.

                   produce s32Index:
                     for z.z_outer:
                       for y.y_outer:
                         for x.x_outer:
                           for z.z_inner in [0, 3]:
                             for y.y_inner in [0, 3]:
                               for x.x_inner in [0, 3]:
                                 produce argmaxResult:
                                   argmaxResult(...) = ...
                                   for r:
                                      argmaxResult(...) = ...
                                 consume argmaxResult:
                                   s32Index(...) = ...

To prevent this, I tried to move the "r" loop outside of "z_inner" using .reorder() on argmaxResult.update(), but it did not work(The compile was OK but the order was not changed). I also thought about including "r" in the tiles with splitting, but ”r” is not visible from "s32Index", so I don't know how to write it.

Does anyone have a solution to this issue?

Thanks.