diff --git a/include/cutlass/numeric_types.h b/include/cutlass/numeric_types.h index b44264bb..3891eaee 100644 --- a/include/cutlass/numeric_types.h +++ b/include/cutlass/numeric_types.h @@ -34,8 +34,17 @@ */ #pragma once -#include "cutlass/numeric_size.h" +#include "cute/util/type_traits.hpp" +#include "cutlass/numeric_size.h" +#include "cutlass/integer_subbyte.h" +#include "cutlass/half.h" +#include "cutlass/bfloat16.h" +#include "cutlass/tfloat32.h" +#include "cutlass/float8.h" +#include "cutlass/uint128.h" +#include "cutlass/exmy_base.h" +#include "cutlass/float_subbyte.h" ///////////////////////////////////////////////////////////////////////////////////////////////// namespace cutlass { @@ -56,6 +65,30 @@ struct index_sequence_helper<0, 0, Next...> { template using make_index_sequence = typename index_sequence_helper::type; +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Default case - no negative zero +template +struct has_negative_zero : CUTE_STL_NAMESPACE::false_type{}; + +// Float types that support negative zero +template <> struct has_negative_zero> : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero> : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero> : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero> : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero : CUTE_STL_NAMESPACE::true_type{}; +template <> struct has_negative_zero : CUTE_STL_NAMESPACE::true_type{}; + +// Helper variable template +template +inline constexpr bool has_negative_zero_v = has_negative_zero::value; ///////////////////////////////////////////////////////////////////////////////////////////////// // @@ -75,14 +108,3 @@ struct get_unpacked_element_type { } // namespace cutlass ///////////////////////////////////////////////////////////////////////////////////////////////// - -#include "cutlass/integer_subbyte.h" -#include "cutlass/half.h" -#include "cutlass/bfloat16.h" -#include "cutlass/tfloat32.h" -#include "cutlass/float8.h" -#include "cutlass/uint128.h" -#include "cutlass/exmy_base.h" -#include "cutlass/float_subbyte.h" -///////////////////////////////////////////////////////////////////////////////////////////////// - diff --git a/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp b/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp index 59f4ba1f..da87b910 100644 --- a/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp +++ b/include/cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp @@ -46,6 +46,7 @@ #include "cutlass/fast_math.h" // cutlass::ceil_div, cutlass::round_up #include "cutlass/kernel_hardware_info.h" // cutlass::KernelHardwareInfo #include "cutlass/numeric_size.h" // cutlass::bits_to_bytes +#include "cutlass/numeric_types.h" // cutlass::has_negative_zero_v #include "cutlass/cuda_host_adapter.hpp" // cutlass::CudaHostAdapter namespace cutlass::transform::kernel { @@ -376,6 +377,15 @@ private: copy_vec_pred(cAgA, cAsA, threadIdx_X, GemmM_within_Cta, GemmK_within_Cta); } + // Construct a sign bit mask for handling negative zeros + ElementAMmaRawUnit sign_mask = ElementAMmaRawUnit{ 0 }; + if constexpr (has_negative_zero_v) { + ElementAMmaRawUnit one_sign_mask = static_cast(~(ElementAMmaRawUnit{ 1 } << (cute::sizeof_bits_v - 1))); + for (int i = 0; i < sizeof(ElementAMmaRawUnit) / sizeof(ElementAUint); ++i) { + sign_mask = static_cast((int32_t)sign_mask | (int32_t)one_sign_mask << (i * cute::sizeof_bits_v)); + } + } + // * Compress // cACsAC is always row major order // TensorEAtomM threads perform the compression, each thread compress one row @@ -401,7 +411,14 @@ private: CUTE_UNROLL for (int elt_log_idx = 0; elt_log_idx < OneChunkSizeA{}; ++elt_log_idx) { ElementAMmaRawUnit elem_A = tAsA[elt_log_idx]; - if ( elem_A != ElementAMmaRawUnit{0} ) { + + // Handle negative 0 + ElementAMmaRawUnit masked_elem_A = elem_A; + if constexpr (has_negative_zero_v) { + masked_elem_A = elem_A & sign_mask; + } + + if (masked_elem_A != ElementAMmaRawUnit{0}) { non_zero_elt_log_idx[non_zero_cnt] = elt_log_idx; tACsAC[non_zero_cnt] = elem_A; non_zero_cnt++; @@ -489,6 +506,7 @@ private: constexpr bool IsRowMajor = cute::is_same_v; using Element = typename TensorSrc::element_type; + CUTE_STATIC_ASSERT(cute::is_static_v, "shape(dSrc) needs to be static"); CUTE_STATIC_ASSERT(cute::is_static_v, "shape(dDst) needs to be static"); CUTE_STATIC_ASSERT(cute::sizeof_bits_v == cute::sizeof_bits_v, @@ -530,7 +548,7 @@ private: for (int iter_col_thr = 0; iter_col_thr < ValueShapeCols; ++iter_col_thr) { const int row_i = (iter_row_blk * ThreadShapeRows + threadIdx_X_row) * ValueShapeRows + iter_row_thr; const int col_i = (col_chunk_i * ThreadShapeCols + threadIdx_X_col) * ValueShapeCols + iter_col_thr; - if constexpr ( (not pred) + if constexpr ( (not pred) ) { dDst(row_i, col_i) = dSrc(row_i, col_i); } diff --git a/include/cutlass/transform/kernel/sparse_gemm_compressor.hpp b/include/cutlass/transform/kernel/sparse_gemm_compressor.hpp index d2679b06..18f8d241 100644 --- a/include/cutlass/transform/kernel/sparse_gemm_compressor.hpp +++ b/include/cutlass/transform/kernel/sparse_gemm_compressor.hpp @@ -35,16 +35,17 @@ #pragma once -#include // std::fill -#include // std::array -#include // std::mt19937 +#include // std::fill +#include // std::array +#include // std::mt19937 -#include "cute/numeric/numeric_types.hpp" // cute::sizeof_bits_v -#include "cute/tensor.hpp" // cute::Tensor, cute::make_tensor -#include "cutlass/arch/arch.h" // cutlass::arch::SmXY -#include "cutlass/gemm/gemm.h" // cutlass::TagToStrideA_t -#include "cutlass/fast_math.h" // cutlass::ceil_div, cutlass::round_up -#include "cutlass/numeric_size.h" // cutlass::bits_to_bytes +#include "cute/numeric/numeric_types.hpp" // cute::sizeof_bits_v +#include "cute/tensor.hpp" // cute::Tensor, cute::make_tensor +#include "cutlass/arch/arch.h" // cutlass::arch::SmXY +#include "cutlass/detail/dependent_false.hpp" // cutlass::detail::dependent_false +#include "cutlass/gemm/gemm.h" // cutlass::TagToStrideA_t +#include "cutlass/fast_math.h" // cutlass::ceil_div, cutlass::round_up +#include "cutlass/numeric_size.h" // cutlass::bits_to_bytes #include "cutlass/transform/kernel/sm90_sparse_gemm_compressor.hpp" diff --git a/test/unit/transform/CMakeLists.txt b/test/unit/transform/CMakeLists.txt index 4cc80945..f00275b3 100644 --- a/test/unit/transform/CMakeLists.txt +++ b/test/unit/transform/CMakeLists.txt @@ -27,6 +27,7 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. add_subdirectory(threadblock) +add_subdirectory(device) add_subdirectory(kernel) add_custom_target( diff --git a/test/unit/transform/device/sm90_sparse_gemm_compressor_legacy.hpp b/test/unit/transform/device/sm90_sparse_gemm_compressor_legacy.hpp index 29280655..c4e7de43 100644 --- a/test/unit/transform/device/sm90_sparse_gemm_compressor_legacy.hpp +++ b/test/unit/transform/device/sm90_sparse_gemm_compressor_legacy.hpp @@ -155,7 +155,9 @@ namespace detail { int offset = chunk_idx * LogicalElemsAPerChunk + subchunk_idx * ElemsARawPerElementAMmaRaw + elem_idx; subchunk_elems[elem_idx] = offset < effective_elems ? tensorA(offset) : ElementA(0); - if (subchunk_elems[elem_idx] != ElementA(0)) { + ElementA zero = static_cast(0); + ElementA minus_zero = static_cast(ElementA(1) << cutlass::sizeof_bits_v - 1); + if (subchunk_elems[elem_idx] != zero && subchunk_elems[elem_idx] != minus_zero) { if (non_zero_cnt >= PhysicalSubChunk) { #ifdef __CUDA_ARCH__ asm volatile ("brkpt;\n" ::);