IQ.Pilot Release Commit @ 0798119

tinygrad_repo/extra/thunder/amd/include/ops/warp/shared/shared.cuh

@@ -0,0 +1,9 @@
+/**
+ * @file
+ * @brief An aggregate header of warp operations on data in shared memory
+ */
+
+#pragma once
+
+#include "tile/tile.cuh"
+#include "vec/vec.cuh"

tinygrad_repo/extra/thunder/amd/include/ops/warp/shared/tile/conversions.cuh

@@ -0,0 +1,37 @@
+/**
+ * @file
+ * @brief Conversions between shared tile types.
+ */
+
+#pragma once
+
+#include "../../../../common/common.cuh"
+#include "../../../../types/types.cuh"
+
+namespace kittens {
+
+/* ----------  SUBTILE  ---------- */
+
+/**
+* @brief Returns a reference to a subtile of the given shared tile.
+*
+* @tparam subtile_height The height of the subtile.
+* @tparam subtile_width The width of the subtile.
+* @tparam ST The type of the input tile, which must satisfy the ducks::st::all concept.
+* @param src The input tile.
+* @param row_idx The row coord of the subtile, in units of subtile_height*16 elements.
+* @param col_idx The col coord of the subtile, in units of subtile_width*16 elements.
+* @return A reference to the subtile.
+*
+* @note The subtile {height, width} must evenly divide the tile {height, width}.
+*/
+template<int subtile_rows, int subtile_cols, ducks::st::all ST>
+__device__ inline st_subtile<ST, subtile_rows, subtile_cols> subtile_inplace(ST &src, int2 rowcol) {
+    using T = typename ST::dtype;
+    static_assert(ST::rows % subtile_rows == 0);
+    static_assert(ST::cols % subtile_cols == 0);
+    static_assert(ST::rows == ST::underlying_rows && ST::cols == ST::underlying_cols); // must be a real ST, no recursive subtiles.
+    return st_subtile<ST, subtile_rows, subtile_cols>(src, rowcol);
+}
+
+} // namespace kittens

tinygrad_repo/extra/thunder/amd/include/ops/warp/shared/tile/tile.cuh

@@ -0,0 +1,8 @@
+/**
+ * @file
+ * @brief An aggregate header for warp operations on shared tiles.
+ */
+
+#pragma once
+
+#include "conversions.cuh"

tinygrad_repo/extra/thunder/amd/include/ops/warp/shared/vec/conversions.cuh

@@ -0,0 +1,55 @@
+/**
+ * @file
+ * @brief Warp-scope conversions on shared vectors.
+ */
+
+#pragma once
+
+#include "../../../../common/common.cuh"
+#include "../../../../types/types.cuh"
+
+
+namespace kittens {
+
+/**
+ * @brief Copies data from one shared vector to another, converting data types if necessary.
+ *
+ * This function copies data from the source shared vector `src` to the destination shared vector `dst`.
+ * If the data types of `src` and `dst` are the same, it performs a direct memory copy. Otherwise, it
+ * converts each element from the source data type to the destination data type using the appropriate
+ * converter before copying.
+ *
+ * @tparam SV1 The type of the destination shared vector, must satisfy the ducks::sv::all concept.
+ * @tparam SV2 The type of the source shared vector, must satisfy the ducks::sv::all concept.
+ * @param[out] dst The destination shared vector.
+ * @param[in] src The source shared vector.
+ * @note The lengths of `src` and `dst` must be equal. This is enforced at compile time.
+ */
+template<ducks::sv::all SV1, ducks::sv::all SV2>
+__device__ static inline void copy(SV1 &dst, const SV2 &src) {
+    static_assert(dst.length == src.length, "Source and destination vectors must have the same length.");
+    #pragma unroll
+    for(int i = kittens::laneid(); i < dst.length; i+=WARP_THREADS) {
+        dst[i] = base_types::convertor<typename SV1::dtype, typename SV2::dtype>::convert(src[i]);
+    }
+}
+
+/* ----------  SUBVEC  ---------- */
+
+/**
+* @brief Returns a reference to a subvec of a given shared vector
+*
+* @tparam subvec_length The length, in elements, of the subvec.
+* @tparam SV The type of the input vector, which must satisfy the ducks::sv::all concept.
+* @param src The input tile.
+* @param vec_idx The coord of the subvec, in units of subvec_length elements.
+* @return A reference to the subvec.
+*
+* @note The subvec length must evenly divide the vector length.
+*/
+template<int subvec_length, ducks::sv::all SV>
+__device__ inline typename SV::template subvec<subvec_length> &subvec_inplace(SV &src, int vec_idx) {
+    return *(typename SV::template subvec<subvec_length>*)(&src[vec_idx*subvec_length]);
+}
+
+}

tinygrad_repo/extra/thunder/amd/include/ops/warp/shared/vec/vec.cuh

@@ -0,0 +1,8 @@
+/**
+ * @file
+ * @brief An aggregate header for warp operations on data stored in shared memory.
+ */
+
+#pragma once
+
+#include "conversions.cuh"