Support sign extending in prmt

This commit is contained in:
Andrzej Janik 2024-02-15 15:07:05 +00:00
parent 8d10f756a9
commit 8fef0e4fe7
5 changed files with 91 additions and 40 deletions

View file

@ -1654,14 +1654,17 @@ fn emit_inst_prmt(
) -> Result<(), TranslateError> {
let builder = ctx.builder.get();
let components = [
((control >> 0) & 0b1111) as u32,
((control >> 4) & 0b1111) as u32,
((control >> 8) & 0b1111) as u32,
((control >> 12) & 0b1111) as u32,
((control >> 0) & 0b0111) as u32,
((control >> 4) & 0b0111) as u32,
((control >> 8) & 0b0111) as u32,
((control >> 12) & 0b0111) as u32,
];
let sext_components = [
((control >> 0) & 0b1000) != 0,
((control >> 4) & 0b1000) != 0,
((control >> 8) & 0b1000) != 0,
((control >> 12) & 0b1000) != 0,
];
if components.iter().any(|&c| c > 7) {
return Err(TranslateError::todo());
}
let llvm_i32 = get_llvm_type(ctx, &ast::Type::Scalar(ast::ScalarType::U32))?;
let llvm_vec4_i8 = get_llvm_type(ctx, &ast::Type::Vector(ast::ScalarType::U8, 4))?;
let src1 = ctx.names.value(arg.src1)?;
@ -1674,9 +1677,24 @@ fn emit_inst_prmt(
unsafe { LLVMConstInt(llvm_i32, components[2] as _, 0) },
unsafe { LLVMConstInt(llvm_i32, components[3] as _, 0) },
];
let mask = unsafe { LLVMConstVector(components_llvm.as_mut_ptr(), 4) };
let shuffle_result =
let mask =
unsafe { LLVMConstVector(components_llvm.as_mut_ptr(), components_llvm.len() as u32) };
let mut shuffle_result =
unsafe { LLVMBuildShuffleVector(builder, src1_vector, src2_vector, mask, LLVM_UNNAMED) };
// In sext case I'd prefer to just emit V_PERM_B32 directly and be done with it,
// but V_PERM_B32 can sext only odd-indexed bytes.
let llvm_i8 = get_llvm_type(ctx, &ast::Type::Scalar(ast::ScalarType::U8))?;
let const_7 = unsafe { LLVMConstInt(llvm_i8, 7, 0) };
for (idx, requires_sext) in sext_components.iter().copied().enumerate() {
if !requires_sext {
continue;
}
let idx = unsafe { LLVMConstInt(llvm_i32, idx as u64, 0) };
let scalar = unsafe { LLVMBuildExtractElement(builder, shuffle_result, idx, LLVM_UNNAMED) };
let shift = unsafe { LLVMBuildAShr(builder, scalar, const_7, LLVM_UNNAMED) };
shuffle_result =
unsafe { LLVMBuildInsertElement(builder, shuffle_result, shift, idx, LLVM_UNNAMED) };
}
ctx.names.register_result(arg.dst, |dst_name| unsafe {
LLVMBuildBitCast(builder, shuffle_result, llvm_i32, dst_name)
});

View file

@ -1097,6 +1097,15 @@ InstSetp: ast::Instruction<ast::ParsedArgParams<'input>> = {
// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#half-precision-comparison-instructions-set
InstSet: ast::Instruction<ast::ParsedArgParams<'input>> = {
"set" <cmp_op:SetpCompareOp> <ftz:".ftz"?> ".f16x2" ".f16x2" <arg:Arg3> => {
let data = ast::SetData {
dst_type: ast::ScalarType::F16x2,
src_type: ast::ScalarType::F16x2,
flush_to_zero: ftz.is_some(),
cmp_op: cmp_op,
};
ast::Instruction::Set(data, arg)
},
"set" <cmp_op:SetpCompareOp> <ftz:".ftz"?> ".u32" ".f16x2" <arg:Arg3> => {
let data = ast::SetData {
dst_type: ast::ScalarType::U32,

View file

@ -271,7 +271,7 @@ test_ptx!(const, [0u16], [10u16, 20, 30, 40]);
test_ptx!(cvt_s16_s8, [0x139231C2u32], [0xFFFFFFC2u32]);
test_ptx!(cvt_f64_f32, [0.125f32], [0.125f64]);
test_ptx!(cvt_f32_f16, [0xa1u16], [0x37210000u32]);
test_ptx!(prmt, [0x70c507d6u32, 0x6fbd4b5cu32], [0x6fbdd65cu32]);
test_ptx!(prmt, [0x70c507d6u32, 0x6fbd4b5cu32], [0x6fbdd65cu32, 0x6FFFD600]);
test_ptx!(
prmt_non_immediate,
[0x70c507d6u32, 0x6fbd4b5cu32],

View file

@ -1,40 +1,60 @@
target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7"
target triple = "amdgcn-amd-amdhsa"
define protected amdgpu_kernel void @prmt(ptr addrspace(4) byref(i64) %"23", ptr addrspace(4) byref(i64) %"24") #0 {
"31":
%"8" = alloca i1, align 1, addrspace(5)
store i1 false, ptr addrspace(5) %"8", align 1
%"9" = alloca i1, align 1, addrspace(5)
store i1 false, ptr addrspace(5) %"9", align 1
define protected amdgpu_kernel void @prmt(ptr addrspace(4) byref(i64) %"32", ptr addrspace(4) byref(i64) %"33") #0 {
"44":
%"10" = alloca i1, align 1, addrspace(5)
store i1 false, ptr addrspace(5) %"10", align 1
%"11" = alloca i1, align 1, addrspace(5)
store i1 false, ptr addrspace(5) %"11", align 1
%"4" = alloca i64, align 8, addrspace(5)
%"5" = alloca i64, align 8, addrspace(5)
%"6" = alloca i32, align 4, addrspace(5)
%"7" = alloca i32, align 4, addrspace(5)
%"10" = load i64, ptr addrspace(4) %"23", align 8
store i64 %"10", ptr addrspace(5) %"4", align 8
%"11" = load i64, ptr addrspace(4) %"24", align 8
store i64 %"11", ptr addrspace(5) %"5", align 8
%"13" = load i64, ptr addrspace(5) %"4", align 8
%"25" = inttoptr i64 %"13" to ptr
%"12" = load i32, ptr %"25", align 4
store i32 %"12", ptr addrspace(5) %"6", align 4
%"8" = alloca i32, align 4, addrspace(5)
%"9" = alloca i32, align 4, addrspace(5)
%"12" = load i64, ptr addrspace(4) %"32", align 8
store i64 %"12", ptr addrspace(5) %"4", align 8
%"13" = load i64, ptr addrspace(4) %"33", align 8
store i64 %"13", ptr addrspace(5) %"5", align 8
%"15" = load i64, ptr addrspace(5) %"4", align 8
%"26" = inttoptr i64 %"15" to ptr
%"33" = getelementptr inbounds i8, ptr %"26", i64 4
%"14" = load i32, ptr %"33", align 4
store i32 %"14", ptr addrspace(5) %"7", align 4
%"17" = load i32, ptr addrspace(5) %"6", align 4
%"18" = load i32, ptr addrspace(5) %"7", align 4
%0 = bitcast i32 %"17" to <4 x i8>
%1 = bitcast i32 %"18" to <4 x i8>
%2 = shufflevector <4 x i8> %0, <4 x i8> %1, <4 x i32> <i32 4, i32 0, i32 6, i32 7>
%"27" = bitcast <4 x i8> %2 to i32
store i32 %"27", ptr addrspace(5) %"7", align 4
%"19" = load i64, ptr addrspace(5) %"5", align 8
%"34" = inttoptr i64 %"15" to ptr
%"14" = load i32, ptr %"34", align 4
store i32 %"14", ptr addrspace(5) %"6", align 4
%"17" = load i64, ptr addrspace(5) %"4", align 8
%"35" = inttoptr i64 %"17" to ptr
%"46" = getelementptr inbounds i8, ptr %"35", i64 4
%"16" = load i32, ptr %"46", align 4
store i32 %"16", ptr addrspace(5) %"7", align 4
%"19" = load i32, ptr addrspace(5) %"6", align 4
%"20" = load i32, ptr addrspace(5) %"7", align 4
%"30" = inttoptr i64 %"19" to ptr
store i32 %"20", ptr %"30", align 4
%0 = bitcast i32 %"19" to <4 x i8>
%1 = bitcast i32 %"20" to <4 x i8>
%2 = shufflevector <4 x i8> %0, <4 x i8> %1, <4 x i32> <i32 4, i32 0, i32 6, i32 7>
%"36" = bitcast <4 x i8> %2 to i32
store i32 %"36", ptr addrspace(5) %"8", align 4
%"22" = load i32, ptr addrspace(5) %"6", align 4
%"23" = load i32, ptr addrspace(5) %"7", align 4
%3 = bitcast i32 %"22" to <4 x i8>
%4 = bitcast i32 %"23" to <4 x i8>
%5 = shufflevector <4 x i8> %3, <4 x i8> %4, <4 x i32> <i32 4, i32 0, i32 6, i32 7>
%6 = extractelement <4 x i8> %5, i32 0
%7 = ashr i8 %6, 7
%8 = insertelement <4 x i8> %5, i8 %7, i32 0
%9 = extractelement <4 x i8> %8, i32 2
%10 = ashr i8 %9, 7
%11 = insertelement <4 x i8> %8, i8 %10, i32 2
%"39" = bitcast <4 x i8> %11 to i32
store i32 %"39", ptr addrspace(5) %"9", align 4
%"24" = load i64, ptr addrspace(5) %"5", align 8
%"25" = load i32, ptr addrspace(5) %"8", align 4
%"42" = inttoptr i64 %"24" to ptr
store i32 %"25", ptr %"42", align 4
%"26" = load i64, ptr addrspace(5) %"5", align 8
%"27" = load i32, ptr addrspace(5) %"9", align 4
%"43" = inttoptr i64 %"26" to ptr
%"48" = getelementptr inbounds i8, ptr %"43", i64 4
store i32 %"27", ptr %"48", align 4
ret void
}

View file

@ -11,13 +11,17 @@
.reg .u64 out_addr;
.reg .u32 temp1;
.reg .u32 temp2;
.reg .u32 temp3;
.reg .u32 temp4;
ld.param.u64 in_addr, [input];
ld.param.u64 out_addr, [output];
ld.u32 temp1, [in_addr];
ld.u32 temp2, [in_addr+4];
prmt.b32 temp2, temp1, temp2, 30212;
st.u32 [out_addr], temp2;
prmt.b32 temp3, temp1, temp2, 30212;
prmt.b32 temp4, temp1, temp2, 32268;
st.u32 [out_addr], temp3;
st.u32 [out_addr+4], temp4;
ret;
}