.. meta:: :description: rocFFT documentation and API reference library :keywords: rocFFT, ROCm, API, documentation .. _codegen: ******************************************************************** Code Generator Design Document for rocFFT ******************************************************************** Proposal ======== Create a new code generator for rocFFT. Rationale --------- The current code generator: * dates from clFFT * is based on string concatenation Ideally, a new code generator: * based on an abstract-syntax-tree (AST) * generates faster, more robust kernels ASTs allow generated code to be transformed and manipulated before being emitted. A concrete example of this for FFT kernels would be: automatically translating a kernel from interleaved to planar format. How the generator is designed and implemented is crucial for both conciseness and ease-of-use. Required kernels (scope) ======================== For rocFFT, we need/want to generate: * Host functions to launch the FFT kernels * Tiled (row/column) + strided + batched Stockham kernels for arbitrary factorization * May want to extend to Cooley-Tukey kernels as well Kernels need to handle all combinations of: * single/double precision (and be extendable to half-float and bfloat) * in-place/out-of-place * planar/interleaved * real/complex * small/large twiddle tables * unit/non-unit stride * transposed output, including with twiddle multiplies for large 1D * fusing with pre and post-processing kernels (e.g. real even-length) Ideally any configuration/runtime parameters required by the kernels would be defined in a single place to avoid repetition between rocFFT and the generator. We have flexibility in handling these combinations at compile-time or run-time. For example, multiple kernels could be generated for single/double precision, but unit/non-unit stride could be handled at runtime. Fundamentally, all multidimensional and batched FFTs can be written in terms of 1D transforms (with affine indexing). As such, an FFT is broken down into: * A *host* function that is aware of dimensions, strides, batches, and tiling. This function would be responsible for determining how the problem will be broken down into GPU thread blocks. * A *global* function that is aware of GPU thread blocks, dimensions, strides, batches, and tiling. This function would be responsible for determining offsets and strides for the device function, and declaring LDS memory buffers. * A *device* function that is passed offsets and strides, and is aware of GPU threads. The device function would perform a (short) 1D transform. A device function may be called so that a thread block is actually transforming multiple batches. As such, indexes (the spatial index in the FFT) should be computed as: .. code-block:: c int fft_index = threadIdx.x % width; Tiling ====== Launching device kernels in a way that traverses memory in tiles will be handled at the host/global level. Kernels need to support reading/writing in columns/rows. These are the block CC/RC/CR flavors (where C and R refer to column and row) of the existing kernels. Strides and batches =================== Host ---- Host/global functions should support arbitrary dimensions, lengths, strides, offsets, and batches. Users should be allowed to store their arrays arbitrarily. For an :math:`N` dimensional dataset, the flat index :math:`a` corresponding to indices :math:`(i_1,\ldots,i_N,i_b)`, where :math:`i_b` is the batch index, is given by .. math:: a(i_1,\ldots,i_N,i_b) = s_b i_b + \sum_{d=1}^N s_d i_d Where :math:`s_d` is the stride along dimension :math:`d`. To support these strides, the device function to compute the FFT along dimension :math:`D` would be passed: .. code-block:: c int offset = 0; offset += batch_index * batch_stride; for (int d=0; d < N; ++d) if (d != D) offset += spatial_index[d] * strides[d]; int stride = strides[D]; For example, in three dimensions, to compute the FFT along the y-dimension given x and z indices ``i`` and ``k`` for batch ``b``, the device function would be passed: .. code-block:: c int offset = 0; offset += b * batch_stride; offset += i * strides[0]; offset += k * strides[2]; int stride = strides[1]; Device ------ Device functions should support arbitrary offsets and strides. Array indexes in device functions should be computed as, eg: .. code-block:: c int fft_index = threadIdx.x % width; int array_index = offset + fft_index * stride; Large twiddle tables ==================== Large 1D transforms are decomposed into multiple transforms. To reduce the size of twiddle tables, rotations can be decomposed into multiple stages as well. For example, the rotation through :math:`2\pi \cdot 280 / 256^2` can be decomposed into :math:`2\pi \cdot 1 / 256 + 2\pi 24 / 256^2`. The resulting twiddle table contains 512 entries instead of 65536 entries. Generated kernels should support these "large twiddle tables". Launching ========= For a specific transform length, the generator is free to choose among several algorithms and related tuning parameters. These choices may influence how the kernel is launched. The generator will create both the kernel and the accompanying struct, which gives indications of how the kernel may be used in both rocFFT and other applications. The generator will populate a function pool with structs of the form .. code-block:: c++ struct ROCFFTKernel { void *device_function = nullptr; std::vector factors; int transforms_per_block = 0; int workgroup_size = 0; // ... }; This moves the responsibility of figuring how a kernel should be launched to the generator. Currently, kernels are launched with: * dimension * number of blocks (batches) * number of threads (threads per batch; kernel parameter) * stream * twiddle table * length(s) * strides * batch count * in/out buffers Implementation ============== The code generator will be implemented in Python using only standard modules. The AST will be represented as a tree structure, with nodes in the tree representing operations, such as assignment, addition, or a block containing multiple operations. Nodes will be represented as objects (e.g., ``Add``) extending the base class ``BaseNode``. Operands will be stored in a simple list called ``args``: .. code-block:: python class BaseNode: args: List[Any] To facilitate building ASTs, the base node will have a constructor that simply stores its arguments as operands: .. code-block:: python class BaseNode: args: List[Any] def __init__(self, *args, **kwargs): self.args = list(args) To facilitate rewriting ASTs, node object's constructors should accept a simple list of argument/operands. This, for example, allows a depth-first tree re-write to be implemented trivially as: .. code-block:: python def depth_first(x, f): '''Depth first traversal of the AST in 'x'. Each node is transformed by 'f(x)'.''' if isinstance(x, BaseNode): y = type(x)(*[ depth_first(a, f) for a in x.args ]) return f(y) return f(x) To emit code, each node must implement ``__str__``. For example: .. code-block:: python class Add(BaseNode): def __str__(self): return ' + '.join([ str(x) for x in self.args ]) Stockham tiling implementation ------------------------------ To support tiling, the *global* function is responsible for loading data from global memory into LDS memory in a tiled manner. Once in LDS memory, a singly strided *device* function performs an interleaved, in-place FFT entirely within LDS. Polymorphism will be used to abstract tiling strategies. Different tiling strategies should extend the ``StockhamTiling`` object and overload the ``load_from_global`` and ``store_to_global`` methods. For example: .. code-block:: python tiling = StockhamTilingRR() scheme = StockhamDeviceKernelUWide() body = StatementList() body += tiling.compute_offsets(...) body += tiling.load_from_global(out=lds, in=global_buffer) body += scheme.fft(lds) body += tiling.store_to_global(out=global_buffer, in=lds) Different tiling strategies may require new template parameters and/or function arguments. Tiling strategies can manipulate the following methods: * ``add_templates`` * ``add_global_arguments`` * ``add_device_arguments`` * ``add_device_call_arguments`` Each of these methods is passed a ``TemplateList`` or ``ArgumentList`` argument, and should return a new template/argument list with any extra parameters added. Large twiddle tables -------------------- Device kernels may need to apply additional twiddles during their execution. These extra twiddle tables are implemented similarly to tiling. Different twiddle table strategies should extend the ``StockhamLargeTwiddles`` object and overload the ``load`` and ``multiply`` methods. Twiddle tables may also require additional templates and arguments. See the Stockham tiling implementation section. Copyright and disclaimer ======================== The information contained herein is for informational purposes only, and is subject to change without notice. While every precaution has been taken in the preparation of this document, it may contain technical inaccuracies, omissions and typographical errors, and AMD is under no obligation to update or otherwise correct this information. Advanced Micro Devices, Inc. makes no representations or warranties with respect to the accuracy or completeness of the contents of this document, and assumes no liability of any kind, including the implied warranties of non-infringement, merchantability or fitness for particular purposes, with respect to the operation or use of AMD hardware, software or other products described herein. No license, including implied or arising by estoppel, to any intellectual property rights is granted by this document. Terms and limitations applicable to the purchase or use of AMD’s products are as set forth in a signed agreement between the parties or in AMD's Standard Terms and Conditions of Sale. AMD is a trademark of Advanced Micro Devices, Inc. Other product names used in this publication are for identification purposes only and may be trademarks of their respective companies. Copyright (C) 2021 - 2024 Advanced Micro Devices, Inc. All rights reserved.