182 lines · plain
1// RUN: mlir-opt --xegpu-wg-to-sg-distribute -split-input-file %s | FileCheck %s2 3gpu.module @test_elementwise_ops {4 5 // CHECK-LABEL: unary_ops_sg_layout_only6 gpu.func @unary_ops_sg_layout_only(%a: memref<24x32xf32>) {7 %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<24x32xf32>8 -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>>9 %load_a = xegpu.load_nd %tdesc_a10 : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>>11 -> vector<24x32xf32>12 // CHECK: math.exp {{.*}} : vector<12x8xf32>13 %exp = math.exp %load_a14 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>}15 : vector<24x32xf32>16 // CHECK: arith.negf {{.*}} : vector<12x8xf32>17 %negf = arith.negf %load_a18 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>}19 : vector<24x32xf32>20 gpu.return21 }22 23 // CHECK-LABEL: unary_ops24 gpu.func @unary_ops(%a: memref<24x32xf32>) {25 %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<24x32xf32>26 -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>27 %load_a = xegpu.load_nd %tdesc_a28 : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>29 -> vector<24x32xf32>30 // CHECK: math.exp {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>} : vector<12x8xf32>31 %exp = math.exp %load_a32 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}33 : vector<24x32xf32>34 // CHECK: arith.negf {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>} : vector<12x8xf32>35 %negf = arith.negf %load_a36 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}37 : vector<24x32xf32>38 gpu.return39 }40 41 // CHECK-LABEL: binary_ops42 gpu.func @binary_ops(%a: memref<24x32xf32>, %b: memref<24x32xf32>) {43 %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<24x32xf32>44 -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>45 %tdesc_b = xegpu.create_nd_tdesc %b[0, 0] : memref<24x32xf32>46 -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>47 %load_a = xegpu.load_nd %tdesc_a48 : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>49 -> vector<24x32xf32>50 %load_b = xegpu.load_nd %tdesc_b51 : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>52 -> vector<24x32xf32>53 // CHECK: arith.addf {{.*}}, {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>}54 // CHECK-SAME: : vector<12x8xf32>55 %addf = arith.addf %load_a, %load_b56 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}57 : vector<24x32xf32>58 // CHECK: math.powf {{.*}}, {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>}59 // CHECK-SAME: : vector<12x8xf32>60 %powf = math.powf %load_a, %load_b61 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}62 : vector<24x32xf32>63 gpu.return64 }65 66 // CHECK-LABEL: ternary_ops67 gpu.func @ternary_ops(%a: memref<24x32xf32>, %b: memref<24x32xf32>, %c: memref<24x32xi1>) {68 %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<24x32xf32>69 -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>70 %tdesc_b = xegpu.create_nd_tdesc %b[0, 0] : memref<24x32xf32>71 -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>72 %tdesc_c = xegpu.create_nd_tdesc %c[0, 0] : memref<24x32xi1>73 -> !xegpu.tensor_desc<24x32xi1, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>74 %load_a = xegpu.load_nd %tdesc_a75 : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>76 -> vector<24x32xf32>77 %load_b = xegpu.load_nd %tdesc_b78 : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>79 -> vector<24x32xf32>80 %load_c = xegpu.load_nd %tdesc_c81 : !xegpu.tensor_desc<24x32xi1, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>82 -> vector<24x32xi1>83 // CHECK: arith.select {{.*}}, {{.*}}, {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>}84 // CHECK-SAME: : vector<12x8xi1>, vector<12x8xf32>85 %select = arith.select %load_c, %load_a, %load_b86 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}87 : vector<24x32xi1>, vector<24x32xf32>88 // CHECK: math.fma {{.*}}, {{.*}}, {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>}89 // CHECK-SAME: : vector<12x8xf32>90 %fma = math.fma %load_a, %load_b, %load_a91 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}92 : vector<24x32xf32>93 gpu.return94 }95 96 // CHECK-LABEL: type_conversion_ops97 gpu.func @type_conversion_ops(%a: memref<24x32xf32>, %b: memref<24x32xi32>) {98 %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<24x32xf32>99 -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>100 %tdesc_b = xegpu.create_nd_tdesc %b[0, 0] : memref<24x32xi32>101 -> !xegpu.tensor_desc<24x32xi32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>102 %load_a = xegpu.load_nd %tdesc_a103 : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>104 -> vector<24x32xf32>105 %load_b = xegpu.load_nd %tdesc_b106 : !xegpu.tensor_desc<24x32xi32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>107 -> vector<24x32xi32>108 // CHECK: arith.truncf {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>}109 // CHECK-SAME: : vector<12x8xf32> to vector<12x8xf16>110 %truncf = arith.truncf %load_a111 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}112 : vector<24x32xf32> to vector<24x32xf16>113 // CHECK: arith.bitcast {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>}114 // CHECK-SAME: : vector<12x8xi32> to vector<12x8xf32>115 %bitcast = arith.bitcast %load_b116 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}117 : vector<24x32xi32> to vector<24x32xf32>118 gpu.return119 }120 121 // CHECK-LABEL: comparison_ops122 gpu.func @comparison_ops(%a: memref<24x32xf32>, %b: memref<24x32xf32>, %c: memref<24x32xi32>, %d: memref<24x32xi32>) {123 %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<24x32xf32>124 -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>125 %tdesc_b = xegpu.create_nd_tdesc %b[0, 0] : memref<24x32xf32>126 -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>127 %tdesc_c = xegpu.create_nd_tdesc %c[0, 0] : memref<24x32xi32>128 -> !xegpu.tensor_desc<24x32xi32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>129 %tdesc_d = xegpu.create_nd_tdesc %d[0, 0] : memref<24x32xi32>130 -> !xegpu.tensor_desc<24x32xi32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>131 %load_a = xegpu.load_nd %tdesc_a132 : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>133 -> vector<24x32xf32>134 %load_b = xegpu.load_nd %tdesc_b135 : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>136 -> vector<24x32xf32>137 %load_c = xegpu.load_nd %tdesc_c138 : !xegpu.tensor_desc<24x32xi32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>139 -> vector<24x32xi32>140 %load_d = xegpu.load_nd %tdesc_d141 : !xegpu.tensor_desc<24x32xi32, #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>>142 -> vector<24x32xi32>143 // CHECK: arith.cmpf ult, {{.*}}, {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>}144 // CHECK-SAME: : vector<12x8xf32>145 %cmpf = arith.cmpf ult, %load_a, %load_b146 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}147 : vector<24x32xf32>148 // CHECK: arith.cmpi eq, {{.*}}, {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>}149 // CHECK-SAME: : vector<12x8xi32>150 %cmpi = arith.cmpi eq, %load_c, %load_d151 {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}152 : vector<24x32xi32>153 gpu.return154 }155 156 // 1 to N decomposition of elementwise operations157 // CHECK-LABEL: elementwise_ops_rr_assignment158 gpu.func @elementwise_ops_rr_assignment(%a: memref<24x32xf32>, %b: memref<24x32xf32>) {159 %tdesc_a = xegpu.create_nd_tdesc %a[0, 0] : memref<24x32xf32>160 -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>>161 %tdesc_b = xegpu.create_nd_tdesc %b[0, 0] : memref<24x32xf32>162 -> !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>>163 %load_a = xegpu.load_nd %tdesc_a164 : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>>165 -> vector<24x32xf32>166 %load_b = xegpu.load_nd %tdesc_b167 : !xegpu.tensor_desc<24x32xf32, #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>>168 -> vector<24x32xf32>169 // CHECK-COUNT-12: arith.negf {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 2], lane_data = [1, 1]>} : vector<2x2xf32>170 // CHECK-NOT: arith.negf171 %negf = arith.negf %load_a172 {layout_result_0 = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}173 : vector<24x32xf32>174 // CHECK-COUNT-12: math.powf {{.*}}, {{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [2, 2], lane_data = [1, 1]>} : vector<2x2xf32>175 // CHECK-NOT: math.powf176 %powf = math.powf %load_a, %load_b177 {layout_result_0 = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}178 : vector<24x32xf32>179 gpu.return180 }181}182