rfc: propose split op for oneDNN Graph API

[xiang1guo]

trafficstars

This RFC is to add a new Split operation to oneDNN graph API opset.

The rendered version: link

[xiang1guo]

Can you help to review this RFC? Thanks! @igorsafo @mgouicem @dzarukin

[chunyuan-w]

1. PyTorch split OP has a split_size_or_sections input:

• When it is list, it is aligned with the size_splits attr here.
• When it is an int, it seems that there's no correspondent attr:

  • PyTorch int split_size_or_sections:

    • It is the size of the chunk and it denotes that the input will be split into equally sized chunks (if possible).
    • Last chunk will be smaller if the tensor size along the given dimension dim is not divisible by split_size.

  • num_splits in oneDNN graph:

    • It is the number of chunks
    • The dimension of input tensor data shape along axis must be evenly divisible by num_splits attribute.

In static shape cases, if the input shape is known, the integration could try to convert int split_size_or_sections to num_splits in case the dimension could be evenly divisible by num_splits. In dynamic shape cases, I feel like it will be hard for the integration to map the OP to oneDNN graph.

2. For PyTorch, each of the output chunk is a view of the original tensor. If it is not the case in oneDNN graph, may I know if this will be mentioned anywhere in the document?

[sanchitintel]

@xiang1guo, we plan to emulate dynamic shape support in IPEX with multiple compiled partitions for various static shapes. Since the num_splits attribute is to be provided at the time of graph creation, IPEX would provide its value for the first input but then it might change for a later input! :(

cc @chunyuan-w

[xiang1guo]

1. PyTorch split OP has a split_size_or_sections input:

• When it is list, it is aligned with the size_splits attr here.

• When it is an int, it seems that there's no correspondent attr:

  • PyTorch int split_size_or_sections:

    • It is the size of the chunk and it denotes that the input will be split into equally sized chunks (if possible).
    • Last chunk will be smaller if the tensor size along the given dimension dim is not divisible by split_size.

  • num_splits in oneDNN graph:

    • It is the number of chunks
    • The dimension of input tensor data shape along axis must be evenly divisible by num_splits attribute.

In static shape cases, if the input shape is known, the integration could try to convert int split_size_or_sections to num_splits in case the dimension could be evenly divisible by num_splits. In dynamic shape cases, I feel like it will be hard for the integration to map the OP to oneDNN graph.

2. For PyTorch, each of the output chunk is a view of the original tensor. If it is not the case in oneDNN graph, may I know if this will be mentioned anywhere in the document?

For 1, as you mentioned in the meeting, we can add another attribute to align the meaning of Pytorch. It helps for mapping, but it still just can't handle the situation when the axis dimension of the input tensor is dynamic. @ZhennanQin @yifeizh2 Do you have any comments?

For 2, yes, graph API doesn't guarantee the view or copy. I will add more descriptions for it. Thanks for your reminder!

[ZhennanQin]

For 1, as you mentioned in the meeting, we can add another attribute to align the meaning of Pytorch. It helps for mapping, but it still just can't handle the situation when the axis dimension of the input tensor is dynamic.

What do you mean by the axis dimension of the input tensor is dynamic? Can you provide an example for better understanding?

[xiang1guo]

For 1, as you mentioned in the meeting, we can add another attribute to align the meaning of Pytorch. It helps for mapping, but it still just can't handle the situation when the axis dimension of the input tensor is dynamic.

What do you mean by the axis dimension of the input tensor is dynamic? Can you provide an example for better understanding?

For dynamic input data [1, 28, ?](? means this dimension can be gotten in the execution stage) and the axis=2.

[ZhennanQin]

For 1, as you mentioned in the meeting, we can add another attribute to align the meaning of Pytorch. It helps for mapping, but it still just can't handle the situation when the axis dimension of the input tensor is dynamic.

What do you mean by the axis dimension of the input tensor is dynamic? Can you provide an example for better understanding?

For dynamic input data [1, 28, ?](? means this dimension can be gotten in the execution stage) and the axis=2.

I see. Thanks for the example. In this case, the number of outputs from split is dynamic, how's this represented in the framework graph?

[xiang1guo]

[xiang1guo]

For 1, as you mentioned in the meeting, we can add another attribute to align the meaning of Pytorch. It helps for mapping, but it still just can't handle the situation when the axis dimension of the input tensor is dynamic.

What do you mean by the axis dimension of the input tensor is dynamic? Can you provide an example for better understanding?

For dynamic input data [1, 28, ?](? means this dimension can be gotten in the execution stage) and the axis=2.

I see. Thanks for the example. In this case, the number of outputs from split is dynamic, how's this represented in the framework graph?

Seems for now we don't have any real case. @chunyuan-w @sanchitintel Do you have any real case about it?

[wuxun-zhang]

@xiang1guo, we plan to emulate dynamic shape support in IPEX with multiple compiled partitions for various static shapes. Since the num_splits attribute is to be provided at the time of graph creation, IPEX would provide its value for the first input but then it might change for a later input! :(

cc @chunyuan-w

Hi @sanchitintel, by saying then it might change for a later input, do you mean that num_splits will be no longer an attribute? How is it handled in bridge?

[sanchitintel]

by saying then it might change for a later input, do you mean that num_splits will be no longer an attribute? How is it handled in bridge?

Hi @wuxun-zhang, if the input-size would change, then num_splits may change as well, but since the framework would have passed a value of num_splits corresponding to the first input encountered at the time of graph creation, I think we would run into a shape-mismatch issue at runtime.

So, it'd be great if we could pass this value at compile time instead. IIRC, something similar is done for DynamicReshape as well.

[sanchitintel]

In this case, the number of outputs from split is dynamic, how's this represented in the framework graph?

@ZhennanQin, currently, the IPEX/PyTorch JIT framework graph is created by tracing with a sample input, so the framework graph's nodes are specific to that input, so, for the case you mentioned, another input with a different input-shape would have generated a different graph.

So, if we were to lower such an op (that would generate a different number of tensors) to oneDNN Graph, then we would encounter a runtime error with an input tensor of a different shape than profiled. (However, if we wouldn't lower such an op to oneDNN Graph, then those ops would be handled by Pytorch/IPEX eager mode.)

@frost-intel @chunyuan-w, do you know how Dynamo handles such cases? Thanks!

[chunyuan-w]

In this case, the number of outputs from split is dynamic, how's this represented in the framework graph?

@ZhennanQin, currently, the IPEX/PyTorch JIT framework graph is created by tracing with a sample input, so the framework graph's nodes are specific to that input, so, for the case you mentioned, another input with a different input-shape would have generated a different graph.

So, if we were to lower such an op (that would generate a different number of tensors) to oneDNN Graph, then we would encounter a runtime error with an input tensor of a different shape than profiled. (However, if we wouldn't lower such an op to oneDNN Graph, then those ops would be handled by Pytorch/IPEX eager mode.)

@frost-intel @chunyuan-w, do you know how Dynamo handles such cases? Thanks!

I would suggest adding an attribute split_size to match the int type split_size_or_sections in PyTorch.

• When the TorchScript IR indicates that this int split_size_or_sections is a Constant, the bridge will set the split_size attr of Split OP to be the Constant value that is known during compilation time.
• If it is not a Constant according to the TorchScript IR, we're not able to map it to the static attr. Another runtime representation of split_size will be needed. A bit similar to the discussion on the exponent of the Pow OP.

[frost-intel]

@frost-intel @chunyuan-w, do you know how Dynamo handles such cases? Thanks!

Dynamo can either use a constant input size similar to PT JIT, or it can have a dynamic size marked as a SymInt object with undetermined size. We haven't explored dynamic sizes in our Dynamo/oneDNN Graph bridge, but an attribute or flag marking a dimension as Dynamic would likely be required.

[sanchitintel]

• If it is not a Constant according to the TorchScript IR, we're not able to map it to the static attr. Another runtime representation of split_size will be needed. A bit similar to the discussion on the exponent of the Pow OP.

Thanks, @chunyuan-w! If split_size_or_sections varies for different input-sizes, then it'd not be constant in TorchScript IR, but would be generated with a combination of ops including aten::size. So, I agree with your proposal, except that in the multiple-compiled-partitions-for-different-static-input-shapes implementation, split_size would also be available at compile time (as opposed to run-time).

Thanks for your inputs, @frost-intel!

[xiang1guo]

@chunyuan-w @sanchitintel @retonym @ZhennanQin @yifeizh2 I just updated the RFC and added an option2 per pytorch team's request. Please help to review the option2 again, if confirmed, our final implementation will be based on this option.

[TaoLv]

The proposal was reviewed internally and will not be implemented in the library due to the complexity of defining a split op to support different usages in different frameworks. Thank you all for the discussion and participation.