From 40e4d4697bb3f848a87c11d8ae38cbc702050afb Mon Sep 17 00:00:00 2001
From: Chris Millette <chris.millette@amd.com>
Date: Thu, 22 Jan 2026 20:30:26 +0000
Subject: [PATCH 1/7] Refactor vector_type to reduce build time

---
 include/ck/utility/data_type.hpp              |   62 +-
 include/ck/utility/dtype_vector.hpp           | 2404 +++--------------
 include/ck/utility/math.hpp                   |   21 +-
 .../ck/utility/statically_indexed_array.hpp   |    1 +
 4 files changed, 445 insertions(+), 2043 deletions(-)

diff --git a/include/ck/utility/data_type.hpp b/include/ck/utility/data_type.hpp
index 8e6f875c399..f37437336ad 100644
--- a/include/ck/utility/data_type.hpp
+++ b/include/ck/utility/data_type.hpp
@@ -34,9 +34,17 @@ using f4_t    = unsigned _BitInt(4);
 using f6_t    = _BitInt(6);          // e2m3 format
 using bf6_t   = unsigned _BitInt(6); // e3m2 format
 
-// scalar_type
-template <typename TV>
-struct scalar_type;
+/**
+ * @brief Mapping of incoming type to local native storage type and vector size
+ * @tparam T Incoming data type
+ */
+template <typename T>
+struct scalar_type
+{
+    // Basic data type mapping to unsigned _BitInt of appropriate size
+    using type                           = unsigned _BitInt(8 * sizeof(T));
+    static constexpr index_t vector_size = 1;
+};
 
 struct f4x2_pk_t
 {
@@ -191,12 +199,6 @@ struct pk_i4_t
     __host__ __device__ constexpr pk_i4_t(type init) : data{init} {}
 };
 
-inline constexpr auto next_pow2(uint32_t x)
-{
-    // Precondition: x > 1.
-    return x > 1u ? (1u << (32u - __builtin_clz(x - 1u))) : x;
-}
-
 // native types: double, float, _Float16, ushort, int32_t, int8_t, uint8_t, f8_fnuz_t, bf8_fnuz_t,
 // native types: bool
 template <typename T>
@@ -208,10 +210,6 @@ inline constexpr bool is_native_type()
            is_same_v<T, _BitInt(8)> || is_same_v<T, unsigned _BitInt(8)> || is_same<T, bool>::value;
 }
 
-// scalar_type
-template <typename TV>
-struct scalar_type;
-
 // is_scalar_type
 template <typename TV>
 struct is_scalar_type
@@ -224,14 +222,13 @@ template <typename X, typename Y>
 using has_same_scalar_type = is_same<typename scalar_type<remove_cvref_t<X>>::type,
                                      typename scalar_type<remove_cvref_t<Y>>::type>;
 
-template <typename T, index_t N>
-struct scalar_type<T __attribute__((ext_vector_type(N)))>
+template <>
+struct scalar_type<bool>
 {
-    using type                           = T;
-    static constexpr index_t vector_size = N;
+    using type                           = bool;
+    static constexpr index_t vector_size = 1;
 };
 
-//
 template <>
 struct scalar_type<double>
 {
@@ -293,35 +290,35 @@ struct scalar_type<int4_t>
 template <>
 struct scalar_type<pk_i4_t>
 {
-    using type                           = pk_i4_t;
+    using type                           = typename pk_i4_t::type;
     static constexpr index_t vector_size = 1;
 };
 
 template <>
 struct scalar_type<f8_fnuz_t>
 {
-    using type                           = f8_fnuz_t::data_type;
+    using type                           = typename f8_fnuz_t::data_type;
     static constexpr index_t vector_size = 1;
 };
 
 template <>
 struct scalar_type<bf8_fnuz_t>
 {
-    using type                           = bf8_fnuz_t::data_type;
+    using type                           = typename bf8_fnuz_t::data_type;
     static constexpr index_t vector_size = 1;
 };
 
 template <>
 struct scalar_type<f8_ocp_t>
 {
-    using type                           = f8_ocp_t::data_type;
+    using type                           = typename f8_ocp_t::data_type;
     static constexpr index_t vector_size = 1;
 };
 
 template <>
 struct scalar_type<bf8_ocp_t>
 {
-    using type                           = bf8_ocp_t::data_type;
+    using type                           = typename bf8_ocp_t::data_type;
     static constexpr index_t vector_size = 1;
 };
 
@@ -329,7 +326,7 @@ struct scalar_type<bf8_ocp_t>
 template <>
 struct scalar_type<e8m0_bexp_t>
 {
-    using type                           = e8m0_bexp_t::type;
+    using type                           = typename e8m0_bexp_t::type;
     static constexpr index_t vector_size = 1;
 };
 #endif
@@ -337,42 +334,35 @@ struct scalar_type<e8m0_bexp_t>
 template <>
 struct scalar_type<f4x2_pk_t>
 {
-    using type                           = f4x2_pk_t::type;
+    using type                           = typename f4x2_pk_t::type;
     static constexpr index_t vector_size = 1;
 };
 
 template <>
 struct scalar_type<f6x32_pk_t>
 {
-    using type                           = f6x32_pk_t::storage_type;
+    using type                           = typename f6x32_pk_t::storage_type;
     static constexpr index_t vector_size = 1;
 };
 
 template <>
 struct scalar_type<bf6x32_pk_t>
 {
-    using type                           = bf6x32_pk_t::storage_type;
+    using type                           = typename bf6x32_pk_t::storage_type;
     static constexpr index_t vector_size = 1;
 };
 
 template <>
 struct scalar_type<f6x16_pk_t>
 {
-    using type                           = f6x16_pk_t::storage_type;
+    using type                           = typename f6x16_pk_t::storage_type;
     static constexpr index_t vector_size = 1;
 };
 
 template <>
 struct scalar_type<bf6x16_pk_t>
 {
-    using type                           = bf6x16_pk_t::storage_type;
-    static constexpr index_t vector_size = 1;
-};
-
-template <>
-struct scalar_type<bool>
-{
-    using type                           = bool;
+    using type                           = typename bf6x16_pk_t::storage_type;
     static constexpr index_t vector_size = 1;
 };
 
diff --git a/include/ck/utility/dtype_vector.hpp b/include/ck/utility/dtype_vector.hpp
index ebdbbb107d7..85f50b11d30 100644
--- a/include/ck/utility/dtype_vector.hpp
+++ b/include/ck/utility/dtype_vector.hpp
@@ -3,197 +3,103 @@
 
 #pragma once
 #include "ck/utility/data_type.hpp"
+#include "ck/utility/math.hpp"
 
 namespace ck {
 
-// vector_type
-template <typename T, index_t N, typename Enable = void>
-struct vector_type;
-
-// Caution: DO NOT REMOVE
-// intentionally have only declaration but no definition to cause compilation failure when trying to
-// instantiate this template. The purpose is to catch user's mistake when trying to make "vector of
-// vectors"
-template <typename T, index_t V, index_t N>
-struct vector_type<T __attribute__((ext_vector_type(V))), N>;
-
-// Caution: DO NOT REMOVE
-// intentionally have only declaration but no definition to cause compilation failure when trying to
-// instantiate this template. The purpose is to catch user's mistake when trying to make "vector of
-// vectors"
-template <typename T, index_t V, index_t N>
-struct vector_type<vector_type<T, V>, N>;
-
-// vector_type_maker
-// This is the right way to handle "vector of vectors": making a bigger vector instead
-template <typename T, index_t N>
-struct vector_type_maker
-{
-    using type = vector_type<T, N>;
-};
-
-template <typename T, index_t N>
-struct scalar_type<vector_type<T, N>>
+/**
+ * @brief Wrapper for native vector type
+ * @tparam T The element type of the vector
+ * @tparam Rank The number of elements in the vector
+ */
+template <typename T, index_t Rank>
+using NativeVectorT = T __attribute__((ext_vector_type(Rank)));
+
+/**
+ * @brief scalar_type trait override for NativeVectorT
+ * @tparam T The vector type
+ * @tparam Rank The number of elements in the vector
+ */
+template <typename T, index_t Rank>
+struct scalar_type<NativeVectorT<T, Rank>>
 {
     using type                           = T;
-    static constexpr index_t vector_size = N;
-};
-
-template <typename T, index_t N0, index_t N1>
-struct vector_type_maker<T __attribute__((ext_vector_type(N1))), N0>
-{
-    using type = vector_type<T, N0 * N1>;
+    static constexpr index_t vector_size = Rank;
 };
 
-template <typename T, index_t N0, index_t N1>
-struct vector_type_maker<vector_type<T, N1>, N0>
-{
-    using type = vector_type<T, N0 * N1>;
-};
+__device__ int static err = 0;
 
-template <typename T, index_t N>
-using vector_type_maker_t = typename vector_type_maker<T, N>::type;
+template <typename T, index_t N, typename Enable = void>
+struct non_native_vector_base;
 
 template <typename T, index_t N>
-__host__ __device__ constexpr auto make_vector_type(Number<N>)
-{
-    return typename vector_type_maker<T, N>::type{};
-}
-
-template <typename T>
-struct vector_type<T, 1, typename ck::enable_if_t<is_native_type<T>()>>
+struct non_native_vector_base<
+    T,
+    N,
+    ck::enable_if_t<sizeof(T) == 1 || sizeof(T) == 2 || sizeof(T) == 4 || sizeof(T) == 8>>
 {
-    using d1_t = T;
-    using type = d1_t;
+    using data_t = typename scalar_type<T>::type; // select data_t based on the size of T
+    static_assert(sizeof(T) == sizeof(data_t), "non_native_vector_base storage size mismatch");
+    using data_v = data_t __attribute__((ext_vector_type(N)));
+    using type   = non_native_vector_base<T, N>;
 
-    union
+    union alignas(math::next_power_of_two<N * sizeof(T)>())
     {
-        T d1_;
-        StaticallyIndexedArray<T, 1> d1x1_;
+        data_v dN; // storage vector;
+        StaticallyIndexedArray_v2<data_t, N> dxN;
+        StaticallyIndexedArray_v2<T, N> dTxN;
+        StaticallyIndexedArray_v2<data_v, 1> dNx1;
     } data_;
 
-    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        return data_.d1x1_;
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
+    __host__ __device__ constexpr non_native_vector_base(data_t a) : data_{data_v(a)} {}
+    __host__ __device__ constexpr non_native_vector_base(T f)
+        : non_native_vector_base(bit_cast<data_t>(f))
     {
-        static_assert(is_same<X, d1_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        return data_.d1x1_;
     }
-};
-
-__device__ int static err = 0;
-template <typename T>
-struct vector_type<T, 2, typename ck::enable_if_t<is_native_type<T>()>>
-{
-    using d1_t = T;
-    typedef T d2_t __attribute__((ext_vector_type(2)));
-
-    using type = d2_t;
-
-    union
-    {
-        d2_t d2_;
-        StaticallyIndexedArray<d1_t, 2> d1x2_;
-        StaticallyIndexedArray<d2_t, 1> d2x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
+    __host__ __device__ constexpr non_native_vector_base() : non_native_vector_base(T{}){};
+    __host__ __device__ constexpr non_native_vector_base(data_v v) : data_{v} {}
 
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
+    __host__ __device__ constexpr operator data_v() const { return data_.dN; }
+    __host__ __device__ constexpr operator data_t() const
     {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x2_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
+        if constexpr(N == 1)
         {
-            return data_.d2x1_;
+            return data_.dxN[Number<0>{}];
         }
         else
         {
-            return err;
+            return data_.dxN; // XXX this should cause an error
         }
     }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
+    __host__ __device__ constexpr operator T() const
     {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x2_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
+        if constexpr(N == 1)
         {
-            return data_.d2x1_;
+            return data_.dTxN[Number<0>{}];
         }
         else
         {
-            return err;
+            return data_.dTxN; // XXX this should cause an error
         }
     }
-};
-
-template <typename T>
-struct vector_type<T, 3, typename ck::enable_if_t<is_native_type<T>()>>
-{
-    using d1_t = T;
-    typedef T d2_t __attribute__((ext_vector_type(2)));
-    typedef T d3_t __attribute__((ext_vector_type(3)));
-
-    using type = d3_t;
-
-    union
-    {
-        d3_t d3_;
-        StaticallyIndexedArray<d1_t, 3> d1x3_;
-        StaticallyIndexedArray<d2_t, 1> d2x1_;
-        StaticallyIndexedArray<d3_t, 1> d3x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
 
     template <typename X>
     __host__ __device__ constexpr const auto& AsType() const
     {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value || is_same<X, d3_t>::value,
+        static_assert(is_same_v<X, data_t> || is_same_v<X, T> || is_same_v<X, data_v>,
                       "Something went wrong, please check src and dst types.");
 
-        if constexpr(is_same<X, d1_t>::value)
+        if constexpr(is_same_v<X, data_t>)
         {
-            return data_.d1x3_;
+            return data_.dxN;
         }
-        else if constexpr(is_same<X, d2_t>::value)
+        else if constexpr(is_same_v<X, T>)
         {
-            return data_.d2x1_;
+            return data_.dTxN;
         }
-        else if constexpr(is_same<X, d3_t>::value)
+        else if constexpr(is_same_v<X, data_v>)
         {
-            return data_.d3x1_;
+            return data_.dNx1;
         }
         else
         {
@@ -204,20 +110,20 @@ struct vector_type<T, 3, typename ck::enable_if_t<is_native_type<T>()>>
     template <typename X>
     __host__ __device__ constexpr auto& AsType()
     {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value || is_same<X, d3_t>::value,
+        static_assert(is_same_v<X, data_t> || is_same_v<X, T> || is_same_v<X, data_v>,
                       "Something went wrong, please check src and dst types.");
 
-        if constexpr(is_same<X, d1_t>::value)
+        if constexpr(is_same_v<X, data_t>)
         {
-            return data_.d1x3_;
+            return data_.dxN;
         }
-        else if constexpr(is_same<X, d2_t>::value)
+        else if constexpr(is_same_v<X, T>)
         {
-            return data_.d2x1_;
+            return data_.dTxN;
         }
-        else if constexpr(is_same<X, d3_t>::value)
+        else if constexpr(is_same_v<X, data_v>)
         {
-            return data_.d3x1_;
+            return data_.dNx1;
         }
         else
         {
@@ -226,138 +132,80 @@ struct vector_type<T, 3, typename ck::enable_if_t<is_native_type<T>()>>
     }
 };
 
-template <typename T>
-struct vector_type<T, 4, typename ck::enable_if_t<is_native_type<T>()>>
+// implementation for f6x16 and f6x32
+template <typename T, index_t N>
+struct non_native_vector_base<
+    T,
+    N,
+    ck::enable_if_t<sizeof(T) == 12 || sizeof(T) == 16 || sizeof(T) == 24 || sizeof(T) == 32>>
 {
-    using d1_t = T;
-    typedef T d2_t __attribute__((ext_vector_type(2)));
-    typedef T d4_t __attribute__((ext_vector_type(4)));
-
-    using type = d4_t;
+    using data_t    = typename scalar_type<T>::type; // select data_t based on declared base type
+    using element_t = typename T::element_type; // select element_t based on declared element type
+    static_assert(sizeof(T) == sizeof(data_t), "non_native_vector_base storage size mismatch");
+    static constexpr size_t size_factor = sizeof(data_t) / sizeof(element_t);
+    using data_v = element_t __attribute__((ext_vector_type(N * size_factor)));
+    using type   = non_native_vector_base<T, N>;
 
-    union
+    union alignas(math::next_power_of_two<N * sizeof(T)>())
     {
-        d4_t d4_;
-        StaticallyIndexedArray<d1_t, 4> d1x4_;
-        StaticallyIndexedArray<d2_t, 2> d2x2_;
-        StaticallyIndexedArray<d4_t, 1> d4x1_;
+        data_v dN; // storage vector;
+        StaticallyIndexedArray_v2<data_t, N> dxN;
+        StaticallyIndexedArray_v2<T, N> dTxN;
+        StaticallyIndexedArray_v2<data_v, 1> dNx1;
     } data_;
 
-    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
+    // Broadcast single value to vector
+    __host__ __device__ constexpr non_native_vector_base(data_t a) : data_{}
     {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value || is_same<X, d4_t>::value,
-                      "Something went wrong, please check src and dst types.");
+        // TODO: consider removing initialization similar to vector_type<T, 256>
 
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x4_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x2_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x1_;
-        }
-        else
-        {
-            return err;
-        }
+        ck::static_for<0, N, 1>{}([&](auto i) {
+            data_.dxN(i) = a; // broadcast value to all elements
+        });
     }
 
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
+    __host__ __device__ constexpr non_native_vector_base(T f)
+        : non_native_vector_base(bit_cast<data_t>(f))
     {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value || is_same<X, d4_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x4_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x2_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x1_;
-        }
-        else
-        {
-            return err;
-        }
     }
-};
 
-template <typename T>
-struct vector_type<T, 5, typename ck::enable_if_t<is_native_type<T>()>>
-{
-    using d1_t = T;
-    typedef T d4_t __attribute__((ext_vector_type(4)));
-    typedef T d5_t __attribute__((ext_vector_type(5)));
-
-    using type = d5_t;
+    __host__ __device__ constexpr non_native_vector_base() : non_native_vector_base(T{}){};
 
-    union
-    {
-        d5_t d5_;
-        StaticallyIndexedArray<d1_t, 5> d1x5_;
-        StaticallyIndexedArray<d4_t, 1> d4x1_;
-        StaticallyIndexedArray<d5_t, 1> d5x1_;
-    } data_;
+    __host__ __device__ constexpr non_native_vector_base(data_v v) : data_{v} {}
 
-    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
+    __host__ __device__ constexpr non_native_vector_base(element_t v) : data_{data_v(v)} {}
 
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
+    __host__ __device__ constexpr operator data_v() const { return data_.dN; }
 
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
+    __host__ __device__ constexpr operator T() const
     {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d4_t>::value || is_same<X, d5_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x5_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x1_;
-        }
-        else if constexpr(is_same<X, d5_t>::value)
+        if constexpr(N == 1)
         {
-            return data_.d5x1_;
+            return data_.dTxN[Number<0>{}];
         }
         else
         {
-            return err;
+            return err; // XXX this should cause an error
         }
     }
 
     template <typename X>
-    __host__ __device__ constexpr auto& AsType()
+    __host__ __device__ constexpr const auto& AsType() const
     {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d4_t>::value || is_same<X, d5_t>::value,
+        static_assert(is_same_v<X, data_t> || is_same_v<X, data_v> || is_same_v<X, T>,
                       "Something went wrong, please check src and dst types.");
 
-        if constexpr(is_same<X, d1_t>::value)
+        if constexpr(is_same_v<X, data_v>)
         {
-            return data_.d1x5_;
+            return data_.dNx1;
         }
-        else if constexpr(is_same<X, d4_t>::value)
+        else if constexpr(is_same_v<X, data_t>)
         {
-            return data_.d4x1_;
+            return data_.dxN;
         }
-        else if constexpr(is_same<X, d5_t>::value)
+        else if constexpr(is_same_v<X, T>)
         {
-            return data_.d5x1_;
+            return data_.dTxN;
         }
         else
         {
@@ -366,1781 +214,333 @@ struct vector_type<T, 5, typename ck::enable_if_t<is_native_type<T>()>>
     }
 };
 
-template <typename T>
-struct vector_type<T, 6, typename ck::enable_if_t<is_native_type<T>()>>
+template <typename T, index_t N>
+struct scalar_type<non_native_vector_base<
+    T,
+    N,
+    ck::enable_if_t<sizeof(T) == 1 || sizeof(T) == 2 || sizeof(T) == 4 || sizeof(T) == 8>>>
 {
-    using d1_t = T;
-    typedef T d2_t __attribute__((ext_vector_type(2)));
-    typedef T d3_t __attribute__((ext_vector_type(3)));
-    typedef T d6_t __attribute__((ext_vector_type(6)));
+    using type                           = typename non_native_vector_base<T, N>::data_t;
+    static constexpr index_t vector_size = N;
+};
 
-    using type = d6_t;
+template <typename T, index_t N>
+struct scalar_type<non_native_vector_base<
+    T,
+    N,
+    ck::enable_if_t<sizeof(T) == 12 || sizeof(T) == 16 || sizeof(T) == 24 || sizeof(T) == 32>>>
+{
+    using type                           = typename non_native_vector_base<T, N>::element_t;
+    static constexpr index_t vector_size = N * non_native_vector_base<T, N>::size_factor;
+};
 
-    union
+/**
+ * @brief Helper struct to determine the storage type for vector_type
+ * @tparam T The element type of the vector
+ * @tparam Rank The number of elements in the vector
+ * @tparam Enable SFINAE helper
+ */
+template <typename T, index_t Rank, typename Enable = void>
+struct vector_type_storage;
+
+/**
+ * @brief Vector storage type for native scalar types.
+ * @tparam T The element type of the vector
+ * @note For Rank = 1 and native types, the storage type is simply T itself (scalar)
+ */
+template <typename T>
+struct vector_type_storage<T, 1, enable_if_t<is_native_type<T>()>>
+{
+    using type = T;
+};
+
+/**
+ * @brief Vector storage type for native vector types.
+ * @tparam T The element type of the vector
+ * @tparam Rank The number of elements in the vector
+ *
+ * Assigns a native vector type based on the element type and rank.
+ * For boolean types, uses a C-style array `T[Rank]`, otherwise uses
+ * the `NativeVectorT<T, Rank>` template specialization.
+ *
+ * @note Special handling note:
+ * Sub-byte sizes such as bool have different sizes in ext_vector_type (via NativeVectorT) vs array
+ * types due to packing. Builtin vector types pack bool elements, while C++ arrays use 1 byte per
+ * bool as a standard (minimum write size = 1 byte). e.g., ext_vector_type(bool, 4) is packed as
+ * minimum 1 byte, while bool[4] is 4 bytes. vector_type::AsType, aliases with
+ * StaticallyIndexedArray_v2 which is C-style array under the hood, so we need to avoid using
+ * ext_vector_type with bool due to potential for data slicing errors.
+ */
+template <typename T, index_t Rank>
+struct vector_type_storage<T, Rank, enable_if_t<is_native_type<T>() && (Rank > 1)>>
+{
+    using type = std::conditional_t<std::is_same_v<bool, T>, T[Rank], NativeVectorT<T, Rank>>;
+    ;
+};
+
+/**
+ * @brief Vector storage type for non-native vector types.
+ * @tparam T The element type of the vector
+ * @tparam Rank The number of elements in the vector
+ * @note For non-native types, the storage type is non_native_vector_base<T, Rank>
+ */
+template <typename T, index_t Rank>
+struct vector_type_storage<T, Rank, enable_if_t<!is_native_type<T>()>>
+{
+    using type = non_native_vector_base<T, Rank>;
+};
+
+/**
+ * @brief Convenience wrapper for vector_type_storage
+ * @tparam T The element type of the vector
+ * @tparam Rank The number of elements in the vector
+ */
+template <typename T, index_t Rank>
+using vector_type_storage_t = typename vector_type_storage<T, Rank>::type;
+
+/**
+ * @brief Trait to check whether one storage class is the same as another (e.g., same scalar, or
+ * same vector class).
+ * @tparam Lhs The source storage type
+ * @tparam Rhs The comparator storage type
+ *
+ * Same storage classes are:
+ * - Same type
+ * - Same template vector types with matching base type (may have different ranks)
+ * - C-style arrays of same base type (may have different ranks)
+ */
+template <typename Lhs, typename Rhs>
+struct is_same_storage_class : public false_type
+{
+};
+
+/**
+ * @brief Same type storage class
+ * @tparam T The storage type
+ */
+template <typename T>
+struct is_same_storage_class<T, T> : public true_type
+{
+};
+
+/**
+ * @brief Template vector types of same base type with different ranks
+ * @tparam VecT The <T, N> vector template class type (e.g., vector_type, NativeVectorT,
+ * non_native_vector_base)
+ * @tparam T The base element type
+ * @tparam LhsRank The rank of the source type
+ * @tparam RhsRank The rank of the comparator type
+ */
+template <template <typename, index_t> class VecT, typename T, index_t LhsRank, index_t RhsRank>
+struct is_same_storage_class<VecT<T, LhsRank>, VecT<T, RhsRank>> : true_type
+{
+};
+
+/**
+ * @brief C-style arrays of same base type with different ranks
+ * @tparam T The base element type
+ * @tparam LhsRank The rank of the source type
+ * @tparam RhsRank The rank of the comparator type
+ */
+template <typename T, index_t LhsRank, index_t RhsRank>
+struct is_same_storage_class<T[LhsRank], T[RhsRank]> : true_type
+{
+};
+
+/**
+ * @brief Convenience evaluator for is_same_storage_class
+ * @tparam Lhs The source storage type
+ * @tparam Rhs The comparator storage type
+ */
+template <typename Lhs, typename Rhs>
+static constexpr bool is_same_storage_class_v = is_same_storage_class<Lhs, Rhs>::value;
+
+/**
+ * @brief Vector type wrapper
+ * @tparam T The element type of the vector
+ * @tparam Rank The number of elements in the vector
+ */
+template <typename T, index_t Rank>
+struct vector_type
+{
+    /// @brief Internal storage type for vector_type.
+    using StorageT = vector_type_storage_t<T, Rank>;
+    using type     = StorageT;
+    StorageT data_;
+
+    /// @brief Default constructor for vector_type
+    __host__ __device__ constexpr vector_type() : data_{} {}
+
+    /// @brief Constructor for native vector initialization
+    __host__ __device__ constexpr vector_type(StorageT v) : data_{v} {}
+
+    /**
+     * @brief Validates whether a type can be used in an AsType cast operation for vector_type
+     * class.
+     *
+     * This function checks if a given type X can be legally used as an alias to either reinterpret
+     * or slice (iterate) through the local storage type StorageT. The validation ensures type
+     * safety and structural compatibility between the source and target vector types.
+     *
+     * @tparam X The target type to validate for AsType casting.
+     *
+     * @return constexpr bool True if the type is valid for AsType casting, false otherwise
+     *
+     * @note Requirements for a valid AsType<X> cast on vector_type<T, Rank>:
+     *       1. The value type of X must match the storage value type (T)
+     *       2. X must be either:
+     *          a) A scalar type (T) where RankX == 1, OR
+     *          b) A vector class that matches the storage vector class (e.g., both are
+     * NativeVectorT or non_native_vector_base) where:
+     *             - RankX is a power of 2, OR
+     *               RankX == 3, OR
+     *               RankX == Storage Rank
+     *             - RankX must be <= Storage Rank
+     * @example
+     * auto srcVec = vector_type<float, 8>{};  // T = float, Rank = 8, native vector storage
+     * auto result = srcVec.AsType<X>();       // Where datatype X could be:
+     * X = NativeVectorT<float, 4>;            // OK: native vector T, RankX = 4 (power of 2)
+     * X = float;                              // OK: scalar T, RankX = 1
+     * X = NativeVectorT<float, 5>;            // ERROR: RankX not a power of 2, ==3, or ==Rank
+     * X = int;                                // ERROR: Invalid scalar cast, T != int
+     * X = float[4];                           // ERROR: Invalid type, storage vector class doesn't
+     * match
+     *                                         //        native vector vs C-array
+     */
+    template <typename X>
+    static constexpr bool is_as_type_cast_valid()
+    {
+        using TraitsX = scalar_type<X>;
+
+        // Checks storage classes match, with same base type (may have different ranks)
+        constexpr bool is_valid_cast = is_same_storage_class_v<StorageT, X>;
+
+        // Validate vector ranks
+        constexpr bool is_valid_rank =
+            (math::is_power_of_two_integer(TraitsX::vector_size) || (TraitsX::vector_size == 3) ||
+             (TraitsX::vector_size == Rank)) &&
+            (TraitsX::vector_size <= Rank);
+
+        return is_valid_cast || is_valid_rank;
+    }
+
+    /**
+     * @brief Allows casting the vector_type to another type X via aliasing or slicing.
+     * Use cases are to expose the internal storage type, or to slice the vector into smaller
+     * vectors for iteration purposes.
+     * @tparam X The target type to validate for AsType casting.
+     * @returns a reference to the reinterpreted data as StaticallyIndexedArray_v2<X, Rank / RankX>.
+     * Rigid control of allowable casts is enforced via static_assert to ensure type safety.
+     * See is_as_type_cast_valid() for requirements.
+     */
+    template <typename X>
+    __host__ __device__ constexpr auto const& AsType() const
+    {
+        // Make this a hard error if the datatype X is not a valid cast.
+        static_assert(is_as_type_cast_valid<X>(), "Datatype X is not a valid AsType cast");
+
+        using TraitsX = scalar_type<X>;
+
+        // Calculate the new rank after slicing.
+        // Note: We might end up with incomplete quantization from slicing
+        // when Rank % TraitsX::vector_size != 0, so take the floor division.
+        constexpr index_t newRank = Rank / TraitsX::vector_size;
+        using ResultT             = StaticallyIndexedArray_v2<X, newRank>;
+
+        // As a rule, the aliasing type should not be larger than the original type.
+        static_assert(sizeof(ResultT) <= sizeof(vector_type),
+                      "Resulting aliasing cannot be larger than original type");
+
+        // Re-cast as vectorized type.
+        return *(bit_cast<ResultT const*>(this));
+    }
+
+    /**
+     * @brief Allows casting the vector_type to another type X via aliasing or slicing.
+     * Use cases are to expose the internal storage type, or to slice the vector into smaller
+     * vectors for iteration purposes.
+     * @tparam X The target type to validate for AsType casting.
+     * @returns a reference to the reinterpreted data as StaticallyIndexedArray_v2<X, Rank / RankX>.
+     * Rigid control of allowable casts is enforced via static_assert to ensure type safety.
+     * See is_as_type_cast_valid() for requirements.
+     */
+    template <typename X>
+    __host__ __device__ constexpr auto& AsType()
     {
-        d6_t d6_;
-        StaticallyIndexedArray<d1_t, 6> d1x6_;
-        StaticallyIndexedArray<d2_t, 3> d2x3_;
-        StaticallyIndexedArray<d3_t, 2> d3x2_;
-        StaticallyIndexedArray<d6_t, 1> d6x1_;
-    } data_;
+        // Make this a hard error if the datatype X is not a valid cast.
+        static_assert(is_as_type_cast_valid<X>(), "Datatype X is not a valid AsType cast");
 
-    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
+        using TraitsX = scalar_type<X>;
 
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
+        // Calculate the new rank after slicing.
+        // Note: We might end up with incomplete quantization from slicing
+        // when Rank % TraitsX::vector_size != 0, so take the floor division.
+        constexpr index_t newRank = Rank / TraitsX::vector_size;
+        using ResultT             = StaticallyIndexedArray_v2<X, newRank>;
 
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d3_t>::value || is_same<X, d6_t>::value,
-                      "Something went wrong, please check src and dst types.");
+        // As a rule, the aliasing type should not be larger than the original type.
+        static_assert(sizeof(ResultT) <= sizeof(vector_type),
+                      "Resulting aliasing cannot be larger than original type");
 
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x6_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x3_;
-        }
-        else if constexpr(is_same<X, d3_t>::value)
-        {
-            return data_.d3x2_;
-        }
-        else if constexpr(is_same<X, d6_t>::value)
-        {
-            return data_.d6x1_;
-        }
-        else
-        {
-            return err;
-        }
+        // Re-cast as vectorized type.
+        return *(bit_cast<ResultT*>(this));
     }
+};
 
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d3_t>::value || is_same<X, d6_t>::value,
-                      "Something went wrong, please check src and dst types.");
+// Caution: DO NOT REMOVE
+// intentionally have only declaration but no definition to cause compilation failure when trying to
+// instantiate this template. The purpose is to catch user's mistake when trying to make "vector of
+// vectors"
+template <typename T, index_t V, index_t N>
+struct vector_type<NativeVectorT<T, V>, N>;
 
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x6_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x3_;
-        }
-        else if constexpr(is_same<X, d3_t>::value)
-        {
-            return data_.d3x2_;
-        }
-        else if constexpr(is_same<X, d6_t>::value)
-        {
-            return data_.d6x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
+// Caution: DO NOT REMOVE
+// intentionally have only declaration but no definition to cause compilation failure when trying to
+// instantiate this template. The purpose is to catch user's mistake when trying to make "vector of
+// vectors"
+template <typename T, index_t V, index_t N>
+struct vector_type<vector_type<T, V>, N>;
 
-template <typename T>
-struct vector_type<T, 7, typename ck::enable_if_t<is_native_type<T>()>>
+/**
+ * @brief scalar_type trait override for vector_type
+ * @tparam T The vector type
+ * @tparam N The number of elements in the vector
+ */
+template <typename T, index_t N>
+struct scalar_type<vector_type<T, N>>
 {
-    using d1_t = T;
-    typedef T d2_t __attribute__((ext_vector_type(2)));
-    typedef T d4_t __attribute__((ext_vector_type(4)));
-    typedef T d7_t __attribute__((ext_vector_type(7)));
-
-    using type = d7_t;
-
-    union
-    {
-        d7_t d7_;
-        StaticallyIndexedArray<d1_t, 7> d1x7_;
-        StaticallyIndexedArray<d2_t, 3> d2x3_;
-        StaticallyIndexedArray<d4_t, 1> d4x1_;
-        StaticallyIndexedArray<d7_t, 1> d7x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d7_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x7_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x3_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x1_;
-        }
-        else if constexpr(is_same<X, d7_t>::value)
-        {
-            return data_.d7x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d7_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x7_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x3_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x1_;
-        }
-        else if constexpr(is_same<X, d7_t>::value)
-        {
-            return data_.d7x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
+    using type                           = typename scalar_type<T>::type;
+    static constexpr index_t vector_size = N;
 };
 
-template <typename T>
-struct vector_type<T, 8, typename ck::enable_if_t<is_native_type<T>()>>
+// vector_type_maker
+// This is the right way to handle "vector of vectors": making a bigger vector instead
+template <typename T, index_t N>
+struct vector_type_maker
 {
-    using d1_t = T;
-    typedef T d2_t __attribute__((ext_vector_type(2)));
-    typedef T d4_t __attribute__((ext_vector_type(4)));
-    typedef T d8_t __attribute__((ext_vector_type(8)));
-
-    using type = d8_t;
-
-    union
-    {
-        d8_t d8_;
-        StaticallyIndexedArray<d1_t, 8> d1x8_;
-        StaticallyIndexedArray<d2_t, 4> d2x4_;
-        StaticallyIndexedArray<d4_t, 2> d4x2_;
-        StaticallyIndexedArray<d8_t, 1> d8x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x8_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x4_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x2_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x8_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x4_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x2_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
+    using type = vector_type<T, N>;
 };
 
-template <typename T>
-struct vector_type<T, 13, typename ck::enable_if_t<is_native_type<T>()>>
+template <typename T, index_t N0, index_t N1>
+struct vector_type_maker<NativeVectorT<T, N1>, N0>
 {
-    using d1_t = T;
-    typedef T d4_t __attribute__((ext_vector_type(4)));
-    typedef T d8_t __attribute__((ext_vector_type(8)));
-    typedef T d13_t __attribute__((ext_vector_type(13)));
-
-    using type = d13_t;
-
-    union
-    {
-        d13_t d13_;
-        StaticallyIndexedArray<d1_t, 13> d1x13_;
-        StaticallyIndexedArray<d4_t, 3> d4x3_;
-        StaticallyIndexedArray<d8_t, 1> d8x1_;
-        StaticallyIndexedArray<d13_t, 1> d13x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d4_t>::value ||
-                          is_same<X, d8_t>::value || is_same<X, d13_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x13_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x3_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x1_;
-        }
-        else if constexpr(is_same<X, d13_t>::value)
-        {
-            return data_.d13x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d4_t>::value ||
-                          is_same<X, d8_t>::value || is_same<X, d13_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x13_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x3_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x1_;
-        }
-        else if constexpr(is_same<X, d13_t>::value)
-        {
-            return data_.d13x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
+    using type = vector_type<T, N0 * N1>;
 };
 
-template <typename T>
-struct vector_type<T, 16, typename ck::enable_if_t<is_native_type<T>()>>
+template <typename T, index_t N0, index_t N1>
+struct vector_type_maker<vector_type<T, N1>, N0>
 {
-    using d1_t = T;
-    typedef T d2_t __attribute__((ext_vector_type(2)));
-    typedef T d4_t __attribute__((ext_vector_type(4)));
-    typedef T d8_t __attribute__((ext_vector_type(8)));
-    typedef T d16_t __attribute__((ext_vector_type(16)));
-
-    using type = d16_t;
-
-    union
-    {
-        d16_t d16_;
-        StaticallyIndexedArray<d1_t, 16> d1x16_;
-        StaticallyIndexedArray<d2_t, 8> d2x8_;
-        StaticallyIndexedArray<d4_t, 4> d4x4_;
-        StaticallyIndexedArray<d8_t, 2> d8x2_;
-        StaticallyIndexedArray<d16_t, 1> d16x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x16_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x8_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x4_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x2_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x16_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x8_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x4_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x2_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
+    using type = vector_type<T, N0 * N1>;
 };
 
-template <typename T>
-struct vector_type<T, 32, typename ck::enable_if_t<is_native_type<T>()>>
-{
-    using d1_t = T;
-    typedef T d2_t __attribute__((ext_vector_type(2)));
-    typedef T d4_t __attribute__((ext_vector_type(4)));
-    typedef T d8_t __attribute__((ext_vector_type(8)));
-    typedef T d16_t __attribute__((ext_vector_type(16)));
-    typedef T d32_t __attribute__((ext_vector_type(32)));
-
-    using type = d32_t;
-
-    union
-    {
-        d32_t d32_;
-        StaticallyIndexedArray<d1_t, 32> d1x32_;
-        StaticallyIndexedArray<d2_t, 16> d2x16_;
-        StaticallyIndexedArray<d4_t, 8> d4x8_;
-        StaticallyIndexedArray<d8_t, 4> d8x4_;
-        StaticallyIndexedArray<d16_t, 2> d16x2_;
-        StaticallyIndexedArray<d32_t, 1> d32x1_;
-    } data_ = {d32_t{0}};
-
-    __attribute__((host)) __attribute__((device)) constexpr vector_type() {}
-
-    __attribute__((host)) __attribute__((device)) constexpr vector_type(type v) { (void)v; }
-
-    // __host__ __device__ constexpr vector_type() : data_{type{0}} {}
-
-    // __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value || is_same<X, d32_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x32_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x16_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x8_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x4_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x2_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value || is_same<X, d32_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x32_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x16_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x8_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x4_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x2_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-template <typename T>
-struct vector_type<T, 64, typename ck::enable_if_t<is_native_type<T>()>>
-{
-    using d1_t = T;
-    typedef T d2_t __attribute__((ext_vector_type(2)));
-    typedef T d4_t __attribute__((ext_vector_type(4)));
-    typedef T d8_t __attribute__((ext_vector_type(8)));
-    typedef T d16_t __attribute__((ext_vector_type(16)));
-    typedef T d32_t __attribute__((ext_vector_type(32)));
-    typedef T d64_t __attribute__((ext_vector_type(64)));
-
-    using type = d64_t;
-
-    union
-    {
-        d64_t d64_;
-        StaticallyIndexedArray<d1_t, 64> d1x64_;
-        StaticallyIndexedArray<d2_t, 32> d2x32_;
-        StaticallyIndexedArray<d4_t, 16> d4x16_;
-        StaticallyIndexedArray<d8_t, 8> d8x8_;
-        StaticallyIndexedArray<d16_t, 4> d16x4_;
-        StaticallyIndexedArray<d32_t, 2> d32x2_;
-        StaticallyIndexedArray<d64_t, 1> d64x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value || is_same<X, d32_t>::value ||
-                          is_same<X, d64_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x64_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x32_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x16_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x8_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x4_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x2_;
-        }
-        else if constexpr(is_same<X, d64_t>::value)
-        {
-            return data_.d64x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value || is_same<X, d32_t>::value ||
-                          is_same<X, d64_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x64_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x32_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x16_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x8_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x4_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x2_;
-        }
-        else if constexpr(is_same<X, d64_t>::value)
-        {
-            return data_.d64x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-template <typename T>
-struct vector_type<T, 128, typename ck::enable_if_t<is_native_type<T>()>>
-{
-    using d1_t = T;
-    typedef T d2_t __attribute__((ext_vector_type(2)));
-    typedef T d4_t __attribute__((ext_vector_type(4)));
-    typedef T d8_t __attribute__((ext_vector_type(8)));
-    typedef T d16_t __attribute__((ext_vector_type(16)));
-    typedef T d32_t __attribute__((ext_vector_type(32)));
-    typedef T d64_t __attribute__((ext_vector_type(64)));
-    typedef T d128_t __attribute__((ext_vector_type(128)));
-
-    using type = d128_t;
-
-    union
-    {
-        d128_t d128_;
-        StaticallyIndexedArray<d1_t, 128> d1x128_;
-        StaticallyIndexedArray<d2_t, 64> d2x64_;
-        StaticallyIndexedArray<d4_t, 32> d4x32_;
-        StaticallyIndexedArray<d8_t, 16> d8x16_;
-        StaticallyIndexedArray<d16_t, 8> d16x8_;
-        StaticallyIndexedArray<d32_t, 4> d32x4_;
-        StaticallyIndexedArray<d64_t, 2> d64x2_;
-        StaticallyIndexedArray<d128_t, 1> d128x1_;
-    } data_ = {d128_t{0}};
-
-    __attribute__((host)) __attribute__((device)) constexpr vector_type() {}
-
-    __attribute__((host)) __attribute__((device)) constexpr vector_type(type v) { (void)v; }
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value || is_same<X, d32_t>::value ||
-                          is_same<X, d64_t>::value || is_same<X, d128_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x128_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x64_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x32_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x16_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x8_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x4_;
-        }
-        else if constexpr(is_same<X, d64_t>::value)
-        {
-            return data_.d64x2_;
-        }
-        else if constexpr(is_same<X, d128_t>::value)
-        {
-            return data_.d128x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value || is_same<X, d32_t>::value ||
-                          is_same<X, d64_t>::value || is_same<X, d128_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x128_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x64_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x32_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x16_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x8_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x4_;
-        }
-        else if constexpr(is_same<X, d64_t>::value)
-        {
-            return data_.d64x2_;
-        }
-        else if constexpr(is_same<X, d128_t>::value)
-        {
-            return data_.d128x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-template <typename T>
-struct vector_type<T, 256, typename ck::enable_if_t<is_native_type<T>()>>
-{
-    using d1_t = T;
-    typedef T d2_t __attribute__((ext_vector_type(2)));
-    typedef T d4_t __attribute__((ext_vector_type(4)));
-    typedef T d8_t __attribute__((ext_vector_type(8)));
-    typedef T d16_t __attribute__((ext_vector_type(16)));
-    typedef T d32_t __attribute__((ext_vector_type(32)));
-    typedef T d64_t __attribute__((ext_vector_type(64)));
-    typedef T d128_t __attribute__((ext_vector_type(128)));
-    typedef T d256_t __attribute__((ext_vector_type(256)));
-
-    using type = d256_t;
-
-    union
-    {
-        d256_t d256_;
-        StaticallyIndexedArray<d1_t, 256> d1x256_;
-        StaticallyIndexedArray<d2_t, 128> d2x128_;
-        StaticallyIndexedArray<d4_t, 64> d4x64_;
-        StaticallyIndexedArray<d8_t, 32> d8x32_;
-        StaticallyIndexedArray<d16_t, 16> d16x16_;
-        StaticallyIndexedArray<d32_t, 8> d32x8_;
-        StaticallyIndexedArray<d64_t, 4> d64x4_;
-        StaticallyIndexedArray<d128_t, 2> d128x2_;
-        StaticallyIndexedArray<d256_t, 1> d256x1_;
-    } data_ = {d256_t{0}};
-
-    __attribute__((host)) __attribute__((device)) constexpr vector_type() {}
-
-    __attribute__((host)) __attribute__((device)) constexpr vector_type(type v) { (void)v; }
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(
-            is_same<X, d1_t>::value || is_same<X, d2_t>::value || is_same<X, d4_t>::value ||
-                is_same<X, d8_t>::value || is_same<X, d16_t>::value || is_same<X, d32_t>::value ||
-                is_same<X, d64_t>::value || is_same<X, d128_t>::value || is_same<X, d256_t>::value,
-            "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x256_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x128_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x64_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x32_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x16_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x8_;
-        }
-        else if constexpr(is_same<X, d64_t>::value)
-        {
-            return data_.d64x4_;
-        }
-        else if constexpr(is_same<X, d128_t>::value)
-        {
-            return data_.d128x2_;
-        }
-        else if constexpr(is_same<X, d256_t>::value)
-        {
-            return data_.d256x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(
-            is_same<X, d1_t>::value || is_same<X, d2_t>::value || is_same<X, d4_t>::value ||
-                is_same<X, d8_t>::value || is_same<X, d16_t>::value || is_same<X, d32_t>::value ||
-                is_same<X, d64_t>::value || is_same<X, d128_t>::value || is_same<X, d256_t>::value,
-            "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x256_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x128_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x64_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x32_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x16_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x8_;
-        }
-        else if constexpr(is_same<X, d64_t>::value)
-        {
-            return data_.d64x4_;
-        }
-        else if constexpr(is_same<X, d128_t>::value)
-        {
-            return data_.d128x2_;
-        }
-        else if constexpr(is_same<X, d256_t>::value)
-        {
-            return data_.d256x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-template <typename T, index_t N, typename Enable = void>
-struct non_native_vector_base;
-
-template <typename T>
-struct nnvb_data_t_selector
-{
-    using type = unsigned _BitInt(8 * sizeof(T));
-};
-
-template <>
-struct nnvb_data_t_selector<f8_ocp_t>
-{
-    using type = f8_ocp_t::data_type;
-};
-
-template <>
-struct nnvb_data_t_selector<bf8_ocp_t>
-{
-    using type = bf8_ocp_t::data_type;
-};
-
-#ifndef CK_CODE_GEN_RTC
-template <>
-struct nnvb_data_t_selector<f8_fnuz_t>
-{
-    using type = f8_fnuz_t::data_type;
-};
-
-template <>
-struct nnvb_data_t_selector<bf8_fnuz_t>
-{
-    using type = bf8_fnuz_t::data_type;
-};
-
-template <>
-struct nnvb_data_t_selector<e8m0_bexp_t>
-{
-    using type = e8m0_bexp_t::type;
-};
-#endif
-
-template <>
-struct nnvb_data_t_selector<f6x16_pk_t>
-{
-    using type = f6x16_pk_t::storage_type;
-};
-
-template <>
-struct nnvb_data_t_selector<f6x32_pk_t>
-{
-    using type = f6x32_pk_t::storage_type;
-};
-
-template <>
-struct nnvb_data_t_selector<bf6x16_pk_t>
-{
-    using type = bf6x16_pk_t::storage_type;
-};
-
-template <>
-struct nnvb_data_t_selector<bf6x32_pk_t>
-{
-    using type = bf6x32_pk_t::storage_type;
-};
-
-template <>
-struct nnvb_data_t_selector<pk_i4_t>
-{
-    using type = pk_i4_t::type;
-};
-
-template <>
-struct nnvb_data_t_selector<f4x2_pk_t>
-{
-    using type = f4x2_pk_t::type;
-};
+template <typename T, index_t N>
+using vector_type_maker_t = typename vector_type_maker<T, N>::type;
 
 template <typename T, index_t N>
-struct non_native_vector_base<
-    T,
-    N,
-    ck::enable_if_t<sizeof(T) == 1 || sizeof(T) == 2 || sizeof(T) == 4 || sizeof(T) == 8>>
+__host__ __device__ constexpr auto make_vector_type(Number<N>)
 {
-    using data_t = typename nnvb_data_t_selector<T>::type; // select data_t based on the size of T
-    static_assert(sizeof(T) == sizeof(data_t), "non_native_vector_base storage size mismatch");
-    using data_v = data_t __attribute__((ext_vector_type(N)));
-    using type   = non_native_vector_base<T, N>;
-
-    union alignas(next_pow2(N * sizeof(T)))
-    {
-        data_v dN; // storage vector;
-        StaticallyIndexedArray<data_t, N> dxN;
-        StaticallyIndexedArray<T, N> dTxN;
-        StaticallyIndexedArray<data_v, 1> dNx1;
-    } data_;
-
-    __host__ __device__ constexpr non_native_vector_base(data_t a) : data_{data_v(a)} {}
-    __host__ __device__ constexpr non_native_vector_base(T f)
-        : non_native_vector_base(bit_cast<data_t>(f))
-    {
-    }
-    __host__ __device__ constexpr non_native_vector_base() : non_native_vector_base(T{}){};
-    __host__ __device__ constexpr non_native_vector_base(data_v v) : data_{v} {}
-
-    __host__ __device__ constexpr operator data_v() const { return data_.dN; }
-    __host__ __device__ constexpr operator data_t() const
-    {
-        if constexpr(N == 1)
-        {
-            return data_.dxN[Number<0>{}];
-        }
-        else
-        {
-            return data_.dxN; // XXX this should cause an error
-        }
-    }
-    __host__ __device__ constexpr operator T() const
-    {
-        if constexpr(N == 1)
-        {
-            return data_.dTxN[Number<0>{}];
-        }
-        else
-        {
-            return data_.dTxN; // XXX this should cause an error
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same_v<X, data_t> || is_same_v<X, T> || is_same_v<X, data_v>,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same_v<X, data_t>)
-        {
-            return data_.dxN;
-        }
-        else if constexpr(is_same_v<X, T>)
-        {
-            return data_.dTxN;
-        }
-        else if constexpr(is_same_v<X, data_v>)
-        {
-            return data_.dNx1;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same_v<X, data_t> || is_same_v<X, T> || is_same_v<X, data_v>,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same_v<X, data_t>)
-        {
-            return data_.dxN;
-        }
-        else if constexpr(is_same_v<X, T>)
-        {
-            return data_.dTxN;
-        }
-        else if constexpr(is_same_v<X, data_v>)
-        {
-            return data_.dNx1;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-// implementation for f6x16 and f6x32
-template <typename T, index_t N>
-struct non_native_vector_base<
-    T,
-    N,
-    ck::enable_if_t<sizeof(T) == 12 || sizeof(T) == 16 || sizeof(T) == 24 || sizeof(T) == 32>>
-{
-    using data_t =
-        typename nnvb_data_t_selector<T>::type; // select data_t based on declared base type
-    using element_t = typename T::element_type; // select element_t based on declared element type
-    static_assert(sizeof(T) == sizeof(data_t), "non_native_vector_base storage size mismatch");
-    static constexpr size_t size_factor = sizeof(data_t) / sizeof(element_t);
-    using data_v = element_t __attribute__((ext_vector_type(N * size_factor)));
-    using type   = non_native_vector_base<T, N>;
-
-    union alignas(next_pow2(N * sizeof(T)))
-    {
-        data_v dN; // storage vector;
-        StaticallyIndexedArray<data_t, N> dxN;
-        StaticallyIndexedArray<T, N> dTxN;
-        StaticallyIndexedArray<data_v, 1> dNx1;
-    } data_;
-
-    // Broadcast single value to vector
-    __host__ __device__ constexpr non_native_vector_base(data_t a) : data_{}
-    {
-        // TODO: consider removing initialization similar to vector_type<T, 256>
-
-        ck::static_for<0, N, 1>{}([&](auto i) {
-            data_.dxN(i) = a; // broadcast value to all elements
-        });
-    }
-
-    __host__ __device__ constexpr non_native_vector_base(T f)
-        : non_native_vector_base(bit_cast<data_t>(f))
-    {
-    }
-
-    __host__ __device__ constexpr non_native_vector_base() : non_native_vector_base(T{}){};
-
-    __host__ __device__ constexpr non_native_vector_base(data_v v) : data_{v} {}
-
-    __host__ __device__ constexpr non_native_vector_base(element_t v) : data_{data_v(v)} {}
-
-    __host__ __device__ constexpr operator data_v() const { return data_.dN; }
-
-    __host__ __device__ constexpr operator T() const
-    {
-        if constexpr(N == 1)
-        {
-            return data_.dTxN[Number<0>{}];
-        }
-        else
-        {
-            return err; // XXX this should cause an error
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same_v<X, data_t> || is_same_v<X, data_v> || is_same_v<X, T>,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same_v<X, data_v>)
-        {
-            return data_.dNx1;
-        }
-        else if constexpr(is_same_v<X, data_t>)
-        {
-            return data_.dxN;
-        }
-        else if constexpr(is_same_v<X, T>)
-        {
-            return data_.dTxN;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-template <typename T, index_t N>
-struct scalar_type<non_native_vector_base<
-    T,
-    N,
-    ck::enable_if_t<sizeof(T) == 1 || sizeof(T) == 2 || sizeof(T) == 4 || sizeof(T) == 8>>>
-{
-    using type                           = typename non_native_vector_base<T, N>::data_t;
-    static constexpr index_t vector_size = N;
-};
-
-template <typename T, index_t N>
-struct scalar_type<non_native_vector_base<
-    T,
-    N,
-    ck::enable_if_t<sizeof(T) == 12 || sizeof(T) == 16 || sizeof(T) == 24 || sizeof(T) == 32>>>
-{
-    using type                           = typename non_native_vector_base<T, N>::element_t;
-    static constexpr index_t vector_size = N * non_native_vector_base<T, N>::size_factor;
-};
-
-// non-native vector_type implementation
-template <typename T>
-struct vector_type<T, 1, typename ck::enable_if_t<!is_native_type<T>()>>
-{
-    using d1_t     = T;
-    using d1_nnv_t = non_native_vector_base<T, 1>;
-    using type     = d1_nnv_t;
-
-    union alignas(next_pow2(1 * sizeof(T)))
-    {
-        d1_t d1_;
-        StaticallyIndexedArray<d1_t, 1> d1x1_;
-        d1_nnv_t d1_nnv_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{d1_t{}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value)
-        {
-            return data_.d1x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value)
-        {
-            return data_.d1x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-template <typename T>
-struct vector_type<T, 2, typename ck::enable_if_t<!is_native_type<T>()>>
-{
-    using d1_t     = T;
-    using d1_nnv_t = non_native_vector_base<T, 1>;
-    using d2_t     = non_native_vector_base<T, 2>;
-
-    using type = d2_t;
-
-    union alignas(next_pow2(2 * sizeof(T)))
-    {
-        d2_t d2_;
-        StaticallyIndexedArray<d1_t, 2> d1x2_;
-        StaticallyIndexedArray<d2_t, 1> d2x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value ||
-                          is_same<X, d2_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value)
-        {
-            return data_.d1x2_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value ||
-                          is_same<X, d2_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value)
-        {
-            return data_.d1x2_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-template <typename T>
-struct vector_type<T, 4, typename ck::enable_if_t<!is_native_type<T>()>>
-{
-    using d1_t     = T;
-    using d1_nnv_t = non_native_vector_base<T, 1>;
-    using d2_t     = non_native_vector_base<T, 2>;
-    using d4_t     = non_native_vector_base<T, 4>;
-
-    using type = d4_t;
-
-    union alignas(next_pow2(4 * sizeof(T)))
-    {
-        d4_t d4_;
-        StaticallyIndexedArray<d1_t, 4> d1x4_;
-        StaticallyIndexedArray<d2_t, 2> d2x2_;
-        StaticallyIndexedArray<d4_t, 1> d4x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value ||
-                          is_same<X, d2_t>::value || is_same<X, d4_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value)
-        {
-            return data_.d1x4_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x2_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value ||
-                          is_same<X, d2_t>::value || is_same<X, d4_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value)
-        {
-            return data_.d1x4_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x2_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-template <typename T>
-struct vector_type<T, 8, typename ck::enable_if_t<!is_native_type<T>()>>
-{
-    using d1_t     = T;
-    using d1_nnv_t = non_native_vector_base<T, 1>;
-    using d2_t     = non_native_vector_base<T, 2>;
-    using d4_t     = non_native_vector_base<T, 4>;
-    using d8_t     = non_native_vector_base<T, 8>;
-
-    using type = d8_t;
-
-    union alignas(next_pow2(8 * sizeof(T)))
-    {
-        d8_t d8_;
-        StaticallyIndexedArray<d1_t, 8> d1x8_;
-        StaticallyIndexedArray<d2_t, 4> d2x4_;
-        StaticallyIndexedArray<d4_t, 2> d4x2_;
-        StaticallyIndexedArray<d8_t, 1> d8x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value ||
-                          is_same<X, d2_t>::value || is_same<X, d4_t>::value ||
-                          is_same<X, d8_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value)
-        {
-            return data_.d1x8_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x4_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x2_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value ||
-                          is_same<X, d2_t>::value || is_same<X, d4_t>::value ||
-                          is_same<X, d8_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value)
-        {
-            return data_.d1x8_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x4_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x2_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-template <typename T>
-struct vector_type<T, 16, typename ck::enable_if_t<!is_native_type<T>()>>
-{
-    using d1_t     = T;
-    using d1_nnv_t = non_native_vector_base<T, 1>;
-    using d2_t     = non_native_vector_base<T, 2>;
-    using d4_t     = non_native_vector_base<T, 4>;
-    using d8_t     = non_native_vector_base<T, 8>;
-    using d16_t    = non_native_vector_base<T, 16>;
-
-    using type = d16_t;
-
-    union alignas(next_pow2(16 * sizeof(T)))
-    {
-        d16_t d16_;
-        StaticallyIndexedArray<d1_t, 16> d1x16_;
-        StaticallyIndexedArray<d2_t, 8> d2x8_;
-        StaticallyIndexedArray<d4_t, 4> d4x4_;
-        StaticallyIndexedArray<d8_t, 2> d8x2_;
-        StaticallyIndexedArray<d16_t, 1> d16x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value ||
-                          is_same<X, d2_t>::value || is_same<X, d4_t>::value ||
-                          is_same<X, d8_t>::value || is_same<X, d16_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value)
-        {
-            return data_.d1x16_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x8_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x4_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x2_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value ||
-                          is_same<X, d2_t>::value || is_same<X, d4_t>::value ||
-                          is_same<X, d8_t>::value || is_same<X, d16_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value || is_same<X, d1_nnv_t>::value)
-        {
-            return data_.d1x16_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x8_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x4_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x2_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-template <typename T>
-struct vector_type<T, 32, typename ck::enable_if_t<!is_native_type<T>()>>
-{
-    using d1_t  = T;
-    using d2_t  = non_native_vector_base<T, 2>;
-    using d4_t  = non_native_vector_base<T, 4>;
-    using d8_t  = non_native_vector_base<T, 8>;
-    using d16_t = non_native_vector_base<T, 16>;
-    using d32_t = non_native_vector_base<T, 32>;
-
-    using type = d32_t;
-
-    union alignas(next_pow2(32 * sizeof(T)))
-    {
-        d32_t d32_;
-        StaticallyIndexedArray<d1_t, 32> d1x32_;
-        StaticallyIndexedArray<d2_t, 16> d2x16_;
-        StaticallyIndexedArray<d4_t, 8> d4x8_;
-        StaticallyIndexedArray<d8_t, 4> d8x4_;
-        StaticallyIndexedArray<d16_t, 2> d16x2_;
-        StaticallyIndexedArray<d32_t, 1> d32x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value || is_same<X, d32_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x32_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x16_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x8_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x4_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x2_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value || is_same<X, d32_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x32_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x16_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x8_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x4_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x2_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
-
-template <typename T>
-struct vector_type<T, 64, typename ck::enable_if_t<!is_native_type<T>()>>
-{
-    using d1_t  = T;
-    using d2_t  = non_native_vector_base<T, 2>;
-    using d4_t  = non_native_vector_base<T, 4>;
-    using d8_t  = non_native_vector_base<T, 8>;
-    using d16_t = non_native_vector_base<T, 16>;
-    using d32_t = non_native_vector_base<T, 32>;
-    using d64_t = non_native_vector_base<T, 64>;
-
-    using type = d64_t;
-
-    union alignas(next_pow2(64 * sizeof(T)))
-    {
-        d64_t d64_;
-        StaticallyIndexedArray<d1_t, 64> d1x64_;
-        StaticallyIndexedArray<d2_t, 32> d2x32_;
-        StaticallyIndexedArray<d4_t, 16> d4x16_;
-        StaticallyIndexedArray<d8_t, 8> d8x8_;
-        StaticallyIndexedArray<d16_t, 4> d16x4_;
-        StaticallyIndexedArray<d32_t, 2> d32x2_;
-        StaticallyIndexedArray<d64_t, 1> d64x1_;
-    } data_;
-
-    __host__ __device__ constexpr vector_type() : data_{type{}} {}
-
-    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
-
-    template <typename X>
-    __host__ __device__ constexpr const auto& AsType() const
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value || is_same<X, d32_t>::value ||
-                          is_same<X, d64_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x64_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x32_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x16_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x8_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x4_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x2_;
-        }
-        else if constexpr(is_same<X, d64_t>::value)
-        {
-            return data_.d64x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-
-    template <typename X>
-    __host__ __device__ constexpr auto& AsType()
-    {
-        static_assert(is_same<X, d1_t>::value || is_same<X, d2_t>::value ||
-                          is_same<X, d4_t>::value || is_same<X, d8_t>::value ||
-                          is_same<X, d16_t>::value || is_same<X, d32_t>::value ||
-                          is_same<X, d64_t>::value,
-                      "Something went wrong, please check src and dst types.");
-
-        if constexpr(is_same<X, d1_t>::value)
-        {
-            return data_.d1x64_;
-        }
-        else if constexpr(is_same<X, d2_t>::value)
-        {
-            return data_.d2x32_;
-        }
-        else if constexpr(is_same<X, d4_t>::value)
-        {
-            return data_.d4x16_;
-        }
-        else if constexpr(is_same<X, d8_t>::value)
-        {
-            return data_.d8x8_;
-        }
-        else if constexpr(is_same<X, d16_t>::value)
-        {
-            return data_.d16x4_;
-        }
-        else if constexpr(is_same<X, d32_t>::value)
-        {
-            return data_.d32x2_;
-        }
-        else if constexpr(is_same<X, d64_t>::value)
-        {
-            return data_.d64x1_;
-        }
-        else
-        {
-            return err;
-        }
-    }
-};
+    return typename vector_type_maker<T, N>::type{};
+}
 
 // fp32
 using float2_t  = typename vector_type<float, 2>::type;
diff --git a/include/ck/utility/math.hpp b/include/ck/utility/math.hpp
index b2ebf4b3713..f43a2c0f9ff 100644
--- a/include/ck/utility/math.hpp
+++ b/include/ck/utility/math.hpp
@@ -222,16 +222,27 @@ template <index_t X>
 __host__ __device__ constexpr auto next_power_of_two()
 {
     // TODO: X need to be 2 ~ 0x7fffffff. 0, 1, or larger than 0x7fffffff will compile fail
-    constexpr index_t Y = 1 << (32 - __builtin_clz(X - 1));
+    constexpr index_t Y = X > 1 ? (1 << (32 - __builtin_clz(X - 1))) : X;
     return Y;
 }
 
 template <index_t X>
-__host__ __device__ constexpr auto next_power_of_two(Number<X> x)
+__host__ __device__ constexpr auto next_power_of_two(Number<X>)
 {
-    // TODO: X need to be 2 ~ 0x7fffffff. 0, 1, or larger than 0x7fffffff will compile fail
-    constexpr index_t Y = 1 << (32 - __builtin_clz(x.value - 1));
-    return Number<Y>{};
+    return Number<next_power_of_two<X>>{};
+}
+
+__host__ __device__ constexpr int32_t integer_log2_floor(int32_t x)
+{
+    // TODO: x need to be 1 ~ 0x7fffffff
+    // __builtin_clz will produce unexpected result if x is 0;
+    return 31 - __builtin_clz(x);
+}
+
+__host__ __device__ constexpr bool is_power_of_two_integer(int32_t x)
+{
+    // TODO: x need to be 1 ~ 0x7fffffff
+    return x == (1 << integer_log2_floor(x));
 }
 
 } // namespace math
diff --git a/include/ck/utility/statically_indexed_array.hpp b/include/ck/utility/statically_indexed_array.hpp
index d0735a32f6d..fc905a43eb0 100644
--- a/include/ck/utility/statically_indexed_array.hpp
+++ b/include/ck/utility/statically_indexed_array.hpp
@@ -8,6 +8,7 @@
 #include "sequence.hpp"
 #include "tuple.hpp"
 
+
 namespace ck {
 
 namespace detail {

From efaf0a7a5ce6f3e93da97783e55675471f2cc0aa Mon Sep 17 00:00:00 2001
From: Christopher Millette <63608002+cgmillette@users.noreply.github.com>
Date: Fri, 23 Jan 2026 11:03:51 -0700
Subject: [PATCH 2/7] Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 include/ck/utility/dtype_vector.hpp | 34 +++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/include/ck/utility/dtype_vector.hpp b/include/ck/utility/dtype_vector.hpp
index 85f50b11d30..8cd73d30ead 100644
--- a/include/ck/utility/dtype_vector.hpp
+++ b/include/ck/utility/dtype_vector.hpp
@@ -212,6 +212,30 @@ struct non_native_vector_base<
             return err;
         }
     }
+
+    template <typename X>
+    __host__ __device__ auto& AsType()
+    {
+        static_assert(is_same_v<X, data_t> || is_same_v<X, data_v> || is_same_v<X, T>,
+                      "Something went wrong, please check src and dst types.");
+
+        if constexpr(is_same_v<X, data_v>)
+        {
+            return data_.dNx1;
+        }
+        else if constexpr(is_same_v<X, data_t>)
+        {
+            return data_.dxN;
+        }
+        else if constexpr(is_same_v<X, T>)
+        {
+            return data_.dTxN;
+        }
+        else
+        {
+            return err;
+        }
+    }
 };
 
 template <typename T, index_t N>
@@ -275,7 +299,6 @@ template <typename T, index_t Rank>
 struct vector_type_storage<T, Rank, enable_if_t<is_native_type<T>() && (Rank > 1)>>
 {
     using type = std::conditional_t<std::is_same_v<bool, T>, T[Rank], NativeVectorT<T, Rank>>;
-    ;
 };
 
 /**
@@ -336,6 +359,13 @@ struct is_same_storage_class<VecT<T, LhsRank>, VecT<T, RhsRank>> : true_type
 {
 };
 
+// Explicit specialization for non_native_vector_base, which has three template parameters
+template <typename T, index_t LhsRank, index_t RhsRank, typename Enable>
+struct is_same_storage_class<
+    non_native_vector_base<T, LhsRank, Enable>,
+    non_native_vector_base<T, RhsRank, Enable>> : true_type
+{
+};
 /**
  * @brief C-style arrays of same base type with different ranks
  * @tparam T The base element type
@@ -421,7 +451,7 @@ struct vector_type
              (TraitsX::vector_size == Rank)) &&
             (TraitsX::vector_size <= Rank);
 
-        return is_valid_cast || is_valid_rank;
+        return is_valid_cast && is_valid_rank;
     }
 
     /**

From 525208da140f20eee6f385ccf27938130a38f00b Mon Sep 17 00:00:00 2001
From: Chris Millette <chris.millette@amd.com>
Date: Fri, 23 Jan 2026 19:57:20 +0000
Subject: [PATCH 3/7] Fixes clang formatting and adjust storage class type
 check

---
 include/ck/utility/dtype_vector.hpp           | 45 +++++++++----------
 .../ck/utility/statically_indexed_array.hpp   |  1 -
 2 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/include/ck/utility/dtype_vector.hpp b/include/ck/utility/dtype_vector.hpp
index 8cd73d30ead..b4494f08ea7 100644
--- a/include/ck/utility/dtype_vector.hpp
+++ b/include/ck/utility/dtype_vector.hpp
@@ -322,8 +322,8 @@ template <typename T, index_t Rank>
 using vector_type_storage_t = typename vector_type_storage<T, Rank>::type;
 
 /**
- * @brief Trait to check whether one storage class is the same as another (e.g., same scalar, or
- * same vector class).
+ * @brief Trait to check whether one vector storage class is the same as another (e.g., same scalar,
+ * or same vector class).
  * @tparam Lhs The source storage type
  * @tparam Rhs The comparator storage type
  *
@@ -333,39 +333,35 @@ using vector_type_storage_t = typename vector_type_storage<T, Rank>::type;
  * - C-style arrays of same base type (may have different ranks)
  */
 template <typename Lhs, typename Rhs>
-struct is_same_storage_class : public false_type
+struct is_same_vector_storage_class : public false_type
 {
 };
 
 /**
- * @brief Same type storage class
- * @tparam T The storage type
+ * @brief Template native vector types of same base type with different ranks
+ * @tparam T The base element type
+ * @tparam LhsRank The rank of the source type
+ * @tparam RhsRank The rank of the comparator type
  */
-template <typename T>
-struct is_same_storage_class<T, T> : public true_type
+template <typename T, index_t LhsRank, index_t RhsRank>
+struct is_same_vector_storage_class<NativeVectorT<T, LhsRank>, NativeVectorT<T, RhsRank>>
+    : true_type
 {
 };
 
 /**
- * @brief Template vector types of same base type with different ranks
- * @tparam VecT The <T, N> vector template class type (e.g., vector_type, NativeVectorT,
- * non_native_vector_base)
+ * @brief Template non-native vector types of same base type with different ranks
  * @tparam T The base element type
  * @tparam LhsRank The rank of the source type
  * @tparam RhsRank The rank of the comparator type
+ * @tparam Enable SFINAE helper
  */
-template <template <typename, index_t> class VecT, typename T, index_t LhsRank, index_t RhsRank>
-struct is_same_storage_class<VecT<T, LhsRank>, VecT<T, RhsRank>> : true_type
-{
-};
-
-// Explicit specialization for non_native_vector_base, which has three template parameters
 template <typename T, index_t LhsRank, index_t RhsRank, typename Enable>
-struct is_same_storage_class<
-    non_native_vector_base<T, LhsRank, Enable>,
-    non_native_vector_base<T, RhsRank, Enable>> : true_type
+struct is_same_vector_storage_class<non_native_vector_base<T, LhsRank, Enable>,
+                                    non_native_vector_base<T, RhsRank, Enable>> : true_type
 {
 };
+
 /**
  * @brief C-style arrays of same base type with different ranks
  * @tparam T The base element type
@@ -373,17 +369,18 @@ struct is_same_storage_class<
  * @tparam RhsRank The rank of the comparator type
  */
 template <typename T, index_t LhsRank, index_t RhsRank>
-struct is_same_storage_class<T[LhsRank], T[RhsRank]> : true_type
+struct is_same_vector_storage_class<T[LhsRank], T[RhsRank]> : true_type
 {
 };
 
 /**
- * @brief Convenience evaluator for is_same_storage_class
+ * @brief Convenience evaluator for is_same_vector_storage_class
  * @tparam Lhs The source storage type
  * @tparam Rhs The comparator storage type
  */
 template <typename Lhs, typename Rhs>
-static constexpr bool is_same_storage_class_v = is_same_storage_class<Lhs, Rhs>::value;
+static constexpr bool is_same_vector_storage_class_v =
+    is_same_vector_storage_class<Lhs, Rhs>::value;
 
 /**
  * @brief Vector type wrapper
@@ -443,7 +440,9 @@ struct vector_type
         using TraitsX = scalar_type<X>;
 
         // Checks storage classes match, with same base type (may have different ranks)
-        constexpr bool is_valid_cast = is_same_storage_class_v<StorageT, X>;
+        constexpr bool is_valid_cast =
+            is_same_vector_storage_class_v<StorageT, X> || // Matching vector storage
+            is_same_v<X, T>;                               // Matching scalar type
 
         // Validate vector ranks
         constexpr bool is_valid_rank =
diff --git a/include/ck/utility/statically_indexed_array.hpp b/include/ck/utility/statically_indexed_array.hpp
index fc905a43eb0..d0735a32f6d 100644
--- a/include/ck/utility/statically_indexed_array.hpp
+++ b/include/ck/utility/statically_indexed_array.hpp
@@ -8,7 +8,6 @@
 #include "sequence.hpp"
 #include "tuple.hpp"
 
-
 namespace ck {
 
 namespace detail {

From 4ce5e93dfecb599e506159baf8d9b9b1a5af0fdb Mon Sep 17 00:00:00 2001
From: Chris Millette <chris.millette@amd.com>
Date: Fri, 23 Jan 2026 23:01:21 +0000
Subject: [PATCH 4/7] Re-home NativeVectorT and ensure partial specialization
 with scalar_type is visible where needed

---
 include/ck/utility/data_type.hpp    | 31 +++++++++++++++++++++--------
 include/ck/utility/dtype_vector.hpp | 26 +++---------------------
 include/ck/utility/type_convert.hpp |  4 ++--
 3 files changed, 28 insertions(+), 33 deletions(-)

diff --git a/include/ck/utility/data_type.hpp b/include/ck/utility/data_type.hpp
index f37437336ad..b2e6061d45e 100644
--- a/include/ck/utility/data_type.hpp
+++ b/include/ck/utility/data_type.hpp
@@ -34,17 +34,20 @@ using f4_t    = unsigned _BitInt(4);
 using f6_t    = _BitInt(6);          // e2m3 format
 using bf6_t   = unsigned _BitInt(6); // e3m2 format
 
+/**
+ * @brief Wrapper for native vector type
+ * @tparam T The element type of the vector
+ * @tparam Rank The number of elements in the vector
+ */
+template <typename T, index_t Rank>
+using NativeVectorT = T __attribute__((ext_vector_type(Rank)));
+
 /**
  * @brief Mapping of incoming type to local native storage type and vector size
  * @tparam T Incoming data type
  */
 template <typename T>
-struct scalar_type
-{
-    // Basic data type mapping to unsigned _BitInt of appropriate size
-    using type                           = unsigned _BitInt(8 * sizeof(T));
-    static constexpr index_t vector_size = 1;
-};
+struct scalar_type;
 
 struct f4x2_pk_t
 {
@@ -97,7 +100,7 @@ struct f6_pk_t
     static constexpr index_t vector_size =
         (packed_size * num_bits_elem) / num_bits_vec_elem; // 3 or 6 element_type units
 
-    using storage_type = element_type __attribute__((ext_vector_type(vector_size)));
+    using storage_type = NativeVectorT<element_type, vector_size>;
     storage_type data_{storage_type(0)}; // packed data
 
     using type = f6_pk_t<BitType, packed_size>;
@@ -118,7 +121,7 @@ struct f6_pk_t
 
     // Broadcast single initialization value to all packed elements
     __host__ __device__ f6_pk_t(const int8_t v)
-        : f6_pk_t(static_cast<int8_t __attribute__((ext_vector_type(packed_size)))>(v))
+        : f6_pk_t(static_cast<NativeVectorT<int8_t, packed_size>>(v))
     {
         // TODO: consider removing initialization similar to vector_type<T, 256>
     }
@@ -366,6 +369,18 @@ struct scalar_type<bf6x16_pk_t>
     static constexpr index_t vector_size = 1;
 };
 
+/**
+ * @brief scalar_type trait override for NativeVectorT
+ * @tparam T The vector type
+ * @tparam Rank The number of elements in the vector
+ */
+template <typename T, index_t Rank>
+struct scalar_type<NativeVectorT<T, Rank>>
+{
+    using type                           = T;
+    static constexpr index_t vector_size = Rank;
+};
+
 template <typename T>
 struct packed_type_info
 {
diff --git a/include/ck/utility/dtype_vector.hpp b/include/ck/utility/dtype_vector.hpp
index b4494f08ea7..7711f2f5330 100644
--- a/include/ck/utility/dtype_vector.hpp
+++ b/include/ck/utility/dtype_vector.hpp
@@ -7,26 +7,6 @@
 
 namespace ck {
 
-/**
- * @brief Wrapper for native vector type
- * @tparam T The element type of the vector
- * @tparam Rank The number of elements in the vector
- */
-template <typename T, index_t Rank>
-using NativeVectorT = T __attribute__((ext_vector_type(Rank)));
-
-/**
- * @brief scalar_type trait override for NativeVectorT
- * @tparam T The vector type
- * @tparam Rank The number of elements in the vector
- */
-template <typename T, index_t Rank>
-struct scalar_type<NativeVectorT<T, Rank>>
-{
-    using type                           = T;
-    static constexpr index_t vector_size = Rank;
-};
-
 __device__ int static err = 0;
 
 template <typename T, index_t N, typename Enable = void>
@@ -40,7 +20,7 @@ struct non_native_vector_base<
 {
     using data_t = typename scalar_type<T>::type; // select data_t based on the size of T
     static_assert(sizeof(T) == sizeof(data_t), "non_native_vector_base storage size mismatch");
-    using data_v = data_t __attribute__((ext_vector_type(N)));
+    using data_v = NativeVectorT<data_t, N>;
     using type   = non_native_vector_base<T, N>;
 
     union alignas(math::next_power_of_two<N * sizeof(T)>())
@@ -143,8 +123,8 @@ struct non_native_vector_base<
     using element_t = typename T::element_type; // select element_t based on declared element type
     static_assert(sizeof(T) == sizeof(data_t), "non_native_vector_base storage size mismatch");
     static constexpr size_t size_factor = sizeof(data_t) / sizeof(element_t);
-    using data_v = element_t __attribute__((ext_vector_type(N * size_factor)));
-    using type   = non_native_vector_base<T, N>;
+    using data_v                        = NativeVectorT<element_t, N * size_factor>;
+    using type                          = non_native_vector_base<T, N>;
 
     union alignas(math::next_power_of_two<N * sizeof(T)>())
     {
diff --git a/include/ck/utility/type_convert.hpp b/include/ck/utility/type_convert.hpp
index b3e399609e3..6a8c312b0e8 100644
--- a/include/ck/utility/type_convert.hpp
+++ b/include/ck/utility/type_convert.hpp
@@ -1841,7 +1841,7 @@ inline __host__ __device__ f6x32_t f6_convert_rne(float32_t x, float scale = 1.0
         float float_array[32];
     } in{x};
 
-    using array_type = uint8_t __attribute__((ext_vector_type(32)));
+    using array_type = NativeVectorT<uint8_t, 32>;
     array_type uint8_array;
 
     // collect the 6-bit values into an array
@@ -2178,7 +2178,7 @@ inline __host__ __device__ bf6x32_t bf6_convert_rne(float32_t x, float scale = 1
         float float_array[32];
     } in{x};
 
-    using array_type = uint8_t __attribute__((ext_vector_type(32)));
+    using array_type = NativeVectorT<uint8_t, 32>;
     array_type uint8_array;
 
     // collect the 6-bit values into an array

From 9ba2fc91c1c886f017ae19532bf658df17af415e Mon Sep 17 00:00:00 2001
From: Chris Millette <chris.millette@amd.com>
Date: Thu, 29 Jan 2026 13:33:38 -0500
Subject: [PATCH 5/7] Fixes scalar_type definition for llvm builtin mma type

---
 include/ck/utility/data_type.hpp | 112 +++++++++++++++++++++++--------
 1 file changed, 85 insertions(+), 27 deletions(-)

diff --git a/include/ck/utility/data_type.hpp b/include/ck/utility/data_type.hpp
index b2e6061d45e..ff0bb10d0cc 100644
--- a/include/ck/utility/data_type.hpp
+++ b/include/ck/utility/data_type.hpp
@@ -34,6 +34,17 @@ using f4_t    = unsigned _BitInt(4);
 using f6_t    = _BitInt(6);          // e2m3 format
 using bf6_t   = unsigned _BitInt(6); // e3m2 format
 
+// native types: double, float, _Float16, ushort, int32_t, int8_t, uint8_t, f8_fnuz_t, bf8_fnuz_t,
+// native types: bool
+template <typename T>
+inline constexpr bool is_native_type()
+{
+    return is_same_v<T, double> || is_same_v<T, float> || is_same_v<T, half_t> ||
+           is_same_v<T, bhalf_t> || is_same_v<T, int32_t> || is_same_v<T, uint32_t> ||
+           is_same_v<T, int8_t> || is_same_v<T, uint8_t> || is_same_v<T, _BitInt(8)> ||
+           is_same_v<T, unsigned _BitInt(8)> || is_same_v<T, bool>;
+}
+
 /**
  * @brief Wrapper for native vector type
  * @tparam T The element type of the vector
@@ -43,11 +54,28 @@ template <typename T, index_t Rank>
 using NativeVectorT = T __attribute__((ext_vector_type(Rank)));
 
 /**
- * @brief Mapping of incoming type to local native storage type and vector size
+ * @brief Mapping of incoming type to local native vector storage type and vector size
  * @tparam T Incoming data type
  */
 template <typename T>
-struct scalar_type;
+struct scalar_type
+{
+    // Basic data type mapping to unsigned _BitInt of appropriate size
+    using type                           = unsigned _BitInt(8 * sizeof(T));
+    static constexpr index_t vector_size = 1;
+};
+
+/**
+ * @brief scalar_type trait override for NativeVectorT
+ * @tparam T The vector type
+ * @tparam Rank The number of elements in the vector
+ */
+template <typename T, index_t Rank>
+struct scalar_type<NativeVectorT<T, Rank>>
+{
+    using type                           = T;
+    static constexpr index_t vector_size = Rank;
+};
 
 struct f4x2_pk_t
 {
@@ -85,6 +113,39 @@ struct f4x2_pk_t
     }
 };
 
+// TODO: Unfortunately, we cannot partially specialize scalar_type for vectors written
+// in the following way:
+// template<typename T, index_t Rank>
+// struct scalar_type<T __attribute__((__vector_size__(sizeof(T) * Rank)))>
+// {
+//     using type                           = T;
+//     static constexpr index_t vector_size = Rank;
+// };
+// The compiler errors out with "partial specialization is not allowed for this type",
+// claiming that the Rank is not a deducible parameter. This might be a compiler bug.
+// Note the above type is classified differently from the NativeVectorT<T, Rank> alias,
+// even though they are functionally equivalent and are trivially constructibe from each other.
+// This is unfortunate, but we have to work around it because some LLVM builtins for some
+// operations (e.g., mma) may return the former type.
+// For now we have to explicitly specialize for each vector size we need. These are used
+// in f6_pk_t below.
+
+/// @brief scalar_type trait override for uint32_t vector of size 3
+template <>
+struct scalar_type<uint32_t __attribute__((__vector_size__(sizeof(uint32_t) * 3)))>
+{
+    using type                           = uint32_t;
+    static constexpr index_t vector_size = 3;
+};
+
+/// @brief scalar_type trait override for uint32_t vector of size 6
+template <>
+struct scalar_type<uint32_t __attribute__((__vector_size__(sizeof(uint32_t) * 6)))>
+{
+    using type                           = uint32_t;
+    static constexpr index_t vector_size = 6;
+};
+
 template <typename BitType, index_t pk_size>
 struct f6_pk_t
 {
@@ -105,16 +166,36 @@ struct f6_pk_t
 
     using type = f6_pk_t<BitType, packed_size>;
 
+    /** This class may trivially constructed by the following vector type alias
+     * for example from a result of an mma operation. This is primarily for internal use.
+     * @note f6x16_pk_t and f6x32_pk_t storage types, may be trivially constructed from
+     *       uint32_t vectors of size 3 and 6 respectively for example from mma operation results.
+     *       Unfortunately, unsigned int __attribute__((ext_vector_type(6))) a.k.a
+     * NativeVectorT<uint32_t, 6> is NOT the same as __attribute__((__vector_size__(6 *
+     * sizeof(unsigned int)))) unsigned int which is returned from the mma ops despite being
+     * functionally equivalent. This class may be trivially constructed from both, so we can steer
+     * the templated ctor below to only consider incoming vectors types other than our two storage
+     * types of interest.
+     */
+    using storage_type_alias =
+        element_type __attribute__((__vector_size__(sizeof(element_type) * vector_size)));
+
     __host__ __device__ constexpr f6_pk_t() {}
     __host__ __device__ constexpr f6_pk_t(const storage_type& init) : data_{init}
     {
         // TODO: consider removing initialization similar to vector_type<T, 256>
     }
 
-    // Initialize from a vector type with the same size as packed_size
-    template <typename T, typename = enable_if_t<scalar_type<T>::vector_size == packed_size>>
+    // Initialize from a vector type with the same size as packed_size.
+    // Exclude storage_type and storage_type_alias because these are trivially constructible.
+    template <
+        typename T,
+        typename = enable_if_t<!is_same_v<T, storage_type> && !is_same_v<T, storage_type_alias> &&
+                               scalar_type<T>::vector_size == packed_size>>
     __host__ __device__ f6_pk_t(const T& v)
     {
+        static_assert(scalar_type<T>::vector_size == packed_size,
+                      "Input vector size must match packed_size.");
         static_for<0, packed_size, 1>{}(
             [&](auto i) { pack(v[static_cast<index_t>(i)], static_cast<index_t>(i)); });
     }
@@ -202,17 +283,6 @@ struct pk_i4_t
     __host__ __device__ constexpr pk_i4_t(type init) : data{init} {}
 };
 
-// native types: double, float, _Float16, ushort, int32_t, int8_t, uint8_t, f8_fnuz_t, bf8_fnuz_t,
-// native types: bool
-template <typename T>
-inline constexpr bool is_native_type()
-{
-    return is_same<T, double>::value || is_same<T, float>::value || is_same<T, half_t>::value ||
-           is_same<T, bhalf_t>::value || is_same<T, int32_t>::value ||
-           is_same<T, uint32_t>::value || is_same<T, int8_t>::value || is_same<T, uint8_t>::value ||
-           is_same_v<T, _BitInt(8)> || is_same_v<T, unsigned _BitInt(8)> || is_same<T, bool>::value;
-}
-
 // is_scalar_type
 template <typename TV>
 struct is_scalar_type
@@ -369,18 +439,6 @@ struct scalar_type<bf6x16_pk_t>
     static constexpr index_t vector_size = 1;
 };
 
-/**
- * @brief scalar_type trait override for NativeVectorT
- * @tparam T The vector type
- * @tparam Rank The number of elements in the vector
- */
-template <typename T, index_t Rank>
-struct scalar_type<NativeVectorT<T, Rank>>
-{
-    using type                           = T;
-    static constexpr index_t vector_size = Rank;
-};
-
 template <typename T>
 struct packed_type_info
 {

From bea883a83c5a374a34c9b0ccc563e8bab397e48e Mon Sep 17 00:00:00 2001
From: Chris Millette <chris.millette@amd.com>
Date: Fri, 30 Jan 2026 15:47:19 -0500
Subject: [PATCH 6/7] Fixes aliasing for vector size of 1

---
 include/ck/utility/dtype_vector.hpp | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/include/ck/utility/dtype_vector.hpp b/include/ck/utility/dtype_vector.hpp
index 7711f2f5330..75008b5e894 100644
--- a/include/ck/utility/dtype_vector.hpp
+++ b/include/ck/utility/dtype_vector.hpp
@@ -454,7 +454,12 @@ struct vector_type
         // Note: We might end up with incomplete quantization from slicing
         // when Rank % TraitsX::vector_size != 0, so take the floor division.
         constexpr index_t newRank = Rank / TraitsX::vector_size;
-        using ResultT             = StaticallyIndexedArray_v2<X, newRank>;
+
+        // Determine the cast type:
+        // - Scalar T if slicing to scalar or vector size of 1,
+        // - X otherwise.
+        using CastT   = conditional_t<TraitsX::vector_size == 1, T, X>;
+        using ResultT = StaticallyIndexedArray_v2<CastT, newRank>;
 
         // As a rule, the aliasing type should not be larger than the original type.
         static_assert(sizeof(ResultT) <= sizeof(vector_type),
@@ -485,7 +490,12 @@ struct vector_type
         // Note: We might end up with incomplete quantization from slicing
         // when Rank % TraitsX::vector_size != 0, so take the floor division.
         constexpr index_t newRank = Rank / TraitsX::vector_size;
-        using ResultT             = StaticallyIndexedArray_v2<X, newRank>;
+
+        // Determine the cast type:
+        // - Scalar T if slicing to scalar or vector size of 1,
+        // - X otherwise.
+        using CastT   = conditional_t<TraitsX::vector_size == 1, T, X>;
+        using ResultT = StaticallyIndexedArray_v2<CastT, newRank>;
 
         // As a rule, the aliasing type should not be larger than the original type.
         static_assert(sizeof(ResultT) <= sizeof(vector_type),

From f551699f908ce0955d21023e47f2b698cfbda292 Mon Sep 17 00:00:00 2001
From: Chris Millette <chris.millette@amd.com>
Date: Fri, 30 Jan 2026 16:23:00 -0500
Subject: [PATCH 7/7] Workaround adjustment to scalar_type<pk_i4_t>::type.
 Skips invalid case for pk_i4_t, but should be addressed in the future.

---
 include/ck/utility/dynamic_buffer.hpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/include/ck/utility/dynamic_buffer.hpp b/include/ck/utility/dynamic_buffer.hpp
index 4e477eed265..00fab270e81 100644
--- a/include/ck/utility/dynamic_buffer.hpp
+++ b/include/ck/utility/dynamic_buffer.hpp
@@ -260,7 +260,10 @@ struct DynamicBuffer
                 x, p_data_, i, is_valid_element, element_space_size_ / PackedSize);
         }
         else if constexpr(GetAddressSpace() == AddressSpaceEnum::Lds &&
-                          is_same<typename scalar_type<remove_cvref_t<T>>::type, int8_t>::value &&
+                          is_same_v<typename scalar_type<remove_cvref_t<T>>::type, int8_t> &&
+                          !is_same_v<remove_cvref_t<T>,
+                                     pk_i4_t> && // TODO: This needs to be fixed for pk_i4_t which
+                                                 // cannot be handled below, but is stored as int8_t
                           workaround_int8_ds_write_issue)
         {
             if(is_valid_element)