Skip to content

Commit

Permalink
[NVPTX] Add intrinsics for shfl instructions.
Browse files Browse the repository at this point in the history 
Summary:
Currently clang emits these instructions via inline (volatile) asm in
the CUDA headers.  Switching to intrinsics will let the optimizer reason
across calls to these intrinsics.

Reviewers: tra

Subscribers: llvm-commits, jholewinski

Differential Revision: http://reviews.llvm.org/D21160

llvm-svn: 272298
  • Loading branch information
Justin Lebar committed Jun 9, 2016
1 parent 007edb8 commit ed2c282
Show file tree
Hide file tree
Showing 3 changed files with 176 additions and 1 deletion.
44 changes: 44 additions & 0 deletions llvm/include/llvm/IR/IntrinsicsNVVM.td
View file
Expand Up @@ -3746,3 +3746,47 @@ def int_ptx_read_pm3 : PTXReadSpecialRegisterIntrinsic_r32

def int_ptx_bar_sync : Intrinsic<[], [llvm_i32_ty], [IntrConvergent]>,
GCCBuiltin<"__builtin_ptx_bar_sync">;

//
// SHUFFLE
//

// shfl.down.b32 dest, val, offset, mask_and_clamp
def int_ptx_shfl_down_i32 :
Intrinsic<[llvm_i32_ty], [llvm_i32_ty, llvm_i32_ty, llvm_i32_ty],
[IntrNoMem, IntrConvergent], "llvm.nvvm.shfl.down.i32">,
GCCBuiltin<"__builtin_ptx_shfl_down_i32">;
def int_ptx_shfl_down_f32 :
Intrinsic<[llvm_float_ty], [llvm_float_ty, llvm_i32_ty, llvm_i32_ty],
[IntrNoMem, IntrConvergent], "llvm.nvvm.shfl.down.f32">,
GCCBuiltin<"__builtin_ptx_shfl_down_f32">;

// shfl.up.b32 dest, val, offset, mask_and_clamp
def int_ptx_shfl_up_i32 :
Intrinsic<[llvm_i32_ty], [llvm_i32_ty, llvm_i32_ty, llvm_i32_ty],
[IntrNoMem, IntrConvergent], "llvm.nvvm.shfl.up.i32">,
GCCBuiltin<"__builtin_ptx_shfl_up_i32">;
def int_ptx_shfl_up_f32 :
Intrinsic<[llvm_float_ty], [llvm_float_ty, llvm_i32_ty, llvm_i32_ty],
[IntrNoMem, IntrConvergent], "llvm.nvvm.shfl.up.f32">,
GCCBuiltin<"__builtin_ptx_shfl_up_f32">;

// shfl.bfly.b32 dest, val, offset, mask_and_clamp
def int_ptx_shfl_bfly_i32 :
Intrinsic<[llvm_i32_ty], [llvm_i32_ty, llvm_i32_ty, llvm_i32_ty],
[IntrNoMem, IntrConvergent], "llvm.nvvm.shfl.bfly.i32">,
GCCBuiltin<"__builtin_ptx_shfl_bfly_i32">;
def int_ptx_shfl_bfly_f32 :
Intrinsic<[llvm_float_ty], [llvm_float_ty, llvm_i32_ty, llvm_i32_ty],
[IntrNoMem, IntrConvergent], "llvm.nvvm.shfl.bfly.f32">,
GCCBuiltin<"__builtin_ptx_shfl_bfly_f32">;

// shfl.idx.b32 dest, val, lane, mask_and_clamp
def int_ptx_shfl_idx_i32 :
Intrinsic<[llvm_i32_ty], [llvm_i32_ty, llvm_i32_ty, llvm_i32_ty],
[IntrNoMem, IntrConvergent], "llvm.nvvm.shfl.idx.i32">,
GCCBuiltin<"__builtin_ptx_shfl_idx_i32">;
def int_ptx_shfl_idx_f32 :
Intrinsic<[llvm_float_ty], [llvm_float_ty, llvm_i32_ty, llvm_i32_ty],
[IntrNoMem, IntrConvergent], "llvm.nvvm.shfl.idx.f32">,
GCCBuiltin<"__builtin_ptx_shfl_idx_f32">;
43 changes: 42 additions & 1 deletion llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
View file
Expand Up @@ -30,7 +30,7 @@ def immDouble1 : PatLeaf<(fpimm), [{


//-----------------------------------
// Synchronization Functions
// Synchronization and shuffle functions
//-----------------------------------
let isConvergent = 1 in {
def INT_CUDA_SYNCTHREADS : NVPTXInst<(outs), (ins),
Expand Down Expand Up @@ -64,6 +64,47 @@ def INT_BARRIER0_OR : NVPTXInst<(outs Int32Regs:$dst), (ins Int32Regs:$pred),
!strconcat("selp.u32 \t$dst, 1, 0, %p2; \n\t",
!strconcat("}}", ""))))))),
[(set Int32Regs:$dst, (int_nvvm_barrier0_or Int32Regs:$pred))]>;

// shfl.{up,down,bfly,idx}.b32
multiclass SHFL<NVPTXRegClass regclass, string mode, Intrinsic IntOp> {
// The last two parameters to shfl can be regs or imms. ptxas is smart
// enough to inline constant registers, so strictly speaking we don't need to
// handle immediates here. But it's easy enough, and it makes our ptx more
// readable.
def reg : NVPTXInst<
(outs regclass:$dst),
(ins regclass:$src, Int32Regs:$offset, Int32Regs:$mask),
!strconcat("shfl.", mode, ".b32 $dst, $src, $offset, $mask;"),
[(set regclass:$dst, (IntOp regclass:$src, Int32Regs:$offset, Int32Regs:$mask))]>;

def imm1 : NVPTXInst<
(outs regclass:$dst),
(ins regclass:$src, i32imm:$offset, Int32Regs:$mask),
!strconcat("shfl.", mode, ".b32 $dst, $src, $offset, $mask;"),
[(set regclass:$dst, (IntOp regclass:$src, imm:$offset, Int32Regs:$mask))]>;

def imm2 : NVPTXInst<
(outs regclass:$dst),
(ins regclass:$src, Int32Regs:$offset, i32imm:$mask),
!strconcat("shfl.", mode, ".b32 $dst, $src, $offset, $mask;"),
[(set regclass:$dst, (IntOp regclass:$src, Int32Regs:$offset, imm:$mask))]>;

def imm3 : NVPTXInst<
(outs regclass:$dst),
(ins regclass:$src, i32imm:$offset, i32imm:$mask),
!strconcat("shfl.", mode, ".b32 $dst, $src, $offset, $mask;"),
[(set regclass:$dst, (IntOp regclass:$src, imm:$offset, imm:$mask))]>;
}

defm INT_SHFL_DOWN_I32 : SHFL<Int32Regs, "down", int_ptx_shfl_down_i32>;
defm INT_SHFL_DOWN_F32 : SHFL<Float32Regs, "down", int_ptx_shfl_down_f32>;
defm INT_SHFL_UP_I32 : SHFL<Int32Regs, "up", int_ptx_shfl_up_i32>;
defm INT_SHFL_UP_F32 : SHFL<Float32Regs, "up", int_ptx_shfl_up_f32>;
defm INT_SHFL_BFLY_I32 : SHFL<Int32Regs, "bfly", int_ptx_shfl_bfly_i32>;
defm INT_SHFL_BFLY_F32 : SHFL<Float32Regs, "bfly", int_ptx_shfl_bfly_f32>;
defm INT_SHFL_IDX_I32 : SHFL<Int32Regs, "idx", int_ptx_shfl_idx_i32>;
defm INT_SHFL_IDX_F32 : SHFL<Float32Regs, "idx", int_ptx_shfl_idx_f32>;

} // isConvergent = 1


Expand Down
90 changes: 90 additions & 0 deletions llvm/test/CodeGen/NVPTX/shfl.ll
View file
@@ -0,0 +1,90 @@
; RUN: llc < %s -march=nvptx64 -mcpu=sm_30 -disable-nvptx-favor-non-generic | FileCheck %s

declare i32 @llvm.nvvm.shfl.down.i32(i32, i32, i32)
declare float @llvm.nvvm.shfl.down.f32(float, i32, i32)
declare i32 @llvm.nvvm.shfl.up.i32(i32, i32, i32)
declare float @llvm.nvvm.shfl.up.f32(float, i32, i32)
declare i32 @llvm.nvvm.shfl.bfly.i32(i32, i32, i32)
declare float @llvm.nvvm.shfl.bfly.f32(float, i32, i32)
declare i32 @llvm.nvvm.shfl.idx.i32(i32, i32, i32)
declare float @llvm.nvvm.shfl.idx.f32(float, i32, i32)

; Try all four permutations of register and immediate parameters with
; shfl.down.

; CHECK-LABEL: .func{{.*}}shfl.down1
define i32 @shfl.down1(i32 %in) {
; CHECK: ld.param.u32 [[IN:%r[0-9]+]]
; CHECK: shfl.down.b32 [[OUT:%r[0-9]+]], [[IN]], 1, 2;
; CHECK: st.param.{{.}}32 {{.*}}, [[OUT]]
%val = call i32 @llvm.nvvm.shfl.down.i32(i32 %in, i32 1, i32 2)
ret i32 %val
}

; CHECK-LABEL: .func{{.*}}shfl.down2
define i32 @shfl.down2(i32 %in, i32 %width) {
; CHECK: ld.param.u32 [[IN1:%r[0-9]+]]
; CHECK: ld.param.u32 [[IN2:%r[0-9]+]]
; CHECK: shfl.down.{{.}}32 %r{{[0-9]+}}, [[IN1]], [[IN2]], 3;
%val = call i32 @llvm.nvvm.shfl.down.i32(i32 %in, i32 %width, i32 3)
ret i32 %val
}

; CHECK-LABEL: .func{{.*}}shfl.down3
define i32 @shfl.down3(i32 %in, i32 %mask) {
; CHECK: ld.param.u32 [[IN1:%r[0-9]+]]
; CHECK: ld.param.u32 [[IN2:%r[0-9]+]]
; CHECK: shfl.down.{{.}}32 %r{{[0-9]+}}, [[IN1]], 4, [[IN2]];
%val = call i32 @llvm.nvvm.shfl.down.i32(i32 %in, i32 4, i32 %mask)
ret i32 %val
}

; CHECK-LABEL: .func{{.*}}shfl.down4
define i32 @shfl.down4(i32 %in, i32 %width, i32 %mask) {
; CHECK: ld.param.u32 [[IN1:%r[0-9]+]]
; CHECK: ld.param.u32 [[IN2:%r[0-9]+]]
; CHECK: ld.param.u32 [[IN3:%r[0-9]+]]
; CHECK: shfl.down.{{.}}32 %r{{[0-9]+}}, [[IN1]], [[IN2]], [[IN3]];
%val = call i32 @llvm.nvvm.shfl.down.i32(i32 %in, i32 %width, i32 %mask)
ret i32 %val
}

; Try shfl.down with floating-point params.
; CHECK-LABEL: .func{{.*}}shfl.down.float
define float @shfl.down.float(float %in) {
; CHECK: ld.param.f32 [[IN:%f[0-9]+]]
; CHECK: shfl.down.b32 [[OUT:%f[0-9]+]], [[IN]], 5, 6;
; CHECK: st.param.{{.}}32 {{.*}}, [[OUT]]
%out = call float @llvm.nvvm.shfl.down.f32(float %in, i32 5, i32 6)
ret float %out
}

; Try the rest of the shfl modes. Hopefully they're declared in such a way
; that if shfl.down works correctly, they also work correctly.
define void @shfl.rest(i32 %in_i32, float %in_float, i32* %out_i32, float* %out_float) {
; CHECK: shfl.up.b32 %r{{[0-9]+}}, %r{{[0-9]+}}, 1, 2;
%up_i32 = call i32 @llvm.nvvm.shfl.up.i32(i32 %in_i32, i32 1, i32 2)
store i32 %up_i32, i32* %out_i32

; CHECK: shfl.up.b32 %f{{[0-9]+}}, %f{{[0-9]+}}, 3, 4;
%up_float = call float @llvm.nvvm.shfl.up.f32(float %in_float, i32 3, i32 4)
store float %up_float, float* %out_float

; CHECK: shfl.bfly.b32 %r{{[0-9]+}}, %r{{[0-9]+}}, 5, 6;
%bfly_i32 = call i32 @llvm.nvvm.shfl.bfly.i32(i32 %in_i32, i32 5, i32 6)
store i32 %bfly_i32, i32* %out_i32

; CHECK: shfl.bfly.b32 %f{{[0-9]+}}, %f{{[0-9]+}}, 7, 8;
%bfly_float = call float @llvm.nvvm.shfl.bfly.f32(float %in_float, i32 7, i32 8)
store float %bfly_float, float* %out_float

; CHECK: shfl.idx.b32 %r{{[0-9]+}}, %r{{[0-9]+}}, 9, 10;
%idx_i32 = call i32 @llvm.nvvm.shfl.idx.i32(i32 %in_i32, i32 9, i32 10)
store i32 %idx_i32, i32* %out_i32

; CHECK: shfl.idx.b32 %f{{[0-9]+}}, %f{{[0-9]+}}, 11, 12;
%idx_float = call float @llvm.nvvm.shfl.idx.f32(float %in_float, i32 11, i32 12)
store float %idx_float, float* %out_float

ret void
}

0 comments on commit ed2c282

Please sign in to comment.