diff options
Diffstat (limited to 'compute_kernel_writer/prototype/include/ckw/KernelWriter.h')
-rw-r--r-- | compute_kernel_writer/prototype/include/ckw/KernelWriter.h | 217 |
1 files changed, 217 insertions, 0 deletions
diff --git a/compute_kernel_writer/prototype/include/ckw/KernelWriter.h b/compute_kernel_writer/prototype/include/ckw/KernelWriter.h new file mode 100644 index 0000000000..a2778a9485 --- /dev/null +++ b/compute_kernel_writer/prototype/include/ckw/KernelWriter.h @@ -0,0 +1,217 @@ +/* + * Copyright (c) 2023 Arm Limited. + * + * SPDX-License-Identifier: MIT + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to + * deal in the Software without restriction, including without limitation the + * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or + * sell copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +#ifndef CKW_PROTOTYPE_INCLUDE_CKW_KERNELWRITER_H +#define CKW_PROTOTYPE_INCLUDE_CKW_KERNELWRITER_H + +#include "ckw/Kernel.h" +#include "ckw/TensorInfo.h" +#include "ckw/TensorOperand.h" +#include "ckw/TileInfo.h" +#include "ckw/TileOperand.h" + +#include <memory> + +namespace ckw +{ + +namespace prototype +{ +struct GpuKernelWriterAttribute; +class IGpuKernelWriter; +} // namespace prototype + +/** Kernel writer. */ +class KernelWriter +{ +public: + // ============================================================================================= + // Constructors and destructor + // ============================================================================================= + + /** Initialize a new instance of kernel writer. + * + * @param[in] kernel The kernel to be written to. + */ + explicit KernelWriter(Kernel &kernel); + + /** Destructor */ + ~KernelWriter(); + + /** No copy constructor. */ + KernelWriter(const KernelWriter &) = delete; + + /** No copy assignment. */ + KernelWriter &operator=(const KernelWriter &) = delete; + + // ============================================================================================= + // Scope management + // ============================================================================================= + + /** Get the current ID space. */ + int32_t id_space() const; + + /** Set the current ID space. */ + KernelWriter &id_space(int32_t id_space); + + /** Switch to and return a new ID space. */ + int32_t next_id_space(); + + // ============================================================================================= + // Tensor and tile declaration + // ============================================================================================= + + /** Define a tensor argument. + * + * @param[in] name The name of the tensor. + * @param[in] info The tensor info. + * + * @return The @ref TensorOperand object. + */ + TensorOperand &create_tensor_argument(const char *name, const TensorInfo &info); + + /** Define a compile-time constant scalar argument. + * + * @param[in] name The name of the tile. + * @param[in] value The value of the tile. + * + * @return The @ref TileOperand object. + */ + TileOperand &create_tile_argument(const char *name, int32_t value); + + /** Declare a new tile. + * + * The name of the tile must be unique in the current ID space. + * + * @param[in] name The name of the tile. + * @param[in] ... The necessary arguments to create a new @ref TileOperand. + * + * @return The @ref TileOperand object. + */ + template <typename... TArgs> + TileOperand &declare_tile(const char *name, TArgs &&...args) + { + const auto var_name = generate_variable_name(name); + auto operand = new TileOperand(var_name, ::std::forward<TArgs>(args)...); + register_operand(operand, true); + + return *operand; + } + + // ============================================================================================= + // Load and store + // ============================================================================================= + + /** Load the data from the tensor memory to the tile using the sampling information. + * + * @param[out] tile The tile to be loaded. + * @param[in] tensor The tensor to be read. + * @param[in] sampler The tensor sampling information. + */ + void op_load(TileOperand &tile, TensorOperand &tensor, const TensorTileSampler &sampler); + + /** Store the tile to the tensor using the specified sampling information. + * + * @param[out] dst The tensor that the tile is written to. + * @param[in] src The tile to be stored. + * @param[in] sampler The tensor sampling information. + */ + void op_store(TensorOperand &tensor, const TileOperand &tile, const TensorTileSampler &sampler); + + // ============================================================================================= + // Data processing + // ============================================================================================= + + /** Write assignment: `<dst> = <src>`. + * + * @param[in] dst The destination tile. + * @param[in] src The source tile. + */ + void op_assign(TileOperand &dst, const TileOperand &src); + + /** Write binary expression: `<dst> = <lhs> <op> <rhs>`. + * + * @param[in] dst The destination tile. + * @param[in] lhs The LHS operand. + * @param[in] rhs The RHS operand. + * @param[in] op The binary operator. + */ + void op_binary_expression(TileOperand &dst, const TileOperand &lhs, const TileOperand &rhs, BinaryOp op); + + /** Write function applied to scalar value: `<dst> = <func>(<src>)`. + * + * @param[in] dst The destination tile. + * @param[in] src The source tile. + * @param[in] func The function to be applied to the source tile. + */ + void op_scalar_function(TileOperand &dst, const TileOperand &src, ScalarUnaryFunction func); + + // ============================================================================================= + // Misc + // ============================================================================================= + + /** Set `dst` the global ID of dimension `dim`. + * + * @param[in] dst The tile to be written to. + * @param[in] dim The global ID dimension. + */ + void op_get_global_id(TileOperand &dst, int32_t dim); + + // ============================================================================================= + // Code generation + // ============================================================================================= + + /** Generate the source code of the kernel. */ + ::std::string generate_code(); + +private: + /** Generate the full variable name based on the original name and the ID space. + * + * @param[in] name The name of the variable. + * + * @return The full variable name. + */ + ::std::string generate_variable_name(const char *name) const; + + /** Register the operand to the kernel. + * + * The operand is uniquely owned by the kernel afterward. + * + * @param[in] operand The operand to be registered. + * @param[in] declaring Whether the tile declaration is generated. + */ + void register_operand(OperandBase *operand, bool declaring); + +private: + Kernel *_kernel; + ::std::unique_ptr<prototype::GpuKernelWriterAttribute> _impl_attr; + ::std::unique_ptr<prototype::IGpuKernelWriter> _impl; + + int32_t _id_space{ 0 }; + int32_t _max_id_space{ 0 }; +}; + +} // namespace ckw + +#endif // CKW_PROTOTYPE_INCLUDE_CKW_KERNELWRITER_H |