(feat): Add a straightforward implementation for tile iterator. (#50)

* Add a straightforward implementation for tile iterator.
* Clean include relations.

Author: haruhi55

include/cell/copy/constants.hpp

@@ -0,0 +1,26 @@
+#pragma once
+
+namespace tiledcuda::cell::copy {
+
+enum class CopyInst {
+    LoadMat = 0,   // ldmatrix for loading data from shared memory to register.
+    StoreMat = 1,  // stmatrix for storing data from register to shared memory.
+    LoadS32 = 2,   // ldsm32 for loading 32-bit data from shared memory.
+    LoadS128 = 3   // ldsm128 for loading 128-bit data from shared memory.
+};
+
+enum class RegLayout {
+    TileWMMA = 0,  // Tile layout for TCU WMMA.
+};
+
+enum class WarpReuse {
+    NoReuse = 0,  // No reuse.
+    Cont = 1,     // Continuous
+    Cir = 2,      // Circular
+    RowCont = 3,  // RowWiseContinuouslyReuse
+    RowCir = 4,   // RowWiseCircularlyReuse
+    ColCont = 5,  // ColWiseContinuouslyReuse
+    ColCir = 6    // ColWiseCircularlyReuse
+};
+
+}  // namespace tiledcuda::cell::copy

include/cell/copy/copy.hpp

@@ -2,7 +2,7 @@
 
 #include "cuda_utils.hpp"
 
-#include <cute/algorithm/copy.hpp>
+#include <cute/tensor.hpp>
 
 namespace tiledcuda::cell::copy {
 

include/cell/copy/mod.hpp

@@ -1,5 +1,7 @@
 #pragma once
 
+#include "cell/copy/constants.hpp"
 #include "cell/copy/copy.hpp"
 #include "cell/copy/dyn_copy.hpp"
+#include "cell/copy/shared_to_register.hpp"
 #include "cell/copy/static_copy.hpp"

include/cell/copy/shared_to_register.hpp

@@ -0,0 +1,60 @@
+#pragma once
+
+#include "cell/copy/constants.hpp"
+#include "types/mod.hpp"
+
+namespace tiledcuda::cell::copy {
+
+namespace detail {
+
+// functor to copy data from shared memory to register file.
+template <typename Shared, typename Reg, typename WarpLayout,
+          CopyInst kCopyInst>
+struct CopyShared2Reg {
+    DEVICE void operator()(const Shared& src, Reg& dst, WarpReuse kMode);
+};
+
+// partial specialization for ldmatrix
+template <typename Shared, typename Reg, typename WarpLayout>
+struct CopyShared2Reg<Shared, Reg, WarpLayout, CopyInst::LoadMat> {
+    DEVICE void operator()(const Shared& src, Reg& dst, WarpReuse kMode) {
+        // implement this
+    }
+};
+
+// functor to copy data from shared memory to register file.
+template <typename Reg, typename Shared, typename InstShape,
+          RegLayout kRegLayout, CopyInst kCopyInst>
+struct CopyReg2Shared {
+    DEVICE void operator()(const Reg& src, Shared& dst);
+};
+
+// partial specialization for wmma 16x16x16, and LDSM32
+template <typename Reg, typename Shared>
+struct CopyReg2Shared<Reg, Shared, InstShape<16, 16, 16>, RegLayout::TileWMMA,
+                      CopyInst::LoadS32> {
+    DEVICE void operator()(const Reg& src, Shared& dst) {}
+};
+}  // namespace detail
+
+/// a warper function for the situation that `Shared` are computed from some
+/// runtime value.
+template <typename Shared, typename Reg, typename WarpLayout>
+DEVICE void copy_tile_s2r(const Shared& src, Reg& dst, const WarpLayout& layout,
+                          WarpReuse kMode) {
+    using Copy =
+        detail::CopyShared2Reg<Shared, Reg, WarpLayout, CopyInst::LoadMat>;
+
+    Copy copy;
+    copy(src, dst, kMode);
+}
+
+template <typename Reg, typename Shared>
+DEVICE void copy_tile_r2s(const Reg& src, Shared& dst) {
+    using Copy = detail::CopyReg2Shared<Reg, Shared, InstShape<16, 16, 16>,
+                                        RegLayout::TileWMMA, CopyInst::LoadS32>;
+    Copy copy;
+    copy(src, dst);
+}
+
+}  // namespace tiledcuda::cell::copy

include/cell/copy/static_copy.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "cell/copy/constants.hpp"
 #include "cuda_utils.hpp"
 #include "types/mod.hpp"
 
@@ -52,69 +53,4 @@ struct R2SCopy2D {
     }
 };
 
-namespace detail {
-
-enum class CopyInst {
-    Ldmatrix = 0,  // ldmatrix for loading data from shared memory to register.
-    Stmatrix = 1,
-    Ldsm32 = 2,
-    Ldsm128 = 3
-};
-
-enum class RegLayout {
-    TcuWmma = 0,  // tile layout for TCU WMMA.
-};
-
-// functor to copy data from shared memory to register file.
-template <typename Shared, typename Reg, CopyInst kCopyInst>
-struct CopyShared2Reg {
-    DEVICE void operator()();
-};
-
-// partial specialization for ldmatrix
-template <typename Shared, typename Reg>
-struct CopyShared2Reg<Shared, Reg, CopyInst::Ldmatrix> {
-    DEVICE void operator()(const Shared& src, Reg& dst) {}
-};
-
-// functor to copy data from shared memory to register file.
-template <typename Reg, typename Shared, typename InstShape,
-          RegLayout kRegLayout, CopyInst kCopyInst>
-struct CopyReg2Shared {
-    DEVICE void operator()();
-};
-
-// partial specialization for wmma 16x16x16, and LDSM32
-template <typename Reg, typename Shared>
-struct CopyReg2Shared<Reg, Shared, InstShape<16, 16, 16>, RegLayout::TcuWmma,
-                      CopyInst::Ldsm32> {
-    DEVICE void operator()(const Reg& src, Shared& dst) {}
-};
-
-}  // namespace detail
-
-/// @brief Copy a tile from shared memory: to register.
-/// @tparam Shared the shared memory tile type.
-/// @tparam Reg the register tile type.
-/// @tparam WarpLayout the warp layout.
-template <typename Shared, typename Reg, typename WarpLayout>
-DEVICE void copy_tile_s2r(const Shared& src, Reg& dst,
-                          const WarpLayout& layout /*for auto type-infer*/) {
-    using Copy =
-        detail::CopyShared2Reg<Shared, Reg, detail::CopyInst::Ldmatrix>;
-
-    Copy copy;
-    copy(src, dst);
-}
-
-template <typename Reg, typename Shared, typename WarpLayout>
-DEVICE void copy_tile_r2s(const Reg& src, Shared& dst,
-                          const WarpLayout& layout /*for auto type infer*/) {
-    using Copy = detail::CopyReg2Shared<Reg, Shared, InstShape<16, 16, 16>,
-                                        detail::RegLayout::TcuWmma,
-                                        detail::CopyInst::Ldsm32>;
-    Copy copy;
-    copy(src, dst);
-}
-
 }  // namespace tiledcuda::cell::copy

include/cell/traits/b2b_gemm.hpp

@@ -1,18 +1,13 @@
 #pragma once
 
-#include "cell/copy/static_copy.hpp"
+#include "cell/copy/mod.hpp"
 #include "cell/traits/base.hpp"
-#include "types/layout.hpp"
-#include "types/tile_shape.hpp"
-
-#include <cute/arch/copy.hpp>
-#include <cute/tensor.hpp>
-
-#include <type_traits>
+#include "types/mod.hpp"
 
 namespace tiledcuda::cell::traits {
 
 using namespace cute;
+
 namespace tl = tiledcuda::cell::tile_layout;
 
 template <typename Element_, typename CtaTileShape,

include/cell/traits/bmm.hpp

@@ -1,14 +1,9 @@
 #pragma once
 
-#include "cell/copy/static_copy.hpp"
+#include "cell/copy/mod.hpp"
 #include "cell/traits/base.hpp"
 #include "types/tile_shape.hpp"
 
-#include <cute/arch/copy.hpp>
-#include <cute/tensor.hpp>
-
-#include <type_traits>
-
 namespace tiledcuda::cell::traits {
 
 using namespace cute;

include/cell/traits/gemm.hpp

@@ -1,13 +1,8 @@
 #pragma once
 
-#include "cell/copy/static_copy.hpp"
+#include "cell/copy/mod.hpp"
 #include "cell/traits/base.hpp"
 
-#include <cute/arch/copy.hpp>
-#include <cute/tensor.hpp>
-
-#include <type_traits>
-
 namespace tiledcuda::cell::traits {
 
 using namespace cute;

include/cell/traits/lstm.hpp

@@ -1,14 +1,9 @@
 #pragma once
 
-#include "cell/copy/static_copy.hpp"
+#include "cell/copy/mod.hpp"
 #include "cell/traits/base.hpp"
 #include "types/tile_shape.hpp"
 
-#include <cute/arch/copy.hpp>
-#include <cute/tensor.hpp>
-
-#include <type_traits>
-
 namespace tiledcuda::cell::traits {
 
 using namespace cute;

include/types/layout.hpp

@@ -36,9 +36,29 @@ static constexpr size_t num_rows = cute::size<0>(Layout_{});
 template <typename Layout_>
 static constexpr size_t num_cols = cute::size<1>(Layout_{});
 
+template <typename Layout_>
+static constexpr size_t row_stride = cute::size<0>(Layout_{}.layout().stride());
+
+template <typename Layout_>
+static constexpr size_t col_stride = cute::size<1>(Layout_{}.layout().stride());
+
 template <typename Layout_>
 static constexpr size_t get_numel = int(size(Layout_{}));
 
+/// We wrap CuTe's `Layout`, which consists of `Shape` and `Stride`, into an
+/// intelligent row-major or column-major layout. In a row-major layout, the
+/// column stride is 1, whereas in a column-major layout, the row stride is 1.
+template <typename Layout_>
+static constexpr bool is_rowmajor = col_stride<Layout_> == 1;
+
+template <const int Shape1, const int Shape2, const int Stride1,
+          const int Stride2>
+HOST_DEVICE auto make_tile_layout() {
+    using Layout = cute::Layout<Shape<Int<Shape1>, Int<Shape2>>,
+                                Stride<Int<Stride1>, Int<Stride2>>>;
+    return Layout{};
+}
+
 HOST_DEVICE auto make_row_major_layout(const int row, const int col,
                                        const int stride) {
     return cute::make_layout(make_shape(row, col),

include/types/shared.hpp

@@ -15,11 +15,21 @@ class SharedTile {
     using Layout = Layout_;
 
     static constexpr int kNumel = tl::get_numel<Layout>;
+
     static constexpr int kRows = tl::num_rows<Layout>;
     static constexpr int kCols = tl::num_cols<Layout>;
 
+    static constexpr int kRowStride = tl::row_stride<Layout>;
+    static constexpr int kColStride = tl::col_stride<Layout>;
+
+    static constexpr bool kIsRowMajor = tl::is_rowmajor<Layout>;
+
     DEVICE SharedTile(DType* data) : data_(data), layout_(Layout{}) {}
 
+    DEVICE DType* mutable_data() { return data_; }
+
+    DEVICE const DType* data() const { return data_; }
+
     // for write access
     DEVICE DType& operator()(int x, int y) { return data_[layout_(x, y)]; }