CUTLASS 3.3.0 (#1167)

* Release 3.3.0

Adds support for mixed precision GEMMs On Hopper and Ampere
Adds support for < 16B aligned GEMMs on Hopper
Enhancements to EVT
Enhancements to Python interface
Enhancements to Sub-byte type handling in CuTe
Several other bug-fixes and performance improvements.

* minor doc update

include/cute/atom/copy_traits.hpp

@@ -92,6 +92,21 @@ struct Copy_Traits<DefaultCopy>
   using RefLayout = SrcLayout;
 };
 
+namespace detail {
+
+template <class Operation,
+          class PtrS, int... Is,
+          class PtrD, int... Id>
+CUTE_HOST_DEVICE constexpr
+void
+copy_explode(PtrS&& s, int_sequence<Is...>,
+             PtrD&& d, int_sequence<Id...>)
+{
+  return Operation::copy(s[Is]..., d[Id]...);
+}
+
+} // end namespace detail
+
 //
 // Generic copy_unpack for any Copy_Traits
 //
@@ -123,9 +138,8 @@ copy_unpack(Copy_Traits<Operation, Args...> const&,
   CUTE_STATIC_ASSERT_V(size(rD) == Int<RegNumDst>{},
     "In CopyAtom, dst layout doesn't vectorize into registers. This dst layout is incompatible with this tiled copy.");
 
-  detail::explode(Operation::copy,
-                  rS, make_int_sequence<RegNumSrc>{},
-                  rD, make_int_sequence<RegNumDst>{});
+  detail::copy_explode<Operation>(rS, make_int_sequence<RegNumSrc>{},
+                                  rD, make_int_sequence<RegNumDst>{});
 }
 
 //

include/cute/atom/copy_traits_sm80.hpp

@@ -51,6 +51,13 @@ struct Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS<S,D>>
 
   // Reference map from (thr,val) to bit
   using RefLayout = SrcLayout;
+
+  // Construct a zfill variant with a given predicate value
+  CUTE_HOST_DEVICE constexpr
+  Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<S,D>>
+  with(bool pred) const {
+    return {pred};
+  }
 };
 
 template <class S, class D>
@@ -66,6 +73,95 @@ struct Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL<S,D>>
 
   // Reference map from (thr,val) to bit
   using RefLayout = SrcLayout;
+
+  // Construct a zfill variant with a given predicate value
+  CUTE_HOST_DEVICE constexpr
+  Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<S,D>>
+  with(bool pred) const {
+    return {pred};
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<S,D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1,Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1,Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Predicate value that determines whether to load or zfill
+  bool pred = false;
+
+  // Overload copy_unpack for zfill variant to pass the predicate into the op
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  CUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits        const& traits,
+              Tensor<TS,SLayout> const& src,
+              Tensor<TD,DLayout>      & dst)
+  {
+    static_assert(is_gmem<TS>::value, "Expected gmem source for cp.async.");
+    static_assert(is_smem<TD>::value, "Expected smem destination for cp.async.");
+
+    Tensor rS = recast<S>(src);
+    Tensor rD = recast<D>(dst);
+
+    CUTE_STATIC_ASSERT_V(size(rS) == Int<1>{},
+      "In CopyAtom, src layout doesn't vectorize into registers. This src layout is incompatible with this tiled copy.");
+    CUTE_STATIC_ASSERT_V(size(rD) == Int<1>{},
+      "In CopyAtom, dst layout doesn't vectorize into registers. This dst layout is incompatible with this tiled copy.");
+
+    SM80_CP_ASYNC_CACHEALWAYS_ZFILL<S,D>::copy(rS[0], rD[0], traits.pred);
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<S,D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1,Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1,Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Predicate value that determines whether to load or zfill
+  bool pred = false;
+
+  // Overload copy_unpack for zfill variant to pass the predicate into the op
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  CUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits        const& traits,
+              Tensor<TS,SLayout> const& src,
+              Tensor<TD,DLayout>      & dst)
+  {
+    static_assert(is_gmem<TS>::value, "Expected gmem source for cp.async.");
+    static_assert(is_smem<TD>::value, "Expected smem destination for cp.async.");
+
+    Tensor rS = recast<S>(src);
+    Tensor rD = recast<D>(dst);
+
+    CUTE_STATIC_ASSERT_V(size(rS) == Int<1>{},
+      "In CopyAtom, src layout doesn't vectorize into registers. This src layout is incompatible with this tiled copy.");
+    CUTE_STATIC_ASSERT_V(size(rD) == Int<1>{},
+      "In CopyAtom, dst layout doesn't vectorize into registers. This dst layout is incompatible with this tiled copy.");
+
+    SM80_CP_ASYNC_CACHEGLOBAL_ZFILL<S,D>::copy(rS[0], rD[0], traits.pred);
+  }
 };
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////

include/cute/atom/copy_traits_sm90_tma.hpp

@@ -79,6 +79,14 @@ struct Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
   copy_unpack_(void const* const dst_ptr,
                Coord const& src_coord, seq<Is...>) const
   {
+#if 0
+    auto [c0,c1,c2,c3,c4] = append<5>(src_coord, 0);
+    printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
+           threadIdx.x, threadIdx.y, threadIdx.z,
+           blockIdx.x, blockIdx.y, blockIdx.z,
+           int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), dst_ptr);
+#endif
+
     SM90_TMA_LOAD::copy(&tma_desc_, tma_load_mbar_,
                         dst_ptr, get<Is>(src_coord)...);
   }
@@ -185,6 +193,14 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
   copy_unpack_(void const* const dst_ptr,
                Coord const& src_coord, seq<Is...>) const
   {
+#if 0
+    auto [c0,c1,c2,c3,c4] = append<5>(src_coord, 0);
+    printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
+           threadIdx.x, threadIdx.y, threadIdx.z,
+           blockIdx.x, blockIdx.y, blockIdx.z,
+           int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), dst_ptr);
+#endif
+
     SM90_TMA_LOAD_MULTICAST::copy(&tma_desc_, tma_load_mbar_, multicast_mask_,
                                   dst_ptr, get<Is>(src_coord)...);
   }
@@ -298,6 +314,14 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
   copy_unpack_(void const* const src_ptr,
                Coord const& dst_coord, seq<Is...>) const
   {
+#if 0
+    auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0);
+    printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
+           threadIdx.x, threadIdx.y, threadIdx.z,
+           blockIdx.x, blockIdx.y, blockIdx.z,
+           int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), src_ptr);
+#endif
+
     SM90_TMA_STORE::copy(&tma_desc_,
                          src_ptr, get<Is>(dst_coord)...);
   }
@@ -354,8 +378,8 @@ struct Copy_Traits<SM90_BULK_COPY_G2S, NumBitsPerTMA, OpArgs...>
                   "Extra arguments not set. Set .with() before use.");
     static_assert(is_gmem<TS>::value, "Expected gmem src for SM90_BULK_COPY_G2S");
     static_assert(is_smem<TD>::value, "Expected smem dst for SM90_BULK_COPY_G2S");
-    SM90_BULK_COPY_G2S::copy(src.data().get(), *get<0>(traits.bulk_load_mbar_),
-                             dst.data().get(), int32_t(NumBitsPerTMA::value / 8));
+    SM90_BULK_COPY_G2S::copy(raw_pointer_cast(src.data()), *get<0>(traits.bulk_load_mbar_),
+                             raw_pointer_cast(dst.data()), int32_t(NumBitsPerTMA::value / 8));
   }
 
   // Record the memory barrier for the instruction
@@ -390,7 +414,7 @@ struct Copy_Traits<SM90_BULK_COPY_S2G, NumBitsPerTMA>
   {
     static_assert(is_smem<TS>::value, "Expected smem src for SM90_BULK_COPY_S2G");
     static_assert(is_gmem<TD>::value, "Expected gmem dst for SM90_BULK_COPY_S2G");
-    SM90_BULK_COPY_S2G::copy(src.data().get(), dst.data().get(), int32_t(NumBitsPerTMA::value / 8));
+    SM90_BULK_COPY_S2G::copy(raw_pointer_cast(src.data()), raw_pointer_cast(dst.data()), int32_t(NumBitsPerTMA::value / 8));
   }
 };
 
@@ -497,7 +521,7 @@ coalesce_256(Tensor<Engine,Layout> const& tensor)
 //   and construct a TMA Descriptor for the resulting instruction
 // At the same time, construct the Tma Tensor's Stride to generate
 //   the TMA coordinates that the instruction consumes.
-// 
+//
 template <class TmaInternalType,
           class GEngine, class GLayout,
           class SShape, class SStride,
@@ -518,7 +542,7 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor,    // The original GM
   // Perform the tiling to the gmem vector again, but with indirections to the gtensor modes
   auto gbasis = make_identity_layout(shape(gtensor));
   auto tile_gbasis_tmp = gbasis.compose(smem_inv_h);
-  
+
   // Instead of the recast (gbasis doesn't have type info), replace the shape with the already-recasted shape
   // tma_box_shape:gmem_mode
   auto tile_gbasis = make_layout(shape(tile_gstride), stride(tile_gbasis_tmp));
@@ -530,8 +554,8 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor,    // The original GM
   // NOTE This is essentially ArithmeticTuple complement...
   // NOTE   in pursuit of implementing an ArithmeticTuple logical_divide for smem_inv_h
   auto tile_gbasis_remaining_stride = filter_tuple(flatten(shape (gtensor_T)), flatten(stride(gtensor_T)),
-                                                   flatten(stride(gbasis)), 
-                                                   [&](auto s, auto d, auto e) 
+                                                   flatten(stride(gbasis)),
+                                                   [&](auto s, auto d, auto e)
   {
     if constexpr (is_constant<1, decltype(s)>::value || is_constant<0, decltype(d)>::value) {
       return cute::tuple<>{};          // If size-1 or stride-0, then don't append
@@ -551,7 +575,7 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor,    // The original GM
   auto tma_gbasis_tile = tile_gbasis.compose(make_layout(wrap(shape(tma_gstride))));
 
   // Append the remaining basis modes that contribute to the TMA with size-1
-  auto tma_gbasis_full = make_layout(tuple_cat(wrap( shape(tma_gbasis_tile)), wrap(repeat<tile_gbasis_remaining_rank>(Int<1>{}))), 
+  auto tma_gbasis_full = make_layout(tuple_cat(wrap( shape(tma_gbasis_tile)), wrap(repeat<tile_gbasis_remaining_rank>(Int<1>{}))),
                                      tuple_cat(wrap(stride(tma_gbasis_tile)), wrap(tile_gbasis_remaining_stride)));
 
   // Group the trailing modes to make this max rank-5 -- TMA rank limitation
@@ -570,7 +594,7 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor,    // The original GM
 
   //
   // TMA desc creation
-  // 
+  //
 
   constexpr int tma_dim = decltype(rank(tma_gbasis))::value;
 
@@ -579,7 +603,7 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor,    // The original GM
   //
 
   void* gmem_address = (void*) raw_pointer_cast(gtensor_T.data());
-  auto  gmem_layout = gtensor_T.layout();
+  auto  gmem_layout  = gtensor_T.layout();
 
   cute::array<uint64_t, 5> gmem_prob_shape  = {1,1,1,1,1};
   cute::array<uint64_t, 5> gmem_prob_stride = {0,0,0,0,0};
@@ -665,20 +689,20 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor,    // The original GM
     //
     // Construct the descriptor
     //
-  
+
     TmaDescriptor tma_desc = {0};
-  
+
     //
     // TMA general info
     //
-  
+
   #if (__CUDACC_VER_MAJOR__ >= 12) && !defined(__CUDACC_RTC__)
-  
+
     CUtensorMapDataType     tma_format      = TMA::to_CUtensorMapDataType<TmaInternalType>();
     CUtensorMapInterleave   tma_interleave  = CU_TENSOR_MAP_INTERLEAVE_NONE;
     CUtensorMapL2promotion  tma_l2Promotion = CU_TENSOR_MAP_L2_PROMOTION_L2_128B;
     CUtensorMapFloatOOBfill tma_oobFill     = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE;
-  
+
     // TMA smem swizzle type
     CUtensorMapSwizzle smem_swizzle = TMA::to_CUtensorMapSwizzle(get_tma_swizzle_bits(swizzle));
     CUresult result = cuTensorMapEncodeTiled(
@@ -694,7 +718,7 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor,    // The original GM
         smem_swizzle,
         tma_l2Promotion,
         tma_oobFill);
-  
+
     if (result != CUDA_SUCCESS) {
       std::cerr << "TMA Desc Addr:   " << &tma_desc
                 << "\nformat         " << tma_format
@@ -711,8 +735,11 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor,    // The original GM
       std::cerr << "Error: Failed to initialize the TMA descriptor " << result << std::endl;
       assert(false);
     }
-  
+
   #endif // (__CUDACC_VER_MAJOR__ >= 12) && !defined(__CUDACC_RTC__)
+  auto recast_ratio = cute::ratio(Int<sizeof_bits<typename GEngine::value_type>::value>{},
+                                  Int<sizeof_bits<             TmaInternalType>::value>{});
+
   // Finally, get the inverse permutation of the E<i> bases for the mocked gmem stride
   // NOTE This is essentially ArithmeticTuple inverse...
   auto gmem_stride_bases = transform_leaf(stride(gbasis), [&](auto ei) {
@@ -727,9 +754,9 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor,    // The original GM
       [[maybe_unused]] auto j = find_if(tma_gbasis_stride, [&](auto tma_stride_j) { return any_of(tma_stride_j, [&](auto dj) { return dj == EI{}; }); });
       if constexpr (decltype(j == rank(tma_gbasis_stride))::value) {
         return Int<0>{};               // If not-found, return arithmetic identity -- no contribution to the TMA
-      } else 
+      } else
       if constexpr (decltype(j == Int<0>{})::value) {
-        auto scale = ratio(size(tma_gstride), size(smem_inv_h)) * basis_get(ei, stride(gtensor));
+        auto scale = recast_ratio * basis_get(ei, stride(gtensor));
         return E<j>{} * scale;         // Return TMA Coord basis -- with a recast scale factor
       } else
       if constexpr (decltype(rank<j>(tma_gbasis_stride) == Int<1>{})::value) {
@@ -959,21 +986,23 @@ template <class TmaInternalType,
           class CopyOp,
           class GEngine, class GLayout,
           class SLayout,
-          class CTA_Tile,
+          class CTA_Tiler,
           class Cluster_Size>
 CUTE_HOST_RTC
 auto
 make_tma_copy(CopyOp                  const& copy_op,
               Tensor<GEngine,GLayout> const& gtensor,
               SLayout                 const& slayout,
-              CTA_Tile                const& cta_tile,
+              CTA_Tiler               const& cta_tiler,
               Cluster_Size            const& cluster_size)
 {
+  auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler);
+  auto cta_t_tile = make_layout(cluster_size);
   return detail::make_tma_copy_tiled<TmaInternalType>(copy_op,
                                                       gtensor,
                                                       slayout,
-                                                      make_layout(cluster_size),
-                                                      make_identity_layout(cta_tile));
+                                                      cta_t_tile,
+                                                      cta_v_tile);
 }
 
 // Explicit defaulting

include/cute/atom/copy_traits_sm90_tma_swizzle.hpp

@@ -37,12 +37,13 @@
 #include <cuda.h>
 #endif
 
-#include "cute/arch/copy_sm90_desc.hpp"
-#include "cute/swizzle_layout.hpp"
+#include <cute/arch/copy_sm90_desc.hpp>
+#include <cute/swizzle_layout.hpp>
 
 namespace cute::detail {
 
 template <int B, int M, int S>
+CUTE_HOST_DEVICE constexpr
 TMA::SmemSwizzleBits
 get_tma_swizzle_bits(Swizzle<B,M,S>)
 {

include/cute/atom/mma_atom.hpp

@@ -155,7 +155,8 @@ struct MMA_Atom<MMA_Traits<Args...>>
 
     if constexpr (has_dereference<FrgTypeA>::value) {
       // If the intended FrgTypeA is a view (of the current tensor), forward the whole
-      static_assert(is_same<get_raw_type_t<ValTypeA>, typename remove_cvref_t<ATensor>::value_type>::value, "Expecting ValTypeA type");
+      static_assert(is_same<ValTypeA, typename remove_cvref_t<ATensor>::value_type>::value
+                      , "Expecting ValTypeA type");
       return make_tensor<FrgTypeA>(std::forward<ATensor>(atensor));
     } else {
       // Else, the intended FrgTypeA is a value type, construct a new tensor with a fragment layout
@@ -176,7 +177,8 @@ struct MMA_Atom<MMA_Traits<Args...>>
 
     if constexpr (has_dereference<FrgTypeB>::value) {
       // If the intended FrgTypeB is a view (of the current tensor), forward the whole
-      static_assert(is_same<ValTypeB, typename remove_cvref_t<BTensor>::value_type>::value, "Expecting ValTypeB type");
+      static_assert(is_same<ValTypeB, typename remove_cvref_t<BTensor>::value_type>::value
+                      , "Expecting ValTypeB type");
       return make_tensor<FrgTypeB>(std::forward<BTensor>(btensor));
     } else {
       // Else, the intended FrgTypeB is a value type, construct a new tensor with a fragment layout
@@ -224,6 +226,11 @@ struct TiledMMA : MMA_Atom
   // thr_idx -> (ThrV,ThrM,ThrN,ThrK)
   using TidLayout = decltype(right_inverse(ThrLayoutVMNK{}));
 
+  CUTE_HOST_DEVICE constexpr auto
+  get_thr_layout_vmnk() const {
+    return ThrLayoutVMNK{};
+  }
+
   // Tile a tensor or a layout from shape
   //   (M,N,...)
   // to shape
@@ -295,8 +302,8 @@ struct TiledMMA : MMA_Atom
   thrfrg_A(ATensor&& atensor)
   {
     CUTE_STATIC_ASSERT_V(rank(atensor) >= Int<2>{});
-    CUTE_STATIC_ASSERT_V(size<0>(atensor) % size<0>(TiledShape_MNK{}) == Int<0>{});
-    CUTE_STATIC_ASSERT_V(size<1>(atensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
+    //CUTE_STATIC_ASSERT_V(size<0>(atensor) % size<0>(TiledShape_MNK{}) == Int<0>{});
+    //UTE_STATIC_ASSERT_V(size<1>(atensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
 
     // Reorder the tensor for the TiledAtom
     auto t_tile = make_tile(left_inverse(get<0>(PermutationsMNK{})),
@@ -353,8 +360,8 @@ struct TiledMMA : MMA_Atom
   thrfrg_B(BTensor&& btensor)
   {
     CUTE_STATIC_ASSERT_V(rank(btensor) >= Int<2>{});
-    CUTE_STATIC_ASSERT_V(size<0>(btensor) % size<1>(TiledShape_MNK{}) == Int<0>{});
-    CUTE_STATIC_ASSERT_V(size<1>(btensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
+    //CUTE_STATIC_ASSERT_V(size<0>(btensor) % size<1>(TiledShape_MNK{}) == Int<0>{});
+    //CUTE_STATIC_ASSERT_V(size<1>(btensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
 
     // Reorder the tensor for the TiledAtom
     auto t_tile = make_tile(left_inverse(get<1>(PermutationsMNK{})),

include/cute/atom/mma_traits_sm90_gmma.hpp

@@ -117,16 +117,22 @@ using Layout_SW128_Atom = typename conditional<tnsp == GMMA::Major::MN,
                                                Layout_K_SW128_Atom<Type>>::type;
 
 //
-// Tensor to LayoutType utility
+// Tensor (position-dependent swizzle) to LayoutType utility
 //
 
-// smem_ptr_swizzle LayoutType
-template <int B, int M, int S, class Shape, class Stride>
+template <class Engine, class Shape, class Stride>
 CUTE_HOST_DEVICE constexpr
 LayoutType
-layout_type(Tensor<ViewEngine<smem_ptr_swizzle<const uint128_t, Swizzle<B,M,S>>>,
-                   Layout<Shape,Stride>> const&)
+layout_type(Tensor<Engine, Layout<Shape,Stride>> const&)
 {
+  static_assert(is_same<uint128_t, typename Engine::value_type>::value,
+                "Expected uint128_t type in LayoutType conversion.");
+
+  using Swizzle = get_swizzle_t<Engine>;
+  constexpr int B = Swizzle::num_bits;
+  constexpr int M = Swizzle::num_base;
+  constexpr int S = Swizzle::num_shft;
+
   static_assert(M == 4,           "Unsupported layout swizzle");
   static_assert(0 <= B && B <= 3, "Unsupported layout swizzle");
   static_assert(S == 3,           "Unsupported layout swizzle");
@@ -140,16 +146,6 @@ layout_type(Tensor<ViewEngine<smem_ptr_swizzle<const uint128_t, Swizzle<B,M,S>>>
   return LayoutType::INTERLEAVE;  // ERROR
 }
 
-// smem_ptr non-swizzled LayoutType
-template <class Shape, class Stride>
-CUTE_HOST_DEVICE constexpr
-LayoutType
-layout_type(Tensor<ViewEngine<smem_ptr<const uint128_t>>,
-                   Layout<Shape,Stride>> const&)
-{
-  return LayoutType::INTERLEAVE;
-}
-
 ///////////////////////////////////////////////////////////////////////////////
 // Construction method for GMMA Descriptors
 ///////////////////////////////////////////////////////////////////////////////
@@ -211,7 +207,7 @@ make_gmma_desc(Tensor<TEngine,TLayout> const& tensor)
   desc.bitfield.layout_type_ = uint8_t(LAYOUT_TYPE);
 
   // Start address (4LSB not included)
-  uint32_t start_address = cast_smem_ptr_to_uint(u128_tensor.data().get());
+  uint32_t start_address = cast_smem_ptr_to_uint(raw_pointer_cast(u128_tensor.data()));
   desc.bitfield.start_address_ = start_address >> 4;
 
   constexpr uint8_t base_offset = 0;
@@ -314,57 +310,67 @@ make_gmma_desc(Tensor<TEngine,TLayout> const& tensor)
 
 struct DescriptorIterator
 {
+  using reference    = GmmaDescriptor;
+  using element_type = GmmaDescriptor;
+  using value_type   = GmmaDescriptor;
+
   GmmaDescriptor desc_;
 
   // Dereference returns the GmmaDescriptor
   CUTE_HOST_DEVICE constexpr
-  GmmaDescriptor const& operator*() const { return desc_; }
+  reference operator*() const { return desc_; }
 
   // Advance and return a new GmmaDescriptor
   template <class Index>
   CUTE_HOST_DEVICE constexpr
-  GmmaDescriptor operator[](Index const& i) const { return *(*this + i); }
+  reference operator[](Index const& i) const { return *(*this + i); }
 
   // Return an advanced iterator
   template <class Index>
   CUTE_HOST_DEVICE constexpr
   DescriptorIterator operator+(Index const& offset) const
   {
-    return { GmmaDescriptor {desc_ + uint64_t(offset)} };
+    return { GmmaDescriptor{desc_ + uint64_t(offset)} };
   }
 
   CUTE_HOST_DEVICE friend void
-  print(DescriptorIterator const&) { printf("GMMA::DescriptorIterator"); }
+  print(DescriptorIterator) { printf("GMMA::DescriptorIterator"); }
 };
 
+template <class T>
+CUTE_HOST_DEVICE constexpr
+GmmaDescriptor
+raw_pointer_cast(DescriptorIterator const& ptr) {
+  return ptr.desc_;
+}
+
+// Recast a DescriptorIterator Tensor to uint64_t, it's RegType in mma_unpack
+template <class NewT>
+CUTE_HOST_DEVICE constexpr
+DescriptorIterator
+recast_ptr(DescriptorIterator const& iter) {
+  static_assert(is_same<NewT, uint64_t>::value, "Can only cast GmmaDescriptorIterator to uint64_t.");
+  return iter;  // Do nothing, it will still dereference to GmmaDescriptor and decay to uint64_t
+}
+
 // The GMMA Traits below have custom fragment type flags for their smem desc tensors.
 // These flags specialize a MakeTensor customization point to correctly make the fragment that is desired.
 template <GMMA::Major>
 struct smem_desc : DescriptorIterator {};
 
-// Recast a DescriptorIterator Tensor to uint64_t, it's RegType
-template <class TLayout, class NewT>
-CUTE_HOST_DEVICE constexpr
-auto
-recast(Tensor<ViewEngine<DescriptorIterator>,TLayout> const& tensor, type_list<NewT>)
-{
-  static_assert(is_same<NewT, uint64_t>::value, "Can only cast descriptors to uint64_t.");
-  return make_tensor(tensor.data(), Layout<_1,_0>{});
-}
-
 } // end namespace GMMA
 
 // Customization point for creating a GMMA::smem_desc Tensor
 template <GMMA::Major MajorMode>
 struct MakeTensor<GMMA::smem_desc<MajorMode>>
 {
-  template <class Engine, class Layout>
+  template <class TEngine, class TLayout>
   CUTE_HOST_DEVICE constexpr auto
-  operator()(Tensor<Engine,Layout> const& smem_tensor)
+  operator()(Tensor<TEngine,TLayout> const& smem_tensor)
   {
-    static_assert(is_smem<Engine>::value, "Expected SMEM Tensor to construct a GMMA Desc Tensor");
+    static_assert(is_smem<TEngine>::value, "Expected SMEM Tensor to construct a GMMA Desc Tensor");
     return make_tensor(GMMA::DescriptorIterator{GMMA::make_gmma_desc<MajorMode>(tensor<0>(smem_tensor))},
-                       recast<uint128_t const>(smem_tensor).layout());
+                       replace<0>(recast<uint128_t const>(smem_tensor).layout(), Layout<_1,_0>{}));
   }
 };