Package: tfopgen
Writing a tensorflow operator requires writing fair amounts of
boilerplate C++ and CUDA code. This script generates code for the CPU
and GPU version of a tensorflow operator. More specifically, given
tensorflowinputs,outputsandattributes, it generates:C++ Header file that defines the operator class, templated on Device.C++ Header file that defines the CPU implementation of the operator.C++ Source file with Shape Function, REGISTER_OP and
REGISTER_KERNEL_BUILDER constructs.Cuda Header that defines the GPU implementation of the operator,
including a CUDA kernel.Cuda Source file with GPU REGISTER_KERNEL_BUILDER’s for the
operator.python unit test case, which constructs random input data, and calls
the operator.Makefile for compiling the operator into a shared library, using g++
and nvcc.RequirementsA tensorflow installation, required for building the operator.pipinstalltensorflowInstallationpipinstalltfopgenUsageThe user should provide a YAML configuration file defining the operator:inputs and optionally, their shapes.outputs and optionally, their outputs.polymorphic type attributes.other attributes.documentation.For example, we can define the outline for aComplexPhaseoperator in thecomplex_phase.ymlfile.---project:astronomylibrary:fouriername:ComplexPhasetype_attrs:-"FT:{float,double}=DT_FLOAT"-"CT:{complex64,complex128}=DT_COMPLEX64"inputs:-["uvw:FT",[null,null,3]]# (ntime, nbl, 3)-["frequency:FT",[null]]# (nchan, )-["lm:FT",[null,2]]# (nsrc, 2)outputs:-["complex_phase:CT",[null,null,null,null]]doc:>Given tensors(1) of (U, V, W) baseline coordinates with shape (ntime, nbl, 3)(2) of (L, M) sky coordinates with shape (nsrc, 2)(3) of frequencies,compute the complex phase with shape (nsrc, ntime, nbl, nchan)We can then run:$tfopgencomplex_phase.ymlto create the following directory structure and files:$treefourier/fourier/├──complex_phase_op_cpu.cpp├──complex_phase_op_cpu.h├──complex_phase_op_gpu.cu├──complex_phase_op_gpu.cuh├──complex_phase_op.h├──Makefile└──test_complex_phase.pyTheprojectandlibraryoptions specify C++ namespaces within
which the operator is created. Additionally, the Makefile will create afourier.soshared library that can be loaded withtf.load_op_library('fourier.so').Any polymorphic type attributes should be supplied. The generator will
template the operators on type attributes. It will also generate
concrete permutations of REGISTER_KERNEL_BUILDER for both the CPU and
GPU op using the actual types supplied in the type attributes (float,
double, complex64 and complex128) below:type_attrs:-"FT:{float,double}=DT_FLOAT"-"CT:{complex64,complex128}=DT_COMPLEX64"The operator inputs and their optional shapes should be specified as a
list containing a string defining the.Inputdirective, and a list
describing the shape of the input tensor. Anullvalue in the shape
will be translated into a pythonNone. If concrete dimensions are specified,
corresponding checks will be generated in the Shape Function associated with the
operator.inputs:-["uvw:FT",[null,null,3]]# (ntime, nbl, 3)-["frequency:FT",[null]]# (nchan, )-["lm:FT",[null,2]]# (nsrc, 2)The operator outputs should similarly defined.outputs:-["complex_phase:CT",[null,null,null,null]]Given these inputs and outputs, CPU and GPU operators are created with
named variables corresponding to the inputs and outputs. Additionally, a
CUDA kernel with the given inputs and outputs is created, as well as a
shape function checking the rank and dimensions of the supplied inputs.Other attributes may be specified (and will be output in the
REGISTER_OP) directive, but are not catered for automatically by the
generator code as the range of attribute behaviour is complex.op_other_attrs:-"iterations:int32>=2",Finally operator documentation may also be supplied.doc:>Given tensors(1) of (U, V, W) baseline coordinates with shape (ntime, nbl, 3)(2) of (L, M) sky coordinates with shape (nsrc, 2)(3) of frequencies,compute the complex phase with shape (nsrc, ntime, nbl, nchan)