Add CKW flow control writing methods

* Structures: if/else/else if, for, return.
* Add corresponding tests.

Partially resolves: COMPMID-6387
Signed-off-by: Viet-Hoa Do <viet-hoa.do@arm.com>
Change-Id: I2912ccaf46f836907f21bb53fa82bcc1f48dd224
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/10199
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: SiCong Li <sicong.li@arm.com>
Reviewed-by: Gunes Bayir <gunes.bayir@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Benchmark: Arm Jenkins <bsgcomp@arm.com>

6 files changed, 186 insertions, 6 deletions

diff --git a/compute_kernel_writer/src/ITile.cpp b/compute_kernel_writer/src/ITile.cpp
new file mode 100644
index 0000000000..eeb7816068
--- /dev/null
+++ b/compute_kernel_writer/src/ITile.cpp
@@ -0,0 +1,35 @@
+/*
+ * 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.
+ */
+
+#include "src/ITile.h"
+
+namespace ckw
+{
+
+bool ITile::is_scalar() const
+{
+    return info().width() == 1 && info().height() == 1;
+}
+
+} // namespace ckw
diff --git a/compute_kernel_writer/src/ITile.h b/compute_kernel_writer/src/ITile.h
index b5585abcdf..73b7315fb5 100644
--- a/compute_kernel_writer/src/ITile.h
+++ b/compute_kernel_writer/src/ITile.h
@@ -21,8 +21,8 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  * SOFTWARE.
  */
-#ifndef CKW_SRC_ITILE
-#define CKW_SRC_ITILE
+#ifndef CKW_SRC_ITILE_H
+#define CKW_SRC_ITILE_H
 
 #include "ckw/TileInfo.h"
 
@@ -100,7 +100,7 @@ public:
 /** Tile base class.
  *  A Tile is a collection of variables (either program variables or constants) used to express a 2D data.
  */
-class ITile: public IScalarAccess
+class ITile : public IScalarAccess
 {
 public:
     virtual ~ITile() = default;
@@ -129,7 +129,13 @@ public:
      * @return true if the tile is assignable
      */
     virtual bool is_assignable() const = 0;
+
+    /** Get whether the tile is scalar, i.e. the width and height are both 1.
+     *
+     * @return true if the tile is scalar.
+     */
+    bool is_scalar() const;
 };
 } // namespace ckw
 
-#endif /* CKW_SRC_ITILE */
+#endif // CKW_SRC_ITILE_H
diff --git a/compute_kernel_writer/src/KernelWriter.cpp b/compute_kernel_writer/src/KernelWriter.cpp
index 21a61d73bf..fb0e62c8ce 100644
--- a/compute_kernel_writer/src/KernelWriter.cpp
+++ b/compute_kernel_writer/src/KernelWriter.cpp
@@ -51,11 +51,35 @@ std::unique_ptr<KernelWriter> KernelWriter::create_instance(TargetArchitecture a
     }
 }
 
+int32_t KernelWriter::new_id_space()
+{
+    _id_space = ++_last_created_id_space;
+
+    return _id_space;
+}
+
 int32_t KernelWriter::id_space() const
 {
     return _id_space;
 }
 
+KernelWriter &KernelWriter::id_space(int32_t value)
+{
+    CKW_ASSERT(value <= _last_created_id_space);
+
+    _id_space = value;
+
+    return *this;
+}
+
+void KernelWriter::write_body(const std::function<void()> &body)
+{
+    const auto curr_id_space = id_space();
+    new_id_space();
+    body();
+    id_space(curr_id_space);
+}
+
 std::string KernelWriter::generate_full_name(const std::string &name) const
 {
     return "G" + std::to_string(id_space()) + "__" + name;
diff --git a/compute_kernel_writer/src/cl/CLHelpers.cpp b/compute_kernel_writer/src/cl/CLHelpers.cpp
index e12e5e1b13..ff4408b1a3 100644
--- a/compute_kernel_writer/src/cl/CLHelpers.cpp
+++ b/compute_kernel_writer/src/cl/CLHelpers.cpp
@@ -26,6 +26,7 @@
 
 #include "ckw/Error.h"
 #include "ckw/types/DataType.h"
+#include "ckw/types/Operators.h"
 #include "ckw/types/TensorStorageType.h"
 #include "src/types/DataTypeHelpers.h"
 
@@ -145,6 +146,21 @@ std::string cl_get_variable_storagetype_as_string(TensorStorageType storage)
     return res;
 }
 
+std::string cl_get_assignment_op_as_string(AssignmentOp op)
+{
+    switch(op)
+    {
+        case AssignmentOp::Increment:
+            return "+=";
+
+        case AssignmentOp::Decrement:
+            return "-=";
+
+        default:
+            CKW_THROW_MSG("Unsupported assignment operator!");
+    }
+}
+
 std::tuple<bool, std::string> cl_get_unary_op(UnaryOp op)
 {
     switch(op)
diff --git a/compute_kernel_writer/src/cl/CLKernelWriter.cpp b/compute_kernel_writer/src/cl/CLKernelWriter.cpp
index 79d0f985d0..90707ccbb1 100644
--- a/compute_kernel_writer/src/cl/CLKernelWriter.cpp
+++ b/compute_kernel_writer/src/cl/CLKernelWriter.cpp
@@ -82,8 +82,7 @@ std::unique_ptr<Kernel> CLKernelWriter::emit_kernel(const std::string &name)
             const auto &tile      = component->tile();
             const auto &tile_info = tile.info();
 
-            CKW_ASSERT(tile_info.height() == 1);
-            CKW_ASSERT(tile_info.width() == 1);
+            CKW_ASSERT(tile.is_scalar());
 
             code += cl_get_variable_datatype_as_string(tile_info.data_type(), 1);
             code += " ";
@@ -315,6 +314,77 @@ void CLKernelWriter::op_ternary(const TileOperand &dst, TernaryOp op, const Tile
     }
 }
 
+void CLKernelWriter::op_if_generic(bool is_else, const TileOperand &lhs, BinaryOp op, const TileOperand &rhs, const std::function<void()> &body)
+{
+    const auto &lhs_tile = to_cl_tile(lhs);
+    const auto &rhs_tile = to_cl_tile(rhs);
+
+    const auto op_name = std::get<1>(cl_get_binary_op(op, lhs_tile.info().data_type()));
+    CKW_ASSERT(op == BinaryOp::Less || op == BinaryOp::LessEqual || op == BinaryOp::Equal || op == BinaryOp::GreaterEqual || op == BinaryOp::Greater);
+
+    CKW_ASSERT(lhs_tile.is_scalar());
+    CKW_ASSERT(rhs_tile.is_scalar());
+
+    if(is_else)
+    {
+        append_code("else ");
+    }
+
+    append_code("if (", lhs_tile.scalar(0, 0).str, " ", op_name, " ", rhs_tile.scalar(0, 0).str, ")\n{\n");
+    write_body(body);
+    append_code("}\n");
+}
+
+void CLKernelWriter::op_if(const TileOperand &lhs, BinaryOp op, const TileOperand &rhs, const std::function<void()> &body)
+{
+    op_if_generic(false, lhs, op, rhs, body);
+}
+
+void CLKernelWriter::op_else_if(const TileOperand &lhs, BinaryOp op, const TileOperand &rhs, const std::function<void()> &body)
+{
+    op_if_generic(true, lhs, op, rhs, body);
+}
+
+void CLKernelWriter::op_else(const std::function<void()> &body)
+{
+    append_code("else\n{\n");
+    write_body(body);
+    append_code("}\n");
+}
+
+void CLKernelWriter::op_for_loop(
+    const TileOperand &var, BinaryOp cond_op, const TileOperand &cond_value,
+    const TileOperand &update_var, AssignmentOp update_op, const TileOperand &update_value,
+    const std::function<void()> &body)
+{
+    const auto &var_tile          = to_cl_tile(var);
+    const auto &cond_value_tile   = to_cl_tile(cond_value);
+    const auto &update_var_tile   = to_cl_tile(update_var);
+    const auto &update_value_tile = to_cl_tile(update_value);
+
+    CKW_ASSERT(var_tile.is_scalar());
+    CKW_ASSERT(cond_value_tile.is_scalar());
+    CKW_ASSERT(update_var_tile.is_scalar());
+    CKW_ASSERT(update_value_tile.is_scalar());
+
+    CKW_ASSERT(var_tile.info().data_type() == cond_value_tile.info().data_type());
+    CKW_ASSERT(update_var_tile.info().data_type() == update_value_tile.info().data_type());
+
+    const auto cond_op_name = std::get<1>(cl_get_binary_op(cond_op, var_tile.info().data_type()));
+    CKW_ASSERT(cond_op == BinaryOp::Less || cond_op == BinaryOp::LessEqual || cond_op == BinaryOp::Equal || cond_op == BinaryOp::GreaterEqual || cond_op == BinaryOp::Greater);
+
+    append_code(
+        "for (; ", var_tile.scalar(0, 0).str, " ", cond_op_name, " ", cond_value_tile.scalar(0, 0).str, "; ",
+        update_var_tile.scalar(0, 0).str, " ", cl_get_assignment_op_as_string(update_op), " ", update_value_tile.scalar(0, 0).str, ")\n{\n");
+    write_body(body);
+    append_code("}\n");
+}
+
+void CLKernelWriter::op_return()
+{
+    append_code("return;\n");
+}
+
 void CLKernelWriter::op_comment(const std::string &text)
 {
 #ifdef COMPUTE_KERNEL_WRITER_DEBUG_ENABLED
diff --git a/compute_kernel_writer/src/cl/CLKernelWriter.h b/compute_kernel_writer/src/cl/CLKernelWriter.h
index d2c84f192e..9fc7e550a7 100644
--- a/compute_kernel_writer/src/cl/CLKernelWriter.h
+++ b/compute_kernel_writer/src/cl/CLKernelWriter.h
@@ -75,6 +75,23 @@ public:
     void op_ternary(const TileOperand &dst, TernaryOp op, const TileOperand &first, const TileOperand &second, const TileOperand &third) override;
 
     // =============================================================================================
+    // Flow control
+    // =============================================================================================
+
+    void op_if(const TileOperand &lhs, BinaryOp op, const TileOperand &rhs, const std::function<void()> &body) override;
+
+    void op_else_if(const TileOperand &lhs, BinaryOp op, const TileOperand &rhs, const std::function<void()> &body) override;
+
+    void op_else(const std::function<void()> &body) override;
+
+    void op_for_loop(
+        const TileOperand &var, BinaryOp cond_op, const TileOperand &cond_value,
+        const TileOperand &update_var, AssignmentOp update_op, const TileOperand &update_value,
+        const std::function<void()> &body) override;
+
+    void op_return() override;
+
+    // =============================================================================================
     // Misc
     // =============================================================================================
 
@@ -177,6 +194,18 @@ private:
         const TileOperand &x, const TileOperand &y, const TileOperand &z, const TileOperand &batch,
         const CLTile &dilation_x, const CLTile &dilation_y);
 
+    /** This function is the generic function to write both `if` and `else if` blocks.
+     *
+     * It is used for both @ref CLKernelWriter::op_if and @ref CLKernelWriter::op_else_if.
+     *
+     * @param[in] is_else True if this is an `else if` block, otherwise this is an `if` block.
+     * @param[in] lhs     The LHS tile of the condition.
+     * @param[in] op      The relational binary operator.
+     * @param[in] rhs     The RHS tile of the condition.
+     * @param[in] body    The function that writes the body of the else-if block.
+     */
+    void op_if_generic(bool is_else, const TileOperand &lhs, BinaryOp op, const TileOperand &rhs, const std::function<void()> &body);
+
     // For attributes
 private:
     /** This string contains the kernel body source code, not the full CL source code.