Arm backend: Fuse RIFE heads without PixelShuffle (#24790)

Accept 13-channel block0 heads with 2x, 4x or 8x bilinear upsampling in
the packed-warp export pass. Read the head directly in the shader and
preserve the integer resize rounding and quantization boundaries.

Reuse the packed flow for the final residual addition so the existing
final-blend fusion can match. This removes three state resize ops and
restores two neural and two compute dispatches for the scale-4 model.
Retain the existing PixelShuffle path and reject unsupported layouts.

Extend C16 concat padding to quarter-size RGB branches. Reuse existing
padding only when its value matches, preserving shared image resources
and the numerical behavior of the unpadded graph.


cc @digantdesai @freddan80 @per @zingo @mansnils @Sebastian-Larsson
@robell @rascani

Signed-off-by: Oscar Andersson <oscar.andersson@arm.com>

Author: oscarandersson8218

examples/arm/QAT_example/rife_vgf/passes/pack_rife_warps_and_concat.py

@@ -262,6 +262,29 @@ def _match_pixel_shuffle(slices, height, width):
     return phase, nodes
 
 
+def _match_previous_flow(concat, customs, nodes):
+    flow_transpose = customs[0].args[0][1]
+    residuals = [n for n in flow_transpose.users if n not in customs]
+    if len(residuals) != 1 or not _is_op(residuals[0], TOSA.RESCALE.default, 5):
+        raise ValueError("Expected one previous-flow rescale.")
+    residual = residuals[0]
+    if (
+        residual.args[1] != torch.int32
+        or len(residual.args[2]) != 1
+        or residual.args[4] != 0
+        or any(residual.kwargs.values())
+    ):
+        raise ValueError("Expected a signed per-tensor INT32 previous-flow rescale.")
+    nodes.add(flow_transpose)
+    allowed = nodes | set(customs) | {concat, residual}
+    if any(user not in allowed for node in nodes for user in node.users):
+        raise ValueError("The block0 tail has additional consumers.")
+    positions = {node: i for i, node in enumerate(concat.graph.nodes)}
+    if any(positions[user] <= positions[concat] for user in residual.users):
+        raise ValueError("Previous flow is needed before the packed shader.")
+    return residual
+
+
 def _match_block0_tail(concat, customs):
     _, height, width, _ = _tensor_shape(concat)
     if height % 4 or width % 4:
@@ -291,26 +314,62 @@ def _match_block0_tail(concat, customs):
         slices.append(source)
     phase, shuffle_nodes = _match_pixel_shuffle(slices, height, width)
     nodes.update(shuffle_nodes)
-    flow_transpose = customs[0].args[0][1]
-    residuals = [n for n in flow_transpose.users if n not in customs]
-    if len(residuals) != 1 or not _is_op(residuals[0], TOSA.RESCALE.default, 5):
-        raise ValueError("Expected one previous-flow rescale.")
-    residual = residuals[0]
-    if (
-        residual.args[1] != torch.int32
-        or len(residual.args[2]) != 1
-        or residual.args[4] != 0
-        or any(residual.kwargs.values())
-    ):
-        raise ValueError("Expected a signed per-tensor INT32 previous-flow rescale.")
-    nodes.add(flow_transpose)
-    allowed = nodes | set(customs) | {concat, residual}
-    if any(user not in allowed for node in nodes for user in node.users):
-        raise ValueError("The block0 tail has additional consumers.")
-    positions = {node: i for i, node in enumerate(concat.graph.nodes)}
-    if any(positions[user] <= positions[concat] for user in residual.users):
-        raise ValueError("Previous flow is needed before the packed shader.")
-    return phase, residual
+    return phase, _match_previous_flow(concat, customs, nodes)
+
+
+def _match_direct_head_tail(concat, customs):
+    _, height, width, _ = _tensor_shape(concat)
+    nodes = set()
+    head = None
+    factor = None
+    for rescale, offset, channels in zip(concat.args[0][5:], (4, 5, 0), (1, 8, 4)):
+        if not _is_op(rescale, TOSA.RESCALE.default, 5):
+            raise ValueError("Expected a resized head rescale.")
+        resize = rescale.args[0]
+        if not _is_op(resize, TOSA.RESIZE.default, 4):
+            raise ValueError("Expected a resized head slice.")
+        source = resize.args[0]
+        if not _is_op(source, TOSA.SLICE.default, 3):
+            raise ValueError("Expected an NHWC head slice.")
+        current_head = source.args[0]
+        batch, h, w, c = _tensor_shape(current_head)
+        if batch != 1 or c != 13 or h <= 0 or w <= 0 or height % h or width % w:
+            raise ValueError("Expected a 13-channel integer-scale head.")
+        current_factor = height // h
+        if current_factor not in (2, 4, 8) or width // w != current_factor:
+            raise ValueError("Expected a uniform 2x, 4x or 8x head resize.")
+        if head is not None and (head is not current_head or factor != current_factor):
+            raise ValueError("Expected one shared head for all state slices.")
+        head, factor = current_head, current_factor
+        denominator = 2 * factor
+        if (
+            current_head.meta["val"].dtype != torch.int8
+            or _shape_values(source.args[1]) != [0, 0, 0, offset]
+            or _shape_values(source.args[2]) != [1, h, w, channels]
+            or resize.meta["val"].dtype != torch.int32
+            or resize.kwargs != {"resize_mode": "bilinear"}
+            or [_shape_values(v) for v in resize.args[1:]]
+            != [
+                [denominator, 2, denominator, 2],
+                [1 - factor, 1 - factor],
+                [factor - 1, factor - 1],
+            ]
+            or rescale.args[1:] != (torch.int8, [1.0 / denominator**2], 0, 0)
+            or any(rescale.kwargs.values())
+        ):
+            raise ValueError("Expected the original signed integer head resize.")
+        nodes.update((source, resize, rescale))
+    return (head, _match_previous_flow(concat, customs, nodes)), factor
+
+
+def _match_head_tail(program, concat, customs):
+    try:
+        tail = _match_block0_tail(concat, customs)
+        parameters = _phase_major_parameters(program, tail[0])
+        return tail, parameters, True, 2
+    except ValueError:
+        tail, factor = _match_direct_head_tail(concat, customs)
+        return tail, None, False, factor
 
 
 def _phase_major_parameters(program, phase):
@@ -348,6 +407,8 @@ def _phase_major_parameters(program, phase):
 
 
 def _reorder_pixel_shuffle_output(program, phase, parameters):
+    if parameters is None:
+        return
     producer, weight, bias, scales = parameters
     graph = phase.graph
     constants = []
@@ -368,7 +429,7 @@ def _reorder_pixel_shuffle_output(program, phase, parameters):
 class PackRifeWarpsAndConcatPass(ArmPass):
     """Pack four private RIFE warps and their C28 concat into one C32 shader.
 
-    Fuse a compatible block0 PixelShuffle and state preparation into the shader.
+    Fuse compatible block0 head resizing and state preparation into the shader.
     Reorder block0 output channels during export so each PixelShuffle read
     accesses 13 contiguous channels. Keep sampling quantization boundaries.
     Four zero weight columns make the following convolution ignore padding.
@@ -489,18 +550,25 @@ def call(self, graph_module: GraphModule) -> PassResult:
             return PassResult(graph_module, False)
 
         phase_parameters = None
+        pixel_shuffle = True
+        resize_scale = 2
         try:
-            tail = _match_block0_tail(concat, customs)
-            phase_parameters = _phase_major_parameters(exported_program, tail[0])
+            tail, phase_parameters, pixel_shuffle, resize_scale = _match_head_tail(
+                exported_program, concat, customs
+            )
         except ValueError as error:
-            logger.debug("Keeping the original PixelShuffle layout: %s", error)
+            logger.debug("Keeping the original head layout: %s", error)
             tail = None
         if tail is None:
             payload = build_packed_warp_payload(payloads, timestep)
             shader_inputs = [*images, *concat.args[0][5:]]
         else:
             payload = build_packed_warp_payload(
-                payloads, timestep, fuse_block0_tail=True
+                payloads,
+                timestep,
+                fuse_block0_tail=True,
+                pixel_shuffle=pixel_shuffle,
+                resize_scale=resize_scale,
             )
             shader_inputs = [*images, tail[0]]
             _reorder_pixel_shuffle_output(exported_program, tail[0], phase_parameters)

examples/arm/QAT_example/rife_vgf/passes/pad_rife_block0_concat.py

@@ -4,89 +4,52 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
-import operator
 
 import torch
 from executorch.backends.arm._passes import ArmPass
 from executorch.backends.arm._passes.tosa_lowering_utils import create_const
 from executorch.backends.arm.tosa.mapping import TosaSpecialDtype
-from executorch.backends.arm.vgf.shaders.grid_sampler import (
-    CUSTOM_SHADER_DOMAIN_NAME,
-    flow_offset_grid_sampler_operator_name,
-)
 from executorch.backends.transforms.utils import create_constant_placeholder
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch.export import ExportedProgram
 from torch.export.graph_signature import InputKind
 from torch.fx import GraphModule
 from torch.fx.passes.shape_prop import _extract_tensor_metadata
 
-from ..shaders.final_warps import BLEND_OPERATOR_NAME
 from .pack_rife_warps_and_concat import (
     _constant_tensor,
     _is_op,
     _shape_values,
     _symmetric_weight,
     _tensor_shape,
-    OPERATOR_NAME,
     TOSA,
 )
 
 logger = logging.getLogger(__name__)
 
 
-def _ignores_alpha(custom, image):
-    if (
-        not _is_op(custom, TOSA.CUSTOM.default, 1)
-        or custom.kwargs.get("domain_name") != CUSTOM_SHADER_DOMAIN_NAME
-    ):
-        return False
-    name = custom.kwargs.get("operator_name")
-    if name in (OPERATOR_NAME, BLEND_OPERATOR_NAME):
-        return image in custom.args[0][:2] and image not in custom.args[0][2:]
-    if (
-        name != flow_offset_grid_sampler_operator_name()
-        or custom.args[0][0] is not image
-    ):
-        return False
-    for output in custom.users:
-        if not _is_op(output, operator.getitem, 2) or output.args[1] != 0:
-            return False
-        for user in output.users:
-            channel_axis = 3
-            crops = [user]
-            if _is_op(user, TOSA.TRANSPOSE.default, 2):
-                if _shape_values(user.args[1]) != [0, 3, 1, 2]:
-                    return False
-                channel_axis = 1
-                crops = list(user.users)
-            for crop in crops:
-                if (
-                    not _is_op(crop, TOSA.SLICE.default, 3)
-                    or _shape_values(crop.args[1])[channel_axis] != 0
-                    or _shape_values(crop.args[2])[channel_axis] != 3
-                ):
-                    return False
-    return True
-
-
-def _reusable_padding(image):
+def _reusable_padding(image, code):
     for candidate in image.users:
         if (
             _is_op(candidate, TOSA.PAD.default, 3)
             and _shape_values(candidate.args[1]) == [0, 0, 0, 0, 0, 0, 0, 1]
-            and all(_ignores_alpha(user, candidate) for user in candidate.users)
+            and _is_op(candidate.args[2], TOSA.CONST.default, 3)
+            and candidate.args[2].args == ([code], [1], torch.int8)
         ):
             return candidate
     return None
 
 
-def _move_padding_before_resize(image, resize, positions):
+def _move_padding_before_resize(image, resize):
+    positions = {node: i for i, node in enumerate(image.graph.nodes)}
     if positions[image] < positions[resize]:
         return
-    padding = image.args[1]
-    if isinstance(padding, torch.fx.Node) and positions[padding] > positions[resize]:
-        resize.prepend(padding)
+    for operand in image.args[1:]:
+        if (
+            isinstance(operand, torch.fx.Node)
+            and positions[operand] > positions[resize]
+        ):
+            resize.prepend(operand)
     resize.prepend(image)
 
 
@@ -142,18 +105,24 @@ def _match(self, concat):
             ):
                 raise ValueError("Expected a private signed 1/4 image rescale.")
             resize = rescale.args[0]
+            if not _is_op(resize, TOSA.RESIZE.default, 4):
+                raise ValueError("Expected a bilinear RGB resize.")
+            parameters = [_shape_values(arg) for arg in resize.args[1:]]
+            if parameters not in (
+                [[2, 4, 2, 4], [1, 1], [-1, -1]],
+                [[2, 8, 2, 8], [3, 3], [-3, -3]],
+            ):
+                raise ValueError("Expected a half-size or quarter-size RGB resize.")
+            factor = parameters[0][1] // 2
             if (
-                not _is_op(resize, TOSA.RESIZE.default, 4)
-                or [_shape_values(arg) for arg in resize.args[1:]]
-                != [[2, 4, 2, 4], [1, 1], [-1, -1]]
-                or resize.kwargs != {"resize_mode": "bilinear"}
+                resize.kwargs != {"resize_mode": "bilinear"}
                 or resize.meta["val"].dtype != torch.int32
                 or set(resize.users) != {rescale}
                 or _tensor_shape(resize.args[0])
-                != (shape[0], shape[1] * 2, shape[2] * 2, 3)
+                != (shape[0], shape[1] * factor, shape[2] * factor, 3)
                 or resize.args[0].meta["val"].dtype != torch.int8
             ):
-                raise ValueError("Expected the original half-size RGB resize.")
+                raise ValueError("Expected a private signed INT8 RGB resize.")
         return conv, weight, int(timestep.flatten()[0])
 
     def call(self, graph_module: GraphModule) -> PassResult:
@@ -185,17 +154,12 @@ def call(self, graph_module: GraphModule) -> PassResult:
                 )
             for rescale, code in zip(concat.args[0][:2], (timestep, 0)):
                 resize = rescale.args[0]
-                image = _reusable_padding(resize.args[0])
-                positions = {node: i for i, node in enumerate(graph.nodes)}
-                anchor = (
-                    image
-                    if image is not None and positions[image] < positions[resize]
-                    else resize
-                )
-                with graph.inserting_before(anchor):
-                    pad_value = create_const(graph, code, torch.int8, resize.args[0])
+                image = _reusable_padding(resize.args[0], code)
                 if image is None:
                     with graph.inserting_before(resize):
+                        pad_value = create_const(
+                            graph, code, torch.int8, resize.args[0]
+                        )
                         padding = graph.call_function(
                             TOSA.CONST_SHAPE.default, ([0, 0, 0, 0, 0, 0, 0, 1],)
                         )
@@ -208,9 +172,7 @@ def call(self, graph_module: GraphModule) -> PassResult:
                         )
                         image.meta = dict(resize.args[0].meta)
                 else:
-                    # The recognized shaders only use RGB; their alpha padding is unused.
-                    image.args = (*image.args[:2], pad_value)
-                    _move_padding_before_resize(image, resize, positions)
+                    _move_padding_before_resize(image, resize)
                 resize.args = (image, *resize.args[1:])
                 # Constant bilinear input becomes 4 * code, then RESCALE divides by 4.
                 for node in (image, resize, rescale):