Support sign extending in prmt

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

View file

@ -1654,14 +1654,17 @@ fn emit_inst_prmt(
) -> Result<(), TranslateError> { ) -> Result<(), TranslateError> {
let builder = ctx.builder.get(); let builder = ctx.builder.get();
let components = [ let components = [
((control >> 0) & 0b1111) as u32, ((control >> 0) & 0b0111) as u32,
((control >> 4) & 0b1111) as u32, ((control >> 4) & 0b0111) as u32,
((control >> 8) & 0b1111) as u32, ((control >> 8) & 0b0111) as u32,
((control >> 12) & 0b1111) 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_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 llvm_vec4_i8 = get_llvm_type(ctx, &ast::Type::Vector(ast::ScalarType::U8, 4))?;
let src1 = ctx.names.value(arg.src1)?; 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[2] as _, 0) },
unsafe { LLVMConstInt(llvm_i32, components[3] as _, 0) }, unsafe { LLVMConstInt(llvm_i32, components[3] as _, 0) },
]; ];
let mask = unsafe { LLVMConstVector(components_llvm.as_mut_ptr(), 4) }; let mask =
let shuffle_result = 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) }; 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 { ctx.names.register_result(arg.dst, |dst_name| unsafe {
LLVMBuildBitCast(builder, shuffle_result, llvm_i32, dst_name) 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 // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#half-precision-comparison-instructions-set
InstSet: ast::Instruction<ast::ParsedArgParams<'input>> = { 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> => { "set" <cmp_op:SetpCompareOp> <ftz:".ftz"?> ".u32" ".f16x2" <arg:Arg3> => {
let data = ast::SetData { let data = ast::SetData {
dst_type: ast::ScalarType::U32, 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_s16_s8, [0x139231C2u32], [0xFFFFFFC2u32]);
test_ptx!(cvt_f64_f32, [0.125f32], [0.125f64]); test_ptx!(cvt_f64_f32, [0.125f32], [0.125f64]);
test_ptx!(cvt_f32_f16, [0xa1u16], [0x37210000u32]); test_ptx!(cvt_f32_f16, [0xa1u16], [0x37210000u32]);
test_ptx!(prmt, [0x70c507d6u32, 0x6fbd4b5cu32], [0x6fbdd65cu32]); test_ptx!(prmt, [0x70c507d6u32, 0x6fbd4b5cu32], [0x6fbdd65cu32, 0x6FFFD600]);
test_ptx!( test_ptx!(
prmt_non_immediate, prmt_non_immediate,
[0x70c507d6u32, 0x6fbd4b5cu32], [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 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" target triple = "amdgcn-amd-amdhsa"
define protected amdgpu_kernel void @prmt(ptr addrspace(4) byref(i64) %"23", ptr addrspace(4) byref(i64) %"24") #0 { define protected amdgpu_kernel void @prmt(ptr addrspace(4) byref(i64) %"32", ptr addrspace(4) byref(i64) %"33") #0 {
"31": "44":
%"8" = alloca i1, align 1, addrspace(5) %"10" = alloca i1, align 1, addrspace(5)
store i1 false, ptr addrspace(5) %"8", align 1 store i1 false, ptr addrspace(5) %"10", align 1
%"9" = alloca i1, align 1, addrspace(5) %"11" = alloca i1, align 1, addrspace(5)
store i1 false, ptr addrspace(5) %"9", align 1 store i1 false, ptr addrspace(5) %"11", align 1
%"4" = alloca i64, align 8, addrspace(5) %"4" = alloca i64, align 8, addrspace(5)
%"5" = alloca i64, align 8, addrspace(5) %"5" = alloca i64, align 8, addrspace(5)
%"6" = alloca i32, align 4, addrspace(5) %"6" = alloca i32, align 4, addrspace(5)
%"7" = alloca i32, align 4, addrspace(5) %"7" = alloca i32, align 4, addrspace(5)
%"10" = load i64, ptr addrspace(4) %"23", align 8 %"8" = alloca i32, align 4, addrspace(5)
store i64 %"10", ptr addrspace(5) %"4", align 8 %"9" = alloca i32, align 4, addrspace(5)
%"11" = load i64, ptr addrspace(4) %"24", align 8 %"12" = load i64, ptr addrspace(4) %"32", align 8
store i64 %"11", ptr addrspace(5) %"5", align 8 store i64 %"12", ptr addrspace(5) %"4", align 8
%"13" = load i64, ptr addrspace(5) %"4", align 8 %"13" = load i64, ptr addrspace(4) %"33", align 8
%"25" = inttoptr i64 %"13" to ptr store i64 %"13", ptr addrspace(5) %"5", align 8
%"12" = load i32, ptr %"25", align 4
store i32 %"12", ptr addrspace(5) %"6", align 4
%"15" = load i64, ptr addrspace(5) %"4", align 8 %"15" = load i64, ptr addrspace(5) %"4", align 8
%"26" = inttoptr i64 %"15" to ptr %"34" = inttoptr i64 %"15" to ptr
%"33" = getelementptr inbounds i8, ptr %"26", i64 4 %"14" = load i32, ptr %"34", align 4
%"14" = load i32, ptr %"33", align 4 store i32 %"14", ptr addrspace(5) %"6", align 4
store i32 %"14", ptr addrspace(5) %"7", align 4 %"17" = load i64, ptr addrspace(5) %"4", align 8
%"17" = load i32, ptr addrspace(5) %"6", align 4 %"35" = inttoptr i64 %"17" to ptr
%"18" = load i32, ptr addrspace(5) %"7", align 4 %"46" = getelementptr inbounds i8, ptr %"35", i64 4
%0 = bitcast i32 %"17" to <4 x i8> %"16" = load i32, ptr %"46", align 4
%1 = bitcast i32 %"18" to <4 x i8> store i32 %"16", ptr addrspace(5) %"7", align 4
%2 = shufflevector <4 x i8> %0, <4 x i8> %1, <4 x i32> <i32 4, i32 0, i32 6, i32 7> %"19" = load i32, ptr addrspace(5) %"6", align 4
%"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
%"20" = load i32, ptr addrspace(5) %"7", align 4 %"20" = load i32, ptr addrspace(5) %"7", align 4
%"30" = inttoptr i64 %"19" to ptr %0 = bitcast i32 %"19" to <4 x i8>
store i32 %"20", ptr %"30", align 4 %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 ret void
} }

View file

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