PyTorch FX exporter dispatching and type promotion scenarios

[justinchuby]

1. An ATen op can be dispatched to multiple onnx-script implementations:

aten::aten_scaled_dot_product_attention ->

def aten_scaled_dot_product_attention(
    query: TFloat,
    key: TFloat,
    value: TFloat,
    attn_mask: Optional[TFloat] = None,
    dropout_p: float = 0.0,
    is_causal: bool = False,
    scale: Optional[float] = None,
): ...

def aten_scaled_dot_product_attention_bool_mask(
    query: TFloat,
    key: TFloat,
    value: TFloat,
    attn_mask: Optional[BOOL] = None,
    dropout_p: float = 0.0,
    is_causal: bool = False,
    scale: Optional[float] = None,
): ...

The exporter will need to dispatch onto one of the implementation based on the dtype of attn_mask.

2. (Easy) Multiple ATen op overloads may map to a single onnx-script implementation:

Both add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor and add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor map to a single

def aten_add(self: TReal, other: TReal, alpha: float = 1.0) -> TReal: ...

Type promotion

Pytorch

• torch_implicit_integral_promotion: Sometimes we don't need to Cast the output back: Cast(int to float) -> Cos -> output [ONNX] Add supported ops into test_fx_op_consistency - 1st batch pytorch/pytorch#100265 (comment)
• torch_implicit_type_promotion: ONNX does not have automatic type promotion.
  • We can either run an fx pass to create explicit type promotion nodes before converting to ONNX, or
  • Create rules to promote types in the exporter during conversion.
  • The exporter can consult the type_constraints in a function's OpSchema to determine whether it needs to cast inputs to the same type.

ONNX

• ONNX tends to support less types than PyTorch. When an ONNX function is selected as the dispatch target but it does not support the input type, the exporter needs to insert cast nodes before and after the function. E.g. Cast -> <the torchlib function> -> Cast to mirror the aten op's behavior.
• The exporter may not be able to dispatch to a function with exact type match but can still make a successful dispatch under case (1). For example, we may have two functions supporting (a) INT64, (b) FLOAT32, and the input is BFLOAT16. The dispatcher should be able to choose (b) as the target and use the process in (1) to insert proper cast nodes.

cc @titaiwangms @thiagocrepaldi @BowenBao

[BowenBao]

From the perspective of each aten op defines its own type promotion logic, as well as the perspective of eager execution. Could you also discuss the pros/cons/viability of dispatcher/evaluator/OnnxFunction handling automatic type promotion?

[BowenBao]

Discussed offline with @titaiwangms.

Scratch notes for case of torchlib owning type promotion

• Identify and maintain a set of type_promotion_kinds. Each type_promotion_kind represents a set of logic of how casting should be done based on arg/output types.
• Decorate torchlib ops with their respective set of type_promotion_kinds, which gets populated as attribute of OnnxFunction and OnnxTraceFunction instance.
• Extend Evaluator (or TypePromotionEvaluator) to insert proper cast nodes based on opschema, actual args/output types, and type_promotion_kinds.

[justinchuby]

Type comparison: https://github.com/pytorch/pytorch/blob/23b7035b3c61b0900e900e39ba7e69997a417ea8/torch/_prims_common/__init__.py#L1073

class ELEMENTWISE_TYPE_PROMOTION_KIND(Enum):
    DEFAULT = (0,)
    NO_OPMATH = (1,)
    INT_TO_FLOAT = (2,)
    ALWAYS_BOOL = (3,)
    COMPLEX_TO_FLOAT = (4,)
    BOOL_TO_LONG = (5,)


class REDUCTION_OUTPUT_TYPE_KIND(Enum):
    SAME = (0,)
    COMPLEX_TO_FLOAT = (1,)  # for complex types outputs corresponding real type
    KEEP_PROMOTED_TYPE = (2,)  # keep output in opmath type, needed for mean
    ALWAYS_BOOL = (3,)

Usage: https://github.com/pytorch/pytorch/blob/23b7035b3c61b0900e900e39ba7e69997a417ea8/torch/_prims_common/__init__.py#L1291-L1333

[thiagocrepaldi]

@justinchuby should we close this?

[justinchuby]

It's half done, with the ONNX part not being urgent right now. I would leave it open.

[thiagocrepaldi]

np.
is it possible to add the missing/completed tasks to help us know what is and is not done?

[justinchuby]

Sure! Done.

[thiagocrepaldi]

@justinchuby should this be closed?

[justinchuby]

Yes. Thanks!

For reference, The two points under ONNX (1) is not needed in practice; (2) is handled by overloads