1 files changed, 19 insertions, 20 deletions
diff --git a/src/dynamic_fusion/sketch/gpu/ckw_driver/GpuCkwComponentArgument.h b/src/dynamic_fusion/sketch/gpu/ckw_driver/GpuCkwComponentArgument.h
index 863989a7bd..7a57c81e5f 100644
--- a/src/dynamic_fusion/sketch/gpu/ckw_driver/GpuCkwComponentArgument.h
+++ b/src/dynamic_fusion/sketch/gpu/ckw_driver/GpuCkwComponentArgument.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023 Arm Limited.
+ * Copyright (c) 2023-2024 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -25,13 +25,9 @@
 #ifndef ACL_SRC_DYNAMIC_FUSION_SKETCH_GPU_CKW_DRIVER_GPUCKWCOMPONENTARGUMENT_H
 #define ACL_SRC_DYNAMIC_FUSION_SKETCH_GPU_CKW_DRIVER_GPUCKWCOMPONENTARGUMENT_H
 
-#include "ckw/TensorTileSampler.h"
-
-namespace ckw
-{
-class TensorOperand;
-class TileOperand;
-} // namespace ckw
+#include "compute_kernel_writer/include/ckw/TensorOperand.h"
+#include "compute_kernel_writer/include/ckw/TensorSampler.h"
+#include "compute_kernel_writer/include/ckw/TileOperand.h"
 
 namespace arm_compute
 {
@@ -44,24 +40,27 @@ namespace dynamic_fusion
 class GpuCkwComponentArgument
 {
 public:
-    /** Initialize a new instance of @ref GpuCkwComponentArgument class for empty virtual tensor. */
-    GpuCkwComponentArgument();
+    /** Default constructor */
+    GpuCkwComponentArgument() = default;
 
     /** Initialize a new instance of @ref GpuCkwComponentArgument class for user tensor.
      *
      * @param[in] tensor The user tensor.
      */
-    explicit GpuCkwComponentArgument(ckw::TensorOperand &tensor);
+    explicit GpuCkwComponentArgument(ckw::TensorOperand tensor);
 
-    /** Set virtual tensor information (tile, sampler) for the argument.
+    /** Bind the tile and sampler to the tensor argument.
      *
-     * If the component is a user tensor, it can be treated as virtual tensor as well
-     * and won't be loaded again using @ref GpuCkwKernelWriter::op_load_once method.
+     * This method can be used to share a tile and sampler associated to a tensor
+     * among different kernel components. For example, when we create the destination
+     * tile and destination sampler for the first time (root component), this method can be
+     * used to bind these two information to the destination tensor so that the following
+     * simple components know the tile size and how to access the elements from memory.
      *
      * @param[in] tile    The tile that has been loaded.
      * @param[in] sampler The tensor sampling information that has been used to load the tile.
      */
-    GpuCkwComponentArgument &init_virtual_tensor(ckw::TileOperand &tile, const ckw::TensorTileSampler &sampler);
+    GpuCkwComponentArgument &init_virtual_tensor(ckw::TileOperand &tile, const ckw::TensorSampler &sampler);
 
     /** Get whether the argument is a user tensor. */
     bool has_tensor() const;
@@ -101,18 +100,18 @@ public:
      *
      * If the tile is not available, throw an error.
      */
-    ckw::TensorTileSampler &tile_sampler();
+    ckw::TensorSampler &tensor_sampler();
 
     /** Get the tensor sampling information for the tile.
      *
      * If the tile is not available, throw an error.
      */
-    const ckw::TensorTileSampler &tile_sampler() const;
+    const ckw::TensorSampler &tensor_sampler() const;
 
 private:
-    ckw::TensorOperand    *_tensor{nullptr};
-    ckw::TileOperand      *_tile{nullptr};
-    ckw::TensorTileSampler _tile_sampler{};
+    ckw::TensorOperand _tensor{};
+    ckw::TileOperand   _tile{};
+    ckw::TensorSampler _sampler{};
 };
 
 } // namespace dynamic_fusion