155 lines · plain
1// RUN: mlir-opt %s -split-input-file -verify-diagnostics | FileCheck %s2 3// CHECK-LABEL: func.func @blockload2d(4// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>,5// CHECK-SAME: %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32, %[[ARG3:.*]]: i32, %[[ARG4:.*]]: i32, %[[ARG5:.*]]: i32)6func.func @blockload2d(%a: !llvm.ptr<1>, %base_width_a: i32, %base_height_a: i32,7 %base_pitch_a: i32, %x: i32, %y: i32) -> vector<8xi16> {8 // CHECK: %[[VAR0:.*]] = xevm.blockload2d %[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]]9 // CHECK-DAG: elem_size_in_bits = 16 : i3210 // CHECK-DAG: tile_width = 16 : i3211 // CHECK-DAG: tile_height = 8 : i3212 // CHECK-DAG: v_blocks = 1 : i3213 // CHECK-DAG: transpose = false14 // CHECK-DAG: pack_register = false15 // CHECK-DAG: cache_control = #xevm.load_cache_control<L1uc_L2uc_L3uc>16 // CHECK: (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi16>17 %loaded_a = xevm.blockload2d %a, %base_width_a, %base_height_a, %base_pitch_a, %x, %y18 <{elem_size_in_bits=16 : i32, tile_width=16 : i32, tile_height=8 : i32, v_blocks=1 : i32,19 transpose=false, pack_register=false, cache_control=#xevm.load_cache_control<L1uc_L2uc_L3uc>}>20 : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi16>21 return %loaded_a : vector<8xi16>22}23 24// -----25// CHECK-LABEL: func.func @blockstore2d(26// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>,27// CHECK-SAME: %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32, %[[ARG3:.*]]: i32, %[[ARG4:.*]]: i32, %[[ARG5:.*]]: i32,28// CHECK-SAME: %[[ARG6:.*]]: vector<8xi32>)29func.func @blockstore2d(%c: !llvm.ptr<1>, %base_width_c: i32, %base_height_c: i32,30 %base_pitch_c: i32, %x: i32, %y: i32, %c_result_casted: vector<8xi32>) {31 // CHECK: xevm.blockstore2d %[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]], %[[ARG6]]32 // CHECK-DAG: elem_size_in_bits = 32 : i3233 // CHECK-DAG: tile_width = 16 : i3234 // CHECK-DAG: tile_height = 8 : i3235 // CHECK: (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi32>)36 xevm.blockstore2d %c, %base_width_c, %base_height_c, %base_pitch_c, %x, %y, %c_result_casted37 <{elem_size_in_bits=32 : i32, tile_width=16 : i32, tile_height=8 : i32}>38 : (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi32>)39 return40}41 42// -----43// CHECK-LABEL: func.func @blockprefetch2d(44// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>,45// CHECK-SAME: %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32, %[[ARG3:.*]]: i32, %[[ARG4:.*]]: i32, %[[ARG5:.*]]: i32)46func.func @blockprefetch2d(%ptr: !llvm.ptr<1>, %base_width: i32, %base_height: i32,47 %base_pitch: i32, %x: i32, %y: i32) {48 // CHECK: xevm.blockprefetch2d %[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]]49 // CHECK-DAG: elem_size_in_bits = 8 : i3250 // CHECK-DAG: tile_width = 32 : i3251 // CHECK-DAG: tile_height = 8 : i3252 // CHECK-DAG: v_blocks = 1 : i3253 // CHECK-DAG: cache_control = #xevm.load_cache_control<L1uc_L2uc_L3uc>54 // CHECK: (!llvm.ptr<1>, i32, i32, i32, i32, i32)55 xevm.blockprefetch2d %ptr, %base_width, %base_height, %base_pitch, %x, %y <{elem_size_in_bits=8 : i32,56 tile_width=32 : i32, tile_height=8 : i32, v_blocks=1 : i32,57 cache_control=#xevm.load_cache_control<L1uc_L2uc_L3uc>}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32)58 return59}60 61// -----62// CHECK-LABEL: func.func @blockload(63// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>)64func.func @blockload(%ptr: !llvm.ptr<1>) -> vector<4xi16> {65 // CHECK: %[[VAR0:.*]] = xevm.blockload %[[ARG0]]66 // CHECK-SAME: cache_control = #xevm.load_cache_control<L1uc_L2uc_L3uc>67 // CHECK-SAME: (!llvm.ptr<1>) -> vector<4xi16>68 %loaded = xevm.blockload %ptr <{cache_control=#xevm.load_cache_control<L1uc_L2uc_L3uc>}>69 : (!llvm.ptr<1>) -> vector<4xi16>70 return %loaded : vector<4xi16>71}72 73// -----74// CHECK-LABEL: func.func @blockstore(75// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>,76// CHECK-SAME: %[[ARG1:.*]]: vector<4xi32>)77func.func @blockstore(%ptr: !llvm.ptr<1>, %value: vector<4xi32>) {78 // CHECK: xevm.blockstore %[[ARG0]], %[[ARG1]]79 // CHECK-SAME: (!llvm.ptr<1>, vector<4xi32>)80 xevm.blockstore %ptr, %value : (!llvm.ptr<1>, vector<4xi32>)81 return82}83 84// -----85// CHECK-LABEL: func.func @mma(86// CHECK-SAME: %[[ARG0:.*]]: vector<8xf32>, %[[ARG1:.*]]: vector<8xi16>, %[[ARG2:.*]]: vector<8xi32>)87func.func @mma(%loaded_c_casted: vector<8xf32>, %loaded_a: vector<8xi16>, %loaded_b_casted: vector<8xi32>) -> vector<8xf32> {88 // CHECK: %0 = xevm.mma %[[ARG1]], %[[ARG2]], %[[ARG0]] {shape = <m = 8, n = 16, k = 16>,89 // CHECK-SAME: types = <d = f32, a = f16, b = f16, c = f32>} : (vector<8xi16>, vector<8xi32>, vector<8xf32>) -> vector<8xf32>90 %c_result = xevm.mma %loaded_a, %loaded_b_casted, %loaded_c_casted { shape=<m=8, n=16, k=16>,91 types=<d=f32, a=f16, b=f16, c=f32> } : (vector<8xi16>, vector<8xi32>, vector<8xf32>) -> vector<8xf32>92 return %c_result : vector<8xf32>93}94 95// -----96// CHECK-LABEL: func.func @memfence()97func.func @memfence() {98 // CHECK: xevm.memfence99 // CHECK-DAG: addrspace = #xevm.addr_space<global>100 // CHECK-DAG: scope = #xevm.mem_scope<workgroup>101 xevm.memfence <{addrspace=#xevm.addr_space<global>, scope=#xevm.mem_scope<workgroup>}>102 return103}104 105// -----106// CHECK-LABEL: func.func @prefetch(107// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>)108func.func @prefetch(%ptr: !llvm.ptr<1>) {109 // CHECK: xevm.prefetch %[[ARG0]]110 // CHECK-SAME: <{cache_control = #xevm.load_cache_control<L1uc_L2uc_L3uc>}> : (!llvm.ptr<1>)111 xevm.prefetch %ptr <{cache_control = #xevm.load_cache_control<L1uc_L2uc_L3uc>}> : (!llvm.ptr<1>)112 return113}114 115// -----116// CHECK-LABEL: @xevm_module [#xevm.target<O = 3, chip = "pvc">] {117gpu.module @xevm_module [#xevm.target<O = 3, chip = "pvc">]{118}119 120// -----121// CHECK-LABEL: @xevm_special_ids122llvm.func @xevm_special_ids() -> i32 {123 // CHECK: xevm.local_id.x : i32124 %1 = xevm.local_id.x : i32125 // CHECK: xevm.local_id.y : i32126 %2 = xevm.local_id.y : i32127 // CHECK: xevm.local_id.z : i32128 %3 = xevm.local_id.z : i32129 // CHECK: xevm.local_size.x : i32130 %4 = xevm.local_size.x : i32131 // CHECK: xevm.local_size.y : i32132 %5 = xevm.local_size.y : i32133 // CHECK: xevm.local_size.z : i32134 %6 = xevm.local_size.z : i32135 // CHECK: xevm.group_id.x : i32136 %7 = xevm.group_id.x : i32137 // CHECK: xevm.group_id.y : i32138 %8 = xevm.group_id.y : i32139 // CHECK: xevm.group_id.z : i32140 %9 = xevm.group_id.z : i32141 // CHECK: xevm.group_count.x : i32142 %10 = xevm.group_count.x : i32143 // CHECK: xevm.group_count.y : i32144 %11 = xevm.group_count.y : i32145 // CHECK: xevm.group_count.z : i32146 %12 = xevm.group_count.z : i32147 // CHECK: xevm.lane_id : i32148 %14 = xevm.lane_id : i32149 // CHECK: xevm.subgroup_size : i32150 %39 = xevm.subgroup_size : i32151 // CHECK: xevm.subgroup_id : i32152 %40 = xevm.subgroup_id : i32153 llvm.return %1 : i32154}155