Support custom scalar types in PyTorch

soumyarooproy · December 5, 2018, 10:29pm

This example:

pytorch/pytorch/blob/master/test/cpp_extensions/complex_registration_extension.cpp

#include <torch/extension.h>

#include <ATen/CPUFloatType.h>
#include <ATen/Type.h>
#include <ATen/core/VariableHooksInterface.h>
#include <ATen/detail/ComplexHooksInterface.h>

#include "ATen/Allocator.h"
#include "ATen/CPUGenerator.h"
#include "ATen/DeviceGuard.h"
#include "ATen/NativeFunctions.h"
#include "ATen/Utils.h"
#include "ATen/WrapDimUtils.h"
#include "ATen/core/Half.h"
#include "ATen/core/TensorImpl.h"
#include "ATen/core/UndefinedTensorImpl.h"
#include "c10/util/Optional.h"

#include <cstddef>
#include <functional>

This file has been truncated. show original

illustrates creating a new tensor type using an existing supported scalar type, which in this case is ScalarType::ComplexFloat.

Should/could supporting a new scalar type, say bfloat, int4, etc., be done through cpp extension or does it need to be added to the aten/c10 code in the PyTorch repo where the existing scalar types are defined (ScalarType.h, etc.)?

In either case, is there a tutorial or checklist for this?

smth · December 6, 2018, 6:58am

some fundamental parts need to be added to aten/C10, similar to how we added ComplexFloat.

The rest can be done out-of-source.

If it’s bfloat16, we are interested in adding it into aten/c10 for some accelerators.

soumyarooproy · December 6, 2018, 7:38am

Thanks.

I see that there is already some boilerplate code pertaining to complex float types. Things like ScalarType, TypeID, deferred/lazy registration of the complex types, etc.

Yes, I was referring to bfloat16. Glow had an issue open about adding bfloat type too (https://github.com/pytorch/glow/issues/1329). But, IIANM, it has been deferred.

I’ll explore more.