permute_pk_int4.hpp Source File

permute_pk_int4.hpp Source File#

Composable Kernel: permute_pk_int4.hpp Source File
// SPDX-License-Identifier: MIT
// Copyright (c), Advanced Micro Devices, Inc. All rights reserved.
#pragma once
 
#include "ck_tile/core/utility/bit_cast.hpp"
namespace ck_tile {
template <typename Tensor>
void permute_vectors_i4x4_b(Tensor& tensor)
{
    auto tensor_row_buf = tensor.data();
    for(size_t idx = 0; idx < tensor.size(); idx += 4)
    {
        int8_t input[8];
 
        for(int k = 0; k < 4; k++)
        {
            int8_t i4x2      = bit_cast<int8_t>(tensor_row_buf[idx + k]);
            input[k * 2 + 0] = (i4x2 >> 4) & 0xf;
            input[k * 2 + 1] = (i4x2 >> 0) & 0xf;
        }
 
        // permute 0x76543210 => 0x75316420
        {
            int8_t hi   = input[2];
            int8_t lo   = input[0];
            int8_t i4x2 = (hi << 4) | lo;
 
            tensor_row_buf[idx + 0] = bit_cast<pk_int4_t>(i4x2);
        }
 
        {
            int8_t hi   = input[6];
            int8_t lo   = input[4];
            int8_t i4x2 = (hi << 4) | lo;
 
            tensor_row_buf[idx + 1] = bit_cast<pk_int4_t>(i4x2);
        }
 
        {
            int8_t hi   = input[3];
            int8_t lo   = input[1];
            int8_t i4x2 = (hi << 4) | lo;
 
            tensor_row_buf[idx + 2] = bit_cast<pk_int4_t>(i4x2);
        }
 
        {
            int8_t hi   = input[7];
            int8_t lo   = input[5];
            int8_t i4x2 = (hi << 4) | lo;
 
            tensor_row_buf[idx + 3] = bit_cast<pk_int4_t>(i4x2);
        }
    }
}
 
} // namespace ck_tile