Add tile declaration capability in KernelWriter

Resolves: COMPMID-5816

Signed-off-by: Gunes Bayir <gunes.bayir@arm.com>
Change-Id: Ibd885707a842550a058252f9d01e072129896055
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/9901
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Viet-Hoa Do <viet-hoa.do@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Benchmark: Arm Jenkins <bsgcomp@arm.com>

11 files changed, 260 insertions, 13 deletions

diff --git a/compute_kernel_writer/include/ckw/Kernel.h b/compute_kernel_writer/include/ckw/Kernel.h
index 0cab713c48..54e7ca33fd 100644
--- a/compute_kernel_writer/include/ckw/Kernel.h
+++ b/compute_kernel_writer/include/ckw/Kernel.h
@@ -25,12 +25,17 @@
 #ifndef CKW_INCLUDE_CKW_KERNEL_H
 #define CKW_INCLUDE_CKW_KERNEL_H
 
-#include "ckw/types/TargetLanguage.h"
 #include <string>
 
 namespace ckw
 {
 
+// Forward Declerations
+class TileInfo;
+class ITileOperand;
+
+enum class TargetLanguage;
+
 /** The kernel that has been emitted by the kernel writer.
  *
  * It contains all the necessary information to compile and execute the kernel.
@@ -38,6 +43,8 @@ namespace ckw
 class Kernel
 {
 public:
+    virtual ~Kernel();
+
     /** Initialize a new instance of @ref Kernel class with all emitted kernel information.
      *
      * @param[in] language    The target language of the kernel.
@@ -51,6 +58,9 @@ public:
     /** Get the source code. */
     const std::string &source_code() const;
 
+    /** Add a tile operand */
+    virtual ITileOperand &add_operand(const std::string &name, const TileInfo &tile_info) = 0;
+
 private:
     TargetLanguage _language;
     std::string    _source_code;
diff --git a/compute_kernel_writer/include/ckw/KernelWriter.h b/compute_kernel_writer/include/ckw/KernelWriter.h
index ba8a6015e6..f1635c6449 100644
--- a/compute_kernel_writer/include/ckw/KernelWriter.h
+++ b/compute_kernel_writer/include/ckw/KernelWriter.h
@@ -25,9 +25,10 @@
 #ifndef CKW_INCLUDE_CKW_KERNELWRITER_H
 #define CKW_INCLUDE_CKW_KERNELWRITER_H
 
-#include "ckw/types/TargetArchitecture.h"
-#include "ckw/types/TargetLanguage.h"
+#include "ckw/ITileOperand.h"
+
 #include <memory>
+#include <set>
 #include <string>
 
 namespace ckw
@@ -35,6 +36,11 @@ namespace ckw
 
 class Kernel;
 
+/** Forward Declerations */
+class TileInfo;
+enum class TargetArchitecture;
+enum class TargetLanguage;
+
 /** A kernel writer.
  *
  * This class is used to construct a new kernel by defining arguments, declaring variable and writing code.
@@ -71,6 +77,7 @@ public:
     // =============================================================================================
 
     /** Write the line comment in debug build.
+     *
      * This function does not take effect on release build.
      *
      * The comment must only contain one line (i.e. no newline character is allowed).
@@ -88,6 +95,31 @@ public:
      * @param[in] name The name of the kernel object to be generated.
      */
     virtual std::unique_ptr<Kernel> emit_kernel(const std::string &name) = 0;
+
+    /** Declare a tile given its name and tile info
+     *
+     * @param[in] name Name of the tile
+     * @param[in] tile_info Shape and data type of the tile
+     *
+     * @returns The created tile operand
+     */
+    virtual ITileOperand &declare_tile(const std::string &name, const TileInfo &tile_info) = 0;
+
+protected:
+    int32_t id_space() const;
+
+    /** Pure virtual function to be overridden by language specific subclasses to add a tile operand to the kernel */
+    virtual ITileOperand &add_operand(const std::string &name, const TileInfo &tile_info) = 0;
+
+    /** Add a tile operand to the operand list */
+    ITileOperand &add_operand(std::unique_ptr<ITileOperand> &operand_ptr);
+
+    /** Generate full variable name by prefixing it with id space */
+    std::string generate_full_name(const std::string &name) const;
+
+private:
+    int32_t _id_space{ 0 };
+    std::set<std::unique_ptr<ITileOperand>> _operands {};
 };
 
 } // namespace ckw
diff --git a/compute_kernel_writer/src/Kernel.cpp b/compute_kernel_writer/src/Kernel.cpp
index ccc68ecefb..5eea1aa548 100644
--- a/compute_kernel_writer/src/Kernel.cpp
+++ b/compute_kernel_writer/src/Kernel.cpp
@@ -23,10 +23,13 @@
  */
 
 #include "ckw/Kernel.h"
+#include "ckw/types/TargetLanguage.h"
 
 namespace ckw
 {
 
+Kernel::~Kernel() = default;
+
 Kernel::Kernel(TargetLanguage language, const std::string &source_code)
     : _language(language), _source_code(source_code)
 {
diff --git a/compute_kernel_writer/src/KernelWriter.cpp b/compute_kernel_writer/src/KernelWriter.cpp
index eb8399a7ef..ab5ede8c77 100644
--- a/compute_kernel_writer/src/KernelWriter.cpp
+++ b/compute_kernel_writer/src/KernelWriter.cpp
@@ -22,15 +22,22 @@
  * SOFTWARE.
  */
 
-#include "ckw/KernelWriter.h"
 #include "ckw/Error.h"
+#include "ckw/ITileOperand.h"
+#include "ckw/KernelWriter.h"
+#include "ckw/types/TargetArchitecture.h"
+#include "ckw/types/TargetLanguage.h"
 #include "src/cl/CLKernelWriter.h"
 
+#include <iterator>
 namespace ckw
 {
 
+KernelWriter::~KernelWriter() = default;
+
 std::unique_ptr<KernelWriter> KernelWriter::create_instance(TargetArchitecture architecture, TargetLanguage language)
 {
+    CKW_UNUSED(architecture);
     switch(language)
     {
         case TargetLanguage::OpenCL:
@@ -43,6 +50,20 @@ std::unique_ptr<KernelWriter> KernelWriter::create_instance(TargetArchitecture a
     }
 }
 
-KernelWriter::~KernelWriter() = default;
+int32_t KernelWriter::id_space() const
+{
+    return _id_space;
+}
+
+ITileOperand &KernelWriter::add_operand(std::unique_ptr<ITileOperand> &operand_ptr)
+{
+    auto it = _operands.insert(std::move(operand_ptr));
+    return *it.first->get();
+}
+
+std::string KernelWriter::generate_full_name(const std::string &name) const
+{
+    return "G" + std::to_string(id_space()) + "__" + name;
+}
 
 } // namespace ckw
diff --git a/compute_kernel_writer/src/cl/CLHelpers.cpp b/compute_kernel_writer/src/cl/CLHelpers.cpp
index d940a5a529..5a3d0fab81 100644
--- a/compute_kernel_writer/src/cl/CLHelpers.cpp
+++ b/compute_kernel_writer/src/cl/CLHelpers.cpp
@@ -88,4 +88,36 @@ std::string cl_get_variable_datatype_as_string(DataType dt, int32_t len)
 
     return res;
 }
-} // namespace ckw
\ No newline at end of file
+
+int32_t width_to_cl_vector_size(int32_t width)
+{
+    switch(width)
+    {
+        case 1:
+            return 1;
+        case 2:
+            return 2;
+        case 3:
+            return 3;
+        case 4:
+            return 4;
+        case 5:
+        case 6:
+        case 7:
+        case 8:
+            return 8;
+        case 9:
+        case 10:
+        case 11:
+        case 12:
+        case 13:
+        case 14:
+        case 15:
+        case 16:
+            return 16;
+        default:
+            CKW_THROW_MSG("Unsupported width to convert to OpenCL vector");
+            return 0;
+    }
+}
+} // namespace ckw
diff --git a/compute_kernel_writer/src/cl/CLHelpers.h b/compute_kernel_writer/src/cl/CLHelpers.h
index 915d59f458..a9a84e2187 100644
--- a/compute_kernel_writer/src/cl/CLHelpers.h
+++ b/compute_kernel_writer/src/cl/CLHelpers.h
@@ -21,10 +21,11 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  * SOFTWARE.
  */
-#ifndef COMPUTE_KERNEL_WRITER_SRC_CL_CLHELPERS_H
-#define COMPUTE_KERNEL_WRITER_SRC_CL_CLHELPERS_H
+#ifndef CKW_SRC_CL_CLHELPERS_H
+#define CKW_SRC_CL_CLHELPERS_H
 
 #include <string>
+#include <cstdint>
 
 /** OpenCL specific helper functions */
 namespace ckw
@@ -48,6 +49,15 @@ bool cl_validate_vector_length(int32_t len);
  * @return the OpenCL datatype as a string
  */
 std::string cl_get_variable_datatype_as_string(DataType dt, int32_t len);
+
+/** Helper function to return the OpenCL vector size that accommodate the the desired width
+ *
+ * @param[in] width The desired width
+ *
+ * @return the OpenCL vector size
+*/
+int32_t width_to_cl_vector_size(int32_t width);
+
 } // namespace ckw
 
 #endif /* COMPUTE_KERNEL_WRITER_SRC_CL_CLHELPERS_H */
diff --git a/compute_kernel_writer/src/cl/CLKernelWriter.cpp b/compute_kernel_writer/src/cl/CLKernelWriter.cpp
index 231d321c22..7faf2e6d60 100644
--- a/compute_kernel_writer/src/cl/CLKernelWriter.cpp
+++ b/compute_kernel_writer/src/cl/CLKernelWriter.cpp
@@ -24,14 +24,14 @@
 
 #include "src/cl/CLKernelWriter.h"
 #include "ckw/Error.h"
+#include "src/cl/CLTile.h"
+#include "src/cl/CLHelpers.h"
+#include <cstdint>
 
 namespace ckw
 {
 
-CLKernelWriter::CLKernelWriter()
-{
-}
-
+CLKernelWriter::CLKernelWriter() = default;
 CLKernelWriter::~CLKernelWriter() = default;
 
 std::unique_ptr<Kernel> CLKernelWriter::emit_kernel(const std::string &name)
@@ -61,4 +61,27 @@ const std::string &CLKernelWriter::body_source_code() const
     return _body_source_code;
 }
 
+ITileOperand &CLKernelWriter::declare_tile(const std::string &name, const TileInfo &tile_info)
+{
+    const std::string fullname = generate_full_name(name);
+
+    const int32_t height = tile_info.height();
+    const int32_t width = tile_info.width();
+    const DataType data_type = tile_info.data_type();
+
+    for(int32_t row = 0; row < height; ++row)
+    {
+        const std::string cl_type = cl_get_variable_datatype_as_string(data_type, width);
+        append_code(cl_type, " ", fullname, std::to_string(row), ";\n");
+    }
+
+    return add_operand(fullname, tile_info);
+}
+
+ITileOperand &CLKernelWriter::add_operand(const std::string &name, const TileInfo &tile_info)
+{
+    std::unique_ptr<ITileOperand> operand = std::make_unique<CLTile>(name, tile_info);
+    return KernelWriter::add_operand(operand);
+}
+
 } // namespace ckw
diff --git a/compute_kernel_writer/src/cl/CLKernelWriter.h b/compute_kernel_writer/src/cl/CLKernelWriter.h
index e6f0641538..d0c4b7c9d4 100644
--- a/compute_kernel_writer/src/cl/CLKernelWriter.h
+++ b/compute_kernel_writer/src/cl/CLKernelWriter.h
@@ -57,6 +57,13 @@ public:
 
     std::unique_ptr<Kernel> emit_kernel(const std::string &name) override;
 
+    /** Declare a tile given name and tile information
+     *
+     * Similar to @ref KernelWriter::declare_tile()
+    */
+    ITileOperand &declare_tile(const ::std::string &name, const TileInfo &tile_info) override;
+
+
 protected:
     /** Append the specified code to the kernel body source code. */
     template <typename T, typename... TArgs>
@@ -76,6 +83,9 @@ protected:
     /** Get the current kernel body source code. */
     const std::string &body_source_code() const;
 
+    /** Add a tile operand to the kernel and return it */
+    ITileOperand &add_operand(const std::string &code, const TileInfo &tile_info) override;
+
 private:
     /** This string contains the kernel body source code, not the full CL source code.
      * The full source code will only be generated when the user calls @ref KernelWriter::emit_kernel.
diff --git a/compute_kernel_writer/validation/Validation.cpp b/compute_kernel_writer/validation/Validation.cpp
index c0bdaaa0dc..e4884fa4aa 100644
--- a/compute_kernel_writer/validation/Validation.cpp
+++ b/compute_kernel_writer/validation/Validation.cpp
@@ -22,11 +22,12 @@
  * SOFTWARE.
  */
 
+#include "tests/CLKernelWriterCommentTest.h"
+#include "tests/CLKernelWriterDeclareTileTest.h"
 #include "tests/CLConstantTileTest.hpp"
 #include "tests/CLTileTest.hpp"
 #include "tests/TensorBitMaskTest.hpp"
 #include "tests/UtilsTest.hpp"
-#include "tests/CLKernelWriterTest.h"
 
 #include <memory>
 #include <vector>
@@ -60,7 +61,10 @@ int32_t main()
     const auto test12 = std::make_unique<CLConstantTileAccessScalarVariableBroadcastYTest>();
     const auto test13 = std::make_unique<CLConstantTileAccessVectorVariablesTest>();
     const auto test14 = std::make_unique<CLConstantTileAccessSubVectorVariablesTest>();
+#ifdef COMPUTE_KERNEL_WRITER_DEBUG_ENABLED
     const auto test15 = std::make_unique<CLKernelWriterCommentTest>();
+#endif /* COMPUTE_KERNEL_WRITER_DEBUG_ENABLED */
+    const auto test16 = std::make_unique<CLKernelWriterDeclareTileTest>();
 
     tests.push_back(test3.get());
     tests.push_back(test4.get());
@@ -74,7 +78,10 @@ int32_t main()
     tests.push_back(test12.get());
     tests.push_back(test13.get());
     tests.push_back(test14.get());
+#ifdef COMPUTE_KERNEL_WRITER_DEBUG_ENABLED
     tests.push_back(test15.get());
+#endif /* COMPUTE_KERNEL_WRITER_DEBUG_ENABLED */
+    tests.push_back(test16.get());
 #endif /* COMPUTE_KERNEL_WRITER_OPENCL_ENABLED */
 
     bool all_test_passed = true;
diff --git a/compute_kernel_writer/validation/tests/CLKernelWriterTest.h b/compute_kernel_writer/validation/tests/CLKernelWriterCommentTest.h
index ff09ea8073..ff09ea8073 100644
--- a/compute_kernel_writer/validation/tests/CLKernelWriterTest.h
+++ b/compute_kernel_writer/validation/tests/CLKernelWriterCommentTest.h
diff --git a/compute_kernel_writer/validation/tests/CLKernelWriterDeclareTileTest.h b/compute_kernel_writer/validation/tests/CLKernelWriterDeclareTileTest.h
new file mode 100644
index 0000000000..5e00084aaa
--- /dev/null
+++ b/compute_kernel_writer/validation/tests/CLKernelWriterDeclareTileTest.h
@@ -0,0 +1,99 @@
+/*
+ * 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_VALIDATION_TESTS_CLKERNELWRITER_H
+#define CKW_VALIDATION_TESTS_CLKERNELWRITER_H
+
+#include "ckw/TileInfo.h"
+#include "ckw/types/DataType.h"
+#include "src/cl/CLKernelWriter.h"
+#include "validation/tests/common/KernelWriterInterceptor.h"
+#include "validation/tests/common/Common.h"
+
+#include <vector>
+
+namespace ckw
+{
+
+using CLKernelWriterDeclareTileConfig = std::tuple<DataType, int32_t, int32_t, std::string>;
+
+class CLKernelWriterDeclareTileTest : public ITest
+{
+public:
+    CLKernelWriterDeclareTileTest()
+    {
+        _configs = {
+            {DataType::Fp32, 4, 4, "float4 G0__a_tile"},
+            {DataType::Uint8, 4, 1, "uchar G0__a_tile"},
+            {DataType::Int8, 4, 2, "char2 G0__a_tile"},
+            {DataType::Bool, 9, 3, "bool3 G0__a_tile"},
+            {DataType::Fp16, 4, 16, "half16 G0__a_tile"},
+            {DataType::Uint32, 1, 8, "uint8 G0__a_tile"},
+            {DataType::Uint16, 2, 3, "ushort3 G0__a_tile"},
+        };
+    }
+
+    bool run() override
+    {
+        bool all_tests_passed = true;
+        int32_t test_idx = 0;
+
+        for(auto _config: _configs)
+        {
+            KernelWriterInterceptor<CLKernelWriter> writer;
+            writer.start_capture_code();
+
+            const DataType data_type = std::get<0>(_config);
+            const int32_t height = std::get<1>(_config);
+            const int32_t width = std::get<2>(_config);
+            const std::string prefix = std::get<3>(_config);
+
+            // expected output
+            std::string expected_code = "";
+            for(int32_t row = 0; row < height; ++row)
+            {
+                expected_code += prefix + std::to_string(row) + ";\n";
+            }
+
+            TileInfo tile_info(data_type, height, width);
+            writer.declare_tile("a_tile", tile_info);
+
+            VALIDATE_TEST(writer.check_added_code(expected_code), all_tests_passed, test_idx++);
+        }
+
+        return all_tests_passed;
+    }
+
+    std::string name() override
+    {
+        return "CLKernelWriterDeclareTileTest";
+    }
+
+private:
+    std::vector<CLKernelWriterDeclareTileConfig> _configs {};
+};
+
+} // namespace ckw
+
+#endif /* CKW_VALIDATION_TESTS_CLKERNELWRITER_H */