Fix pk_int4 cast and add pk_int4 dtype in ck tile (#1854)

* Fix pk_int4 cast and add pk_int4 dtype in ck tile

* fixes

* Improvements

* fix typo

[ROCm/composable_kernel commit: 9ee69dd297]
This commit is contained in:
Bartłomiej Kocot
2025-02-04 10:32:07 +01:00
committed by GitHub
parent 2753e26d39
commit 5835ed012d
12 changed files with 406 additions and 73 deletions

View File

@@ -2,3 +2,4 @@ add_subdirectory(image_to_column)
add_subdirectory(gemm)
add_subdirectory(batched_gemm)
add_subdirectory(grouped_gemm)
add_subdirectory(data_type)

View File

@@ -0,0 +1,4 @@
# Currently ck_tile is only built on gfx9
if(GPU_TARGETS MATCHES "gfx9")
add_gtest_executable(test_ck_tile_pk_int4 test_pk_int4.cpp)
endif()

View File

@@ -0,0 +1,65 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
#include "gtest/gtest.h"
#include <hip/hip_runtime.h>
#include "ck_tile/core.hpp"
using ck_tile::bf16_t;
using ck_tile::bf16x2_t;
using ck_tile::fp16x2_t;
using ck_tile::fp32x2_t;
using ck_tile::half_t;
using ck_tile::pk_int4_t;
TEST(PackedInt4, ConvertToFloat)
{
#ifdef CK_TILE_USE_PK4_LAYOUT_SHUFFLE
constexpr float first_input_val = 7.f;
constexpr float second_input_val = -1.f;
#else
constexpr float first_input_val = -1.f;
constexpr float second_input_val = 7.f;
#endif
uint8_t data = 0b11110111; // {-1, 7}
pk_int4_t in = ck_tile::bit_cast<int8_t>(data);
fp32x2_t out = ck_tile::pk_int4_t_to_fp32x2_t(in);
EXPECT_EQ(out.x, first_input_val);
EXPECT_EQ(out.y, second_input_val);
}
TEST(PackedInt4, ConvertToHalf)
{
#ifdef CK_TILE_USE_PK4_LAYOUT_SHUFFLE
const half_t first_input_val = ck_tile::type_convert<half_t>(7.f);
const half_t second_input_val = ck_tile::type_convert<half_t>(-1.f);
#else
const half_t first_input_val = ck_tile::type_convert<half_t>(-1.f);
const half_t second_input_val = ck_tile::type_convert<half_t>(7.f);
#endif
uint8_t data = 0b11110111; // {-1, 7}
pk_int4_t in = ck_tile::bit_cast<int8_t>(data);
fp16x2_t out = ck_tile::pk_int4_t_to_halfx2_t(in);
EXPECT_EQ(out.x, first_input_val);
EXPECT_EQ(out.y, second_input_val);
}
TEST(PackedInt4, ConvertToBHalf)
{
#ifdef CK_TILE_USE_PK4_LAYOUT_SHUFFLE
const bf16_t first_input_val = ck_tile::type_convert<bf16_t>(7.f);
const bf16_t second_input_val = ck_tile::type_convert<bf16_t>(-1.f);
#else
const bf16_t first_input_val = ck_tile::type_convert<bf16_t>(-1.f);
const bf16_t second_input_val = ck_tile::type_convert<bf16_t>(7.f);
#endif
uint8_t data = 0b11110111; // {-1, 7}
pk_int4_t in = ck_tile::bit_cast<int8_t>(data);
bf16x2_t out = ck_tile::pk_int4_t_to_bfloat16x2_t(in);
EXPECT_EQ(out.x, first_input_val);
EXPECT_EQ(out.y, second_input_val);
}

View File

@@ -50,3 +50,4 @@ endif()
add_gtest_executable(test_type_convert_const type_convert_const.cpp)
add_gtest_executable(test_bhalf test_bhalf.cpp)
add_gtest_executable(test_pk_i4 test_pk_i4.cpp)

View File

@@ -0,0 +1,77 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
#include <bitset>
#include <cinttypes>
#include <cstdint>
#include <iomanip>
#include "gtest/gtest.h"
#include <hip/hip_runtime.h>
#include "ck/host_utility/hip_check_error.hpp"
#include "ck/utility/data_type.hpp"
#include "ck/utility/math_v2.hpp"
#include "ck/utility/get_id.hpp"
#include "ck/library/utility/device_memory.hpp"
#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp"
using ck::bhalf2_t;
using ck::bhalf_t;
using ck::float2_t;
using ck::half2_t;
using ck::half4_t;
using ck::half_t;
using ck::pk_i4_t;
using ck::pk_i4x4_t;
TEST(PackedInt4, ConvertToFloat)
{
#ifdef CK_USE_PK4_LAYOUT_SHUFFLE
constexpr float first_input_val = 7.f;
constexpr float second_input_val = -1.f;
#else
constexpr float first_input_val = -1.f;
constexpr float second_input_val = 7.f;
#endif
uint8_t data = 0b11110111; // {-1, 7}
pk_i4_t in = ck::bit_cast<int8_t>(data);
float2_t out = ck::type_convert<float2_t>(in);
EXPECT_EQ(out.x, first_input_val);
EXPECT_EQ(out.y, second_input_val);
}
TEST(PackedInt4, ConvertToHalf)
{
#ifdef CK_USE_PK4_LAYOUT_SHUFFLE
constexpr half_t first_input_val = ck::type_convert<half_t>(7.f);
constexpr half_t second_input_val = ck::type_convert<half_t>(-1.f);
#else
constexpr half_t first_input_val = ck::type_convert<half_t>(-1.f);
constexpr half_t second_input_val = ck::type_convert<half_t>(7.f);
#endif
uint8_t data = 0b11110111; // {-1, 7}
pk_i4_t in = ck::bit_cast<int8_t>(data);
half2_t out = ck::type_convert<half2_t>(in);
EXPECT_EQ(out.x, first_input_val);
EXPECT_EQ(out.y, second_input_val);
}
TEST(PackedInt4, ConvertToBHalf)
{
#ifdef CK_USE_PK4_LAYOUT_SHUFFLE
const bhalf_t first_input_val = ck::type_convert<bhalf_t>(7.f);
const bhalf_t second_input_val = ck::type_convert<bhalf_t>(-1.f);
#else
const bhalf_t first_input_val = ck::type_convert<bhalf_t>(-1.f);
const bhalf_t second_input_val = ck::type_convert<bhalf_t>(7.f);
#endif
uint8_t data = 0b11110111; // {-1, 7}
pk_i4_t in = ck::bit_cast<int8_t>(data);
bhalf2_t out = ck::type_convert<bhalf2_t>(in);
EXPECT_EQ(out.x, first_input_val);
EXPECT_EQ(out.y, second_input_val);
}