mlir/test/Conversion/AMDGPUToROCDL/dpp.mlir

   1 // RUN: mlir-opt -convert-amdgpu-to-rocdl=chipset=gfx908 %s | FileCheck %s
   2 // RUN: mlir-opt -convert-amdgpu-to-rocdl=chipset=gfx90a %s | FileCheck %s
   3 // RUN: mlir-opt -convert-amdgpu-to-rocdl=chipset=gfx942 %s | FileCheck %s
   4
   5 func.func @test_dpp(%arg0: i32, %arg1: i32) -> i32 {
   6   // CHECK-LABEL: func @test_dpp
   7   // CHECK: rocdl.update.dpp %arg0, %arg1 with 257, 10, 15, false : i32
   8   // CHECK: return %0 : i32
   9   %0 = amdgpu.dpp %arg0 %arg1 row_shl ( 0x1 : i32 ) { row_mask = 0xa : i32, bound_ctrl = false } : i32
  10     return %0 : i32
  11 }
  12
  13 func.func @quad_dpp(%arg0: i32, %arg1: i32) -> i32 {
  14   // CHECK-LABEL: func @quad_dpp
  15   // CHECK: rocdl.update.dpp %arg0, %arg1 with 145, 1, 1, true : i32
  16   // CHECK: return %0 : i32
  17   %0 = amdgpu.dpp %arg0 %arg1 quad_perm ( [1,0,1,2] ) { row_mask = 0x1 : i32, bank_mask = 0x1 : i32, bound_ctrl = true } : i32
  18     return %0 : i32
  19 }
  20
  21 func.func @wave_shr_dpp(%arg0: i32, %arg1: i32) -> i32 {
  22   // CHECK-LABEL: func @wave_shr_dpp
  23   // CHECK: rocdl.update.dpp %arg0, %arg1 with 312, 10, 1, true : i32
  24   // CHECK: return %0 : i32
  25   %0 = amdgpu.dpp %arg0 %arg1 wave_shr { row_mask = 0xa : i32, bank_mask = 0x1 : i32, bound_ctrl = true } : i32
  26     return %0 : i32
  27 }
  28
  29 func.func @row_half_mirror_update_dpp(%arg0: i32, %arg1: i32) -> i32 {
  30   // CHECK-LABEL: func @row_half_mirror_update_dpp
  31   // CHECK: rocdl.update.dpp %arg0, %arg1 with 321, 15, 1, false : i32
  32   // CHECK: return %0 : i32
  33 %0 = amdgpu.dpp %arg0 %arg1 row_half_mirror { bank_mask = 0x1 : i32 } : i32
  34     return %0 : i32
  35 }
  36
  37 func.func @wave_rol_update_dpp(%arg0: i32, %arg1: i32) -> i32 {
  38   // CHECK-LABEL: func @wave_rol_update_dpp
  39   // CHECK: rocdl.update.dpp %arg0, %arg1 with 308, 10, 1, false : i32
  40   // CHECK: return %0 : i32
  41   %0 = amdgpu.dpp %arg0 %arg1 wave_rol { row_mask = 0xa : i32, bank_mask = 0x1 : i32 } : i32
  42     return %0 : i32
  43 }
  44
  45 func.func @row_bcast_dpp_f32(%arg0: f32, %arg1: f32) -> f32 {
  46   // CHECK-LABEL: func @row_bcast_dpp_f32
  47   // CHECK: rocdl.update.dpp %arg0, %arg1 with 322, 15, 15, true : f32
  48   // CHECK: return %0 : f32
  49   %0 = amdgpu.dpp %arg0 %arg1 row_bcast_15 { bound_ctrl = true } : f32
  50     return %0 : f32
  51 }
  52
  53 func.func @test_dpp_f32(%arg0: f32, %arg1: f32) -> f32 {
  54   // CHECK-LABEL: func @test_dpp_f32
  55   // CHECK: rocdl.update.dpp %arg0, %arg1 with 320, 1, 4, true : f32
  56   // CHECK: return %0 : f32
  57   %0 = amdgpu.dpp %arg0 %arg1 row_mirror { row_mask = 0x1 : i32, bank_mask = 0x4 : i32, bound_ctrl = true } : f32
  58     return %0 : f32
  59 }
  60
  61 func.func @quad_perm_update_dpp_f32(%arg0: f32, %arg1: f32) -> f32 {
  62   // CHECK-LABEL: func @quad_perm_update_dpp_f32
  63   // CHECK: rocdl.update.dpp %arg0, %arg1 with  88, 15, 1, false : f32
  64   // CHECK: return %0 : f32
  65   %0 = amdgpu.dpp %arg0 %arg1 quad_perm ( [0,2,1,1] ) { bank_mask = 0x1 : i32 } : f32
  66     return %0 : f32
  67 }
  68
  69 func.func @quad_perm_dpp(%arg0: i64, %arg1: i64) -> i64 {
  70   // CHECK-LABEL: func @quad_perm_dpp
  71   // CHECK: rocdl.update.dpp %arg0, %arg1 with 88, 15, 15, false : i64
  72   // CHECK: return %0 : i64
  73   %0 = amdgpu.dpp %arg0 %arg1 quad_perm ( [0,2,1,1] ) : i64
  74     return %0 : i64
  75 }
  76
  77 func.func @row_bcast_dpp(%arg0: f64, %arg1: f64) -> f64 {
  78   // CHECK-LABEL: func @row_bcast_dpp
  79   // CHECK: rocdl.update.dpp %arg0, %arg1 with 323, 4, 1, false : f64
  80   // CHECK: return %0 : f64
  81   %0 = amdgpu.dpp %arg0 %arg1 row_bcast_31 { row_mask = 0x4 : i32, bank_mask = 0x1 : i32} : f64
  82     return %0 : f64
  83 }
  84
  85 func.func @test_dpp_f16(%arg0: f16, %arg1: f16) -> f16 {
  86   // CHECK-LABEL:  func @test_dpp_f16
  87   // CHECK: llvm.bitcast %arg1 : f16 to i16
  88   // CHECK: llvm.mlir.undef : vector<2xi16>
  89   // CHECK: llvm.mlir.constant(0 : i32) : i32
  90   // CHECK: llvm.insertelement %0, %1[%2 : i32] : vector<2xi16>
  91   // CHECK: llvm.bitcast %3 : vector<2xi16> to i32
  92   // CHECK: llvm.bitcast %arg0 : f16 to i16
  93   // CHECK: llvm.mlir.undef : vector<2xi16>
  94   // CHECK: llvm.mlir.constant(0 : i32) : i32
  95   // CHECK: llvm.insertelement %5, %6[%7 : i32] : vector<2xi16>
  96   // CHECK: llvm.bitcast %8 : vector<2xi16> to i32
  97   // CHECK: rocdl.update.dpp %9, %4 with 273, 15, 3, false : i32
  98   // CHECK: llvm.trunc %10 : i32 to i16
  99   // CHECK: llvm.bitcast %11 : i16 to f16
 100   // CHECK: return %12 : f16
 101   %0 = amdgpu.dpp %arg0 %arg1 row_shr ( 0x1 : i32 ){ bank_mask = 0x3 : i32 } : f16
 102     return %0 : f16
 103 }
 104
 105 func.func @row_shl_dpp_i16(%arg0: i16, %arg1: i16) -> i16 {
 106   // CHECK-LABEL: func @row_shl_dpp_i16
 107   // CHECK: llvm.mlir.undef : vector<2xi16>
 108   // CHECK: llvm.mlir.constant(0 : i32) : i32
 109   // CHECK: llvm.insertelement %arg1, %0[%1 : i32] : vector<2xi16>
 110   // CHECK: llvm.bitcast %2 : vector<2xi16> to i32
 111   // CHECK: llvm.mlir.undef : vector<2xi16>
 112   // CHECK: llvm.mlir.constant(0 : i32) : i32
 113   // CHECK: llvm.insertelement %arg0, %4[%5 : i32] : vector<2xi16>
 114   // CHECK: llvm.bitcast %6 : vector<2xi16> to i32
 115   // CHECK: rocdl.update.dpp %7, %3 with 298, 10, 1, false : i32
 116   // CHECK: llvm.trunc %8 : i32 to i16
 117   // CHECK: return %9 : i16
 118   %0 = amdgpu.dpp %arg0 %arg1 row_ror ( 0xa : i32 ) { row_mask = 0xa : i32, bank_mask = 0x1 : i32 } : i16
 119     return %0 : i16
 120 }
 121
 122 func.func @row_bcast_update_dpp_f16(%arg0: f16, %arg1: f16) -> f16 {
 123   // CHECK-LABEL: func @row_bcast_update_dpp_f16
 124   // CHECK: llvm.bitcast %arg1 : f16 to i16
 125   // CHECK: llvm.mlir.undef : vector<2xi16>
 126   // CHECK: llvm.mlir.constant(0 : i32) : i32
 127   // CHECK: llvm.insertelement %0, %1[%2 : i32] : vector<2xi16>
 128   // CHECK: llvm.bitcast %arg0 : f16 to i16
 129   // CHECK: llvm.mlir.undef : vector<2xi16>
 130   // CHECK: llvm.mlir.constant(0 : i32) : i32
 131   // CHECK  llvm.insertelement %5, %6[%7 : i32] : vector<2xi16>
 132   // CHECK: llvm.bitcast %8 : vector<2xi16> to i32
 133   // CHECK: rocdl.update.dpp %9, %4 with 322, 15, 15, true : i32
 134   // CHECK: llvm.trunc %10 : i32 to i16
 135   // CHECK: llvm.bitcast %11 : i16 to f16
 136   // CHECK: return %12 : f16
 137   %0 = amdgpu.dpp %arg0 %arg1 row_bcast_15 { bound_ctrl = true } : f16
 138     return %0 : f16
 139 }