Class MemoryFormatOpsPass

exir/passes/memory_format_ops_pass.py:24–91

This pass replaces ops which takes torch.memory_format as an argument with 'equivalent' op which takes dim_order. This is towards the larger ExecuTorch goal to move away from torch.memory_format. There is a 1:1 mapping between the aten op and the new edge dialect dim_order op.

class MemoryFormatOpsPass(ExportPass):
    """
    This pass replaces ops which takes torch.memory_format as an argument with
    'equivalent' op which takes dim_order. This is towards the larger ExecuTorch
    goal to move away from torch.memory_format. There is a 1:1 mapping between
    the aten op and the new edge dialect dim_order op.
    """

    def call_operator(self, op, args, kwargs, meta):
        if not (isinstance(op, EdgeOpOverload) and op in DimOrderOpsMap):
            return super().call_operator(
                op,
                args,
                kwargs,
                meta,
            )

        # new kwargs with dim_order, and no memory_format for the new op
        nkwargs = dict(copy.deepcopy(kwargs))  # orig kwargs are immutable

        # Get the target memory format for the EdgeOp, defaulting to
        # preserve_format (clone() with no memory_format kwarg preserves
        # the input's layout instead of forcing contiguous).
        mem_format = nkwargs.pop("memory_format", torch.preserve_format)

        # Get input tensor and ndim
        input_tensor: Optional[torch.Tensor] = None
        if isinstance(args[0], ProxyValue) and args[0].is_tensor():
            input_tensor = args[0].to_tensor()
            ndim = input_tensor.dim()
        elif isinstance(args[0], torch.Tensor):
            input_tensor = args[0]
            ndim = input_tensor.dim()
        elif isinstance(args[0], torch.fx.immutable_collections.immutable_list):
            ndim = len(args[0])
        else:
            assert (
                0
            ), f"Expecting a Tensor, a ProxyValue, or a Sequence, but got {type(args[0])}"

        # Derive dim_order based on memory format
        dim_order: List[int]
        if mem_format in (None, torch.preserve_format):
            # preserve_format: inherit dim_order from input tensor
            if input_tensor is not None:
                dim_order = [int(d) for d in input_tensor.dim_order()]
            else:
                # Fallback to contiguous if no single input tensor is available
                # (e.g. list inputs like torch.stack).
                dim_order = list(range(ndim))
        else:
            # Explicit memory format (contiguous_format, channels_last, etc.)
            dim_order = get_dim_order(mem_format, ndim)

        nkwargs["dim_order"] = dim_order
        logger.debug(
            f"{op.__name__} = rank: {ndim}, memory_format: {mem_format}."
            f" {DimOrderOpsMap[op].__name__} = dim_order: {nkwargs['dim_order']}"