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.