Automatic Type Promotion in OneFlow

6 min readNov 5, 2021

Written by Zheng Zekang; Translated by Wang Kaiyan, Dong Wenwen

This article illustrates the implementation of automatic type promotion in the OneFlow. Although type promotion is an operation that we occasionally use, if the output type is not handled correctly, it can lead to an overflow of results and incorrect results.

Introduction

Let’s take a brief look at the code written in PyTorch, and you can guess what the final type of output is.

The answer is:

We can find that for the same multiply operation, some of the result data types are promoted to a higher level, some are not and remain as int8. This is a type promotion system, where some type promotion rules are customized within the system to derive the data type of the final result based on the input data type.

Python Array API Standard

Here we can learn more about Python Array’s type promotion rules: https://data-apis.org/array-api/latest/API_specification/type_promotion.html

As you can see from the figure above:

Promotion of different types is given by their join.
Dashed lines indicate that behavior for Python scalars is undefined on overflow.
Boolean, integer and float types are not connected, indicating mixed-kind promotion is undefined.

Regarding the first rule, we can look at int8 and uint8, both of which eventually point to int16, indicating that the final type is raised to int16 after both operations.

And according to this rule, we can list a type promotion table (this table is significant and will be used subsequently in the PyTorch source code)

Taking the unsigned int series and the signed int series as examples, the tables listed are:

More type promotion rules can be found in the previously mentioned link

The horizontal and vertical coordinates represent the input data type respectively, and the values of the table represent the data type after type promotion:

i1 : 8-bit signed integer (i.e., int8 )
i2 : 16-bit signed integer (i.e., int16 )
i4 : 32-bit signed integer (i.e., int32 )
i8 : 64-bit signed integer (i.e., int64 )

Same for unsigned int

Type Promotion for Python Array and Scalar

The above are the type promotion rules for operations between arrays and arrays, but not for arrays and scalars (that is, a separate int, float value).

If they both belong to the same data type family (e.g. both are int family, containing int8, int32, int64), then the final data type follows the data type of the array
If they do not belong to the same data type family (e.g. one is int32 and one is float), then type promotion is performed

Let’s look at these two examples:

It should be noted that the behavior of Array and Scalar will be consistent with the behavior of Array and 0d Array.

We can test the previous two examples again, with the difference that we have changed the scalar to a 0d Array

Regarding the behavior of operations with Scalar, the standard of PyTorch is in line with that of Python Array API, but Numpy is different in that it will take a reasonable type promotion based on the range of data of scalar:

Personally, I prefer Scalar to be a separate behavior in type promotion, while Scalar Tensor and Tensor are the same behavior

Other Cases

In addition to the rules mentioned above, PyTorch has the following two cases:

PyTorch requires two inputs of exactly the same data type, e.g. torch.dot

2. There exists a minimum data type for the input, such as torch.sum, and passing any int series data type will result in a final output of torch.int64.

The above briefly introduces PyTorch’s type promotion rules, for more examples, you can refer to the official documentation: https://pytorch.org/docs/master/tensor_attributes.html#torch.torch.dtype

How Does PyTorch Realize Type Promotion?

As for the Kernel in the actual operation, the input and output data types have the same template parameters, so there is no specialization of a function where the input is int32 and the output is float32 or some other type.

PyTorch internally will first infer a reasonable dtype, then insert a to op, type promote the input tensor, and then go to the Kernel for the actual operation. Next we'll look at the source code of PyTorch:

Code involved:

https://github.com/pytorch/pytorch/blob/master/c10/core/ScalarType.h

https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Activation.cpp

https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/TensorIterator.cpp

https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/TypeProperties.cpp

ScalarType.h

In this header file the relevant data types and a type-promoted two-dimensional matrix are defined. This allows us to enter two data types and get the promoted data type according to the index.

Activation.cpp

https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Activation.cpp#L24 Let’s take one of the activation functions threshold as an example.

Here a build function is called which accepts a TensorIteratorConfig. This Config class is used to formulate various properties, as can be seen here promote_inputs_to_common_dtype is called and set to true.

TensorIterator.cpp

We can see the build function here: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/TensorIterator.cpp#L1321

At line 1340, the build function internally calls the compute_type function.

This function performs a series of type derivations starting at line 260.

TensorIterator is a container class (there is a similar container NpyIter in Numpy) for storing the output, input tensor, which uses multiple for loops to derive a common_dtype.

Make a conditional judgement at the end: promote_inputs_to_common_dtype_ is true, the current Tensor is not an output Tensor and the input dtype is not equal to the derived common_dtype, then do a type promotion.

The Implementation in OneFlow

OneFlow puts the type promotion logic in the functional front-end part of C++, similarly designing a TensorProcessor class. The interface is designed as follows:

Taking the binary operation Functor base class as an example, when it is actually called, we can do this as follows.

PromoteInputsToCommonDtype is used to set the relevant properties
The AddInputs function adds the Tensor that needs to be involved in the type promotion to the container
The Apply function performs the actual type promotion logic, etc

There are several other functions in tensor_processor.cpp, the functions of which are briefly described here:

CheckHasDifferentInputDType iterate over the input Tensor to check if the input Tensor has a different dtype
ComputeCommonDType derive a reasonable promoted dtype based on the input dtype
CastToSameType insert a cast operation to the input tensor

The logic of the Apply function is as follows:

Type promotion is implemented inside the if logic, while the else logic corresponds to the second of the other cases mentioned earlier, promoting the Tensor type to a set minimum data type. For the sum operator, we set the minimum data type to int64 by implementing this:

Conclusion

Type promotion is an operation that we occasionally use, and if the output data type is not handled correctly, it can lead to an overflow of results and incorrect results. This seems like a simple operation, but it took two or three weeks of actual research and puzzling over the details.

I hope this article will help you in your deep learning projects😊. If you want to experience the functions of OneFlow, you can follow the method described in this article. If you have any questions or comments💡 about use, please feel free to leave a comment in the comments section below. Please do the same if you have any comments, remarks or suggestions for improvement. In future articles, we’ll introduce more functions of OneFlow.

Related articles：

Welcome to visit OneFlow on GitHub and follow us on Twitter and LinkedIn.

Also, welcome to join our Discord group to discuss and ask OneFlow related questions, and connect with OneFlow contributors and users all around the world.