Merge branch 'develop' into users/jlichtne/ALMIOPEN-1994-bump-rock-hash

Author: JonathanLichtnerAMD

projects/composablekernel/CMakeLists.txt

@@ -741,6 +741,10 @@ SET(BUILD_DEV ON CACHE BOOL "BUILD_DEV")
 if(BUILD_DEV)
     add_compile_options(-Werror)
     add_compile_options(-Weverything)
+    add_compile_options(-Wno-lifetime-safety-intra-tu-suggestions)
+    add_compile_options(-Wno-lifetime-safety-cross-tu-suggestions)
+    add_compile_options(-Wno-lifetime-safety-lifetimebound-violation)
+    add_compile_options(-Wno-unknown-warning-option)
 endif()
 message(STATUS "CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}")
 

projects/composablekernel/cmake/EnableCompilerWarnings.cmake

@@ -50,6 +50,10 @@ else()
             -Wsign-compare
             -Wno-extra-semi-stmt
             -Wno-unused-template
+            -Wno-lifetime-safety-intra-tu-suggestions
+            -Wno-lifetime-safety-cross-tu-suggestions
+            -Wno-lifetime-safety-lifetimebound-violation
+            -Wno-unknown-warning-option
         )
         if (CMAKE_${COMPILER}_COMPILER_ID MATCHES "Clang")
             list(APPEND CMAKE_COMPILER_WARNINGS
@@ -76,6 +80,10 @@ else()
                 -Wno-unsafe-buffer-usage
                 -Wno-unused-lambda-capture
                 -Wno-nvcc-compat
+                -Wno-lifetime-safety-intra-tu-suggestions
+                -Wno-lifetime-safety-cross-tu-suggestions
+                -Wno-lifetime-safety-lifetimebound-violation
+                -Wno-unknown-warning-option
             )
             if(CK_CXX_STANDARD GREATER_EQUAL 20)
                 list(APPEND CMAKE_COMPILER_WARNINGS -Wno-c++20-compat)

projects/composablekernel/dispatcher/CMakeLists.txt

@@ -59,7 +59,7 @@ endif()
 # Compiler warnings
 if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang")
     target_compile_options(ck_tile_dispatcher PRIVATE
-        -Wall -Wextra -Wpedantic
+        -Wall -Wextra -Wpedantic -Wno-lifetime-safety-intra-tu-suggestions -Wno-lifetime-safety-cross-tu-suggestions -Wno-lifetime-safety-lifetimebound-violation -Wno-unknown-warning-option
     )
 elseif(CMAKE_CXX_COMPILER_ID MATCHES "MSVC")
     target_compile_options(ck_tile_dispatcher PRIVATE

projects/composablekernel/example/ck_tile/01_fmha/CMakeLists.txt

@@ -2,8 +2,8 @@
 # SPDX-License-Identifier: MIT
 
 set(INST_TARGETS ${SUPPORTED_GPU_TARGETS})
-# Currently only gfx9 and gfx12 archs are supported by FMHA
-list(FILTER INST_TARGETS INCLUDE REGEX "gfx9|gfx12")
+# Currently only gfx9, gfx11, and gfx12 archs are supported by FMHA
+list(FILTER INST_TARGETS INCLUDE REGEX "gfx9|gfx1[12]")
 if(NOT INST_TARGETS)
   message(WARNING "Skipping Tile Engine FMHA compilation: No supported GPU targets (gfx9, gfx11, gfx12) found in SUPPORTED_GPU_TARGETS: ${SUPPORTED_GPU_TARGETS}")
   return()

projects/composablekernel/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py

@@ -28,6 +28,13 @@
 FMHA_BWD_KERNEL_HEADER = """// SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.\n
 // auto generated by generate.py
+#if defined(__HIP_DEVICE_COMPILE__) &&                                                       \\
+    (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) ||                  \\
+     defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) ||                  \\
+     defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__))
+#undef CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT
+#define CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT 1
+#endif
 #include "fmha_bwd.hpp"
 
 """

projects/composablekernel/include/ck_tile/core/config.hpp

@@ -173,13 +173,15 @@
 #endif
 
 // buffer atomic add: floating point
+#ifndef CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT
 #ifndef __HIP_DEVICE_COMPILE__ // for host code
 #define CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT 1
 #elif defined(__gfx9__) || defined(__gfx12__) // for GPU code
 #define CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT 1
 #else // for GPU code
 #define CK_TILE_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT 0
 #endif
+#endif
 
 #if(defined(__gfx90a__) || defined(__gfx94__)) // for GPU code
 #define CK_TILE_USE_AMD_BUFFER_ATOMIC_MAX_FLOAT64 1

projects/composablekernel/include/ck_tile/core/numeric/mxfp_scale.hpp

@@ -103,7 +103,101 @@ struct Packed4Scale
     }
 };
 
+template <typename ScaleType>
+struct Packed8Scale
+{
+    using scale_type     = ScaleType;
+    using raw_type       = uint64_t;
+    using raw_scale_type = typename ScaleType::raw_type;
+
+    static constexpr int num_pack = 8;
+    union
+    {
+        raw_type data_;
+        raw_scale_type scales_[num_pack]; // Direct byte/element access
+    };
+
+    // Constructors
+    CK_TILE_HOST_DEVICE constexpr Packed8Scale() = default;
+    CK_TILE_HOST_DEVICE constexpr Packed8Scale(raw_type val) : data_(val) {}
+    CK_TILE_HOST_DEVICE constexpr Packed8Scale(
+        float s0, float s1, float s2, float s3, float s4, float s5, float s6, float s7)
+    {
+        set_scales_from_float(s0, s1, s2, s3, s4, s5, s6, s7);
+    }
+
+    CK_TILE_HOST_DEVICE constexpr Packed8Scale(ScaleType s0,
+                                               ScaleType s1,
+                                               ScaleType s2,
+                                               ScaleType s3,
+                                               ScaleType s4,
+                                               ScaleType s5,
+                                               ScaleType s6,
+                                               ScaleType s7)
+    {
+        set_scales(s0, s1, s2, s3, s4, s5, s6, s7);
+    }
+
+    CK_TILE_HOST_DEVICE constexpr void set_scales_from_float(
+        float s0, float s1, float s2, float s3, float s4, float s5, float s6, float s7)
+    {
+        set_scales(ScaleType(s0),
+                   ScaleType(s1),
+                   ScaleType(s2),
+                   ScaleType(s3),
+                   ScaleType(s4),
+                   ScaleType(s5),
+                   ScaleType(s6),
+                   ScaleType(s7));
+    }
+
+    CK_TILE_HOST_DEVICE constexpr void set_scales(ScaleType s0,
+                                                  ScaleType s1,
+                                                  ScaleType s2,
+                                                  ScaleType s3,
+                                                  ScaleType s4,
+                                                  ScaleType s5,
+                                                  ScaleType s6,
+                                                  ScaleType s7)
+    {
+        data_ = 0;
+        pack_scale(s0, 7);
+        pack_scale(s1, 6);
+        pack_scale(s2, 5);
+        pack_scale(s3, 4);
+        pack_scale(s4, 3);
+        pack_scale(s5, 2);
+        pack_scale(s6, 1);
+        pack_scale(s7, 0);
+    }
+
+    CK_TILE_HOST_DEVICE constexpr operator raw_type() const { return data_; }
+    CK_TILE_HOST_DEVICE constexpr raw_type& data() [[clang::lifetimebound]] { return data_; }
+    CK_TILE_HOST_DEVICE constexpr raw_type data() const { return data_; }
+
+    CK_TILE_HOST_DEVICE constexpr float unpack_to_float(int i) const
+    {
+        return static_cast<float>(unpack_scale(i));
+    }
+
+    CK_TILE_HOST_DEVICE constexpr ScaleType unpack_scale(int i) const
+    {
+        return ScaleType(scales_[i]);
+    }
+
+    CK_TILE_HOST_DEVICE constexpr void pack_from_float(float scale, int i)
+    {
+        pack_scale(ScaleType(scale), i);
+    }
+
+    CK_TILE_HOST_DEVICE constexpr void pack_scale(ScaleType scale, int i)
+    {
+        scales_[i] = scale.get();
+    }
+};
+
 // Type alias for e8m0_t scales
 using Packed4Scale_E8M0 = Packed4Scale<e8m0_t>;
+using Packed8Scale_E8M0 = Packed8Scale<e8m0_t>;
 
 } // namespace ck_tile

projects/composablekernel/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_kernel.hpp

@@ -363,7 +363,8 @@ struct StreamKKernel
 
             // Determine the total size along the K dimension the workgroup is using in this
             // iteration (used to construct tensor views).
-            index_t k_size = num_loop_sk * TilePartitioner::KPerBlock;
+            index_t k_size = amd_wave_read_first_lane(
+                kargs.tile_partitioner.get_k_size(num_loop_sk, local_iter_end));
 
             // Get the K offsets for the A and B tensors
             auto [i_k_a, i_k_b] = GetKOffsets<ALayout, BLayout>(

projects/composablekernel/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner.hpp

@@ -139,6 +139,19 @@ struct StreamKTilePartitionerBase
     CK_TILE_DEVICE auto
     get_output_tile_index(index_t tile_idx) const noexcept -> tuple<index_t, index_t>;
 
+    /**
+     * @brief Calculates the total size along the K dimension the workgroup is using in this
+     * Stream-K loop iteration
+     *
+     * @param num_macro_tiles  The number of macro tiles along the K dimension this workgroup is
+     * assigned.
+     * @param local_iter_end The workgroup's non-inclusive end iteration that is local to its
+     * current tile.
+     * @return index_t  The K dimension size for the current Stream-K loop iteration.
+     */
+    CK_TILE_DEVICE index_t get_k_size(index_t num_macro_tiles,
+                                      index_t local_iter_end) const noexcept;
+
     /**
      * @brief Calculates the total space needed for the partials and flags buffers.
      *
@@ -208,6 +221,17 @@ struct StreamKTilePartitionerBase
      */
     CK_TILE_HOST_DEVICE index_t get_n() const noexcept;
 
+    /**
+     * @brief Returns the k dimension for the GEMM problem.
+     */
+    CK_TILE_HOST_DEVICE index_t get_k() const noexcept;
+
+    /**
+     * @brief Returns the remainder along the k dimension when k is not evenly divisible by
+     * KPerBlock.
+     */
+    CK_TILE_HOST_DEVICE index_t get_remainder_along_k() const noexcept;
+
     /**
      * @brief Returns an estimate of the number of workgroups writing to the same macro tile in C.
      */
@@ -244,6 +268,8 @@ struct StreamKTilePartitionerBase
     index_t extra_iters_;
     index_t total_dp_iters_;
     index_t n_;
+    index_t k_;
+    index_t remainder_along_k_;
 };
 
 /**

projects/composablekernel/include/ck_tile/ops/gemm/kernel/streamk_gemm/streamk_gemm_tile_partitioner_impl.hpp

@@ -8,10 +8,11 @@ namespace ck_tile {
 template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategyType>
 StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategyType>::StreamKTilePartitionerBase(
     index_t m, index_t n, index_t k, index_t max_active_wgs)
-    : max_active_wgs_{max_active_wgs}, n_{n}
+    : max_active_wgs_{max_active_wgs}, n_{n}, k_{k}
 {
-    iters_per_tile_ = integer_divide_ceil(k, KPerBlock);
-    num_tiles_      = integer_divide_ceil(m, MPerBlock) * integer_divide_ceil(n_, NPerBlock);
+    iters_per_tile_    = integer_divide_ceil(k, KPerBlock);
+    num_tiles_         = integer_divide_ceil(m, MPerBlock) * integer_divide_ceil(n_, NPerBlock);
+    remainder_along_k_ = k % KPerBlock;
 
     bool big_enough         = num_tiles_ > max_active_wgs_;
     index_t remainder_tiles = num_tiles_ % max_active_wgs_;
@@ -250,6 +251,21 @@ StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategyType>::get_n() c
     return n_;
 }
 
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategyType>
+CK_TILE_HOST_DEVICE index_t
+StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategyType>::get_k() const noexcept
+{
+    return k_;
+}
+
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategyType>
+CK_TILE_HOST_DEVICE index_t
+StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategyType>::get_remainder_along_k()
+    const noexcept
+{
+    return remainder_along_k_;
+}
+
 template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategyType>
 CK_TILE_HOST index_t
 StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategyType>::estimate_num_wgs_per_tile()
@@ -334,6 +350,29 @@ StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategyType>::remap_xcd
     return block_1d_id;
 }
 
+template <typename BlockGemmShapeType, StreamKReductionStrategy ReductionStrategyType>
+CK_TILE_DEVICE index_t
+StreamKTilePartitionerBase<BlockGemmShapeType, ReductionStrategyType>::get_k_size(
+    index_t num_macro_tiles, index_t local_iter_end) const noexcept
+{
+    // Determine if this workgroup is responsible for the last macro tile in the K dimension
+    bool last_tile = get_iters_per_tile() == local_iter_end;
+    index_t k_size;
+    // If there is no remainder or if the workgroup was not assigned the last macro tile along K,
+    // then their k_size will be a multiple of KPerBlock.
+    if(!remainder_along_k_ || !last_tile)
+    {
+        k_size = num_macro_tiles * KPerBlock;
+    }
+    // Otherwise, there's a remainder. So, k_size is not a multiple of KPerBlock.
+    else
+    {
+        k_size = (num_macro_tiles - 1) * KPerBlock + remainder_along_k_;
+    }
+
+    return k_size;
+}
+
 template <typename BlockGemmShapeType,
           StreamKReductionStrategy ReductionStrategyType,
           bool Persistent>