Open nonnull-ca opened 1 year ago
Trivial issues:
- aten::rsub unimplemented.
- aten::eq unimplemented.
Take a look into src/pointwise_ops.cpp
many stuff implemented there and some are trivial to do
No support for
aten::native_batch_norm
Actually it is usually called from norm_layer
and it is on its way for transformers, see https://github.com/artyom-beilis/pytorch_dlprim/discussions/16
No
aten::index_select
support
It looks like something that requires custom but simple GPU kernel... it is probably more complex in backward. If you can give me a simple example (mnist style) of when/how it is used in real world I can look into it.
No half-precision support
Biggest issue there is gemm and winograd kernels - writing their optimised versions is going to be hard. Without it half-precision isn't that useful beside storage. bfloat16 is similar... Currently I just don't have capacity to handle it. I think it is somewhat lower priority since most can be done just using regular float
Not your fault, but there's a surprising number of callerside changes that need to happen when using the private device backend.
Obviously :-)
but these are least of the problems at this point
Is there documentation of the aten API anywhere?
The functions signatures that need to be implemented (of course not all of them) can be found in ATen/RegistrationDeclarations.h
- you can find it in $VIRTUAL_ENV/lib/python3.*/site-packages/torch/include/ATen/RegistrationDeclarations.h
Other than that there some docs in the source code and most functions that appear in aten::
have their exact python copy and have public docs.
I spent a chunk of yesterday seeing how far I could go with getting this running. I ran into a bunch of different issues of varying severity - I don't know if you would prefer them all in one issue or split.
Trivial issues:
aten::rsub
unimplemented.aten::eq
unimplemented.(I'll probably do up a PR for these two later today.)
Simple issues:
1.0 - someTensor
).torch.tensor(1.0) - someTensor
for now.aten::rsub
being unimplemented.aten::any_all_out
unimplemented.Complex issues:
aten::native_batch_norm
without mean & variance. (This assert. Note that this is actually being hit when callinglayer_norm
.)aten::index_select
support.native_batch_norm
, but I suspect there's more than just that.aten
is fully implemented this is nowhere near a dropin replacement backend.torch.cuda.device_count()
/torch.cuda.device
/ etc.'privateusedevice:0'
instead....and that's where I ran out of time.
Is there documentation of the aten API anywhere?