Skip to content

Adding new unary operators to pytorch

Pearu Peterson edited this page Sep 18, 2019 · 1 revision

To implement a new unary operation to pytrorch, one needs to change the following files (for required changes, use some existing unary operation (say, exp) as a template and do copy-paste-rename):

  1. aten/src/ATen/native/native_functions.yaml
  2. aten/src/ATen/native/UnaryOps.h
  3. aten/src/ATen/native/UnaryOps.cpp
  4. aten/src/ATen/native/cpu/UnaryOpsKernel.cpp
  5. aten/src/ATen/cpu/vml.h
  6. aten/src/ATen/cpu/vec256/vec256_{base, double, float, ...}.h

The following files will be re-generated (requires using environment variable GEN_TO_SOURCE=1):

  1. aten/src/ATen/core/TensorMethods.h
  2. aten/src/ATen/core/TensorBody.h
  3. aten/src/ATen/core/OpsAlreadyMovedToC10.cpp

The may exist several libraries that implement the same math function: C++ std, sleef, mkl, vml, etc. Enable the one that is fastest, for instance, for some functions sleef implementation is more performant that of mkl while pytorch uses mkl by default when available.

Clone this wiki locally