mlir/test/Conversion/ArithToLLVM/convert-nd-vector-to-llvmir.mlir

   1 // RUN: mlir-opt -pass-pipeline="builtin.module(func.func(convert-arith-to-llvm))" %s -split-input-file | FileCheck %s
   2
   3 // CHECK-LABEL: @vec_bin
   4 func.func @vec_bin(%arg0: vector<2x2x2xf32>) -> vector<2x2x2xf32> {
   5   // CHECK: llvm.mlir.undef : !llvm.array<2 x array<2 x vector<2xf32>>>
   6
   7   // This block appears 2x2 times
   8   // CHECK-NEXT: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<2 x array<2 x vector<2xf32>>>
   9   // CHECK-NEXT: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<2 x array<2 x vector<2xf32>>>
  10   // CHECK-NEXT: llvm.fadd %{{.*}} : vector<2xf32>
  11   // CHECK-NEXT: llvm.insertvalue %{{.*}}[0, 0] : !llvm.array<2 x array<2 x vector<2xf32>>>
  12
  13   // We check the proper indexing of extract/insert in the remaining 3 positions.
  14   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<2 x array<2 x vector<2xf32>>>
  15   // CHECK: llvm.insertvalue %{{.*}}[0, 1] : !llvm.array<2 x array<2 x vector<2xf32>>>
  16   // CHECK: llvm.extractvalue %{{.*}}[1, 0] : !llvm.array<2 x array<2 x vector<2xf32>>>
  17   // CHECK: llvm.insertvalue %{{.*}}[1, 0] : !llvm.array<2 x array<2 x vector<2xf32>>>
  18   // CHECK: llvm.extractvalue %{{.*}}[1, 1] : !llvm.array<2 x array<2 x vector<2xf32>>>
  19   // CHECK: llvm.insertvalue %{{.*}}[1, 1] : !llvm.array<2 x array<2 x vector<2xf32>>>
  20   %0 = arith.addf %arg0, %arg0 : vector<2x2x2xf32>
  21   return %0 : vector<2x2x2xf32>
  22 }
  23
  24 // CHECK-LABEL: @sexti
  25 func.func @sexti_vector(%arg0 : vector<1x2x3xi32>, %arg1 : vector<1x2x3xi64>) {
  26   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xi64>>>
  27   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi32>>>
  28   // CHECK: llvm.sext %{{.*}} : vector<3xi32> to vector<3xi64>
  29   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
  30   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi32>>>
  31   // CHECK: llvm.sext %{{.*}} : vector<3xi32> to vector<3xi64>
  32   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
  33   %0 = arith.extsi %arg0: vector<1x2x3xi32> to vector<1x2x3xi64>
  34   return
  35 }
  36
  37 // CHECK-LABEL: @zexti
  38 func.func @zexti_vector(%arg0 : vector<1x2x3xi32>, %arg1 : vector<1x2x3xi64>) {
  39   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xi64>>>
  40   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi32>>>
  41   // CHECK: llvm.zext %{{.*}} : vector<3xi32> to vector<3xi64>
  42   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
  43   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi32>>>
  44   // CHECK: llvm.zext %{{.*}} : vector<3xi32> to vector<3xi64>
  45   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
  46   %0 = arith.extui %arg0: vector<1x2x3xi32> to vector<1x2x3xi64>
  47   return
  48 }
  49
  50 // CHECK-LABEL: @sitofp
  51 func.func @sitofp_vector(%arg0 : vector<1x2x3xi32>) -> vector<1x2x3xf32> {
  52   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xf32>>>
  53   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi32>>>
  54   // CHECK: llvm.sitofp %{{.*}} : vector<3xi32> to vector<3xf32>
  55   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xf32>>>
  56   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi32>>>
  57   // CHECK: llvm.sitofp %{{.*}} : vector<3xi32> to vector<3xf32>
  58   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xf32>>>
  59   %0 = arith.sitofp %arg0: vector<1x2x3xi32> to vector<1x2x3xf32>
  60   return %0 : vector<1x2x3xf32>
  61 }
  62
  63 // CHECK-LABEL: @uitofp
  64 func.func @uitofp_vector(%arg0 : vector<1x2x3xi32>) -> vector<1x2x3xf32> {
  65   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xf32>>>
  66   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi32>>>
  67   // CHECK: llvm.uitofp %{{.*}} : vector<3xi32> to vector<3xf32>
  68   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xf32>>>
  69   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi32>>>
  70   // CHECK: llvm.uitofp %{{.*}} : vector<3xi32> to vector<3xf32>
  71   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xf32>>>
  72   %0 = arith.uitofp %arg0: vector<1x2x3xi32> to vector<1x2x3xf32>
  73   return %0 : vector<1x2x3xf32>
  74 }
  75
  76 // CHECK-LABEL: @fptosi
  77 func.func @fptosi_vector(%arg0 : vector<1x2x3xf32>) -> vector<1x2x3xi32> {
  78   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xi32>>>
  79   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xf32>>>
  80   // CHECK: llvm.fptosi %{{.*}} : vector<3xf32> to vector<3xi32>
  81   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi32>>>
  82   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xf32>>>
  83   // CHECK: llvm.fptosi %{{.*}} : vector<3xf32> to vector<3xi32>
  84   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi32>>>
  85   %0 = arith.fptosi %arg0: vector<1x2x3xf32> to vector<1x2x3xi32>
  86   return %0 : vector<1x2x3xi32>
  87 }
  88
  89 // CHECK-LABEL: @fptoui
  90 func.func @fptoui_vector(%arg0 : vector<1x2x3xf32>) -> vector<1x2x3xi32> {
  91   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xi32>>>
  92   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xf32>>>
  93   // CHECK: llvm.fptoui %{{.*}} : vector<3xf32> to vector<3xi32>
  94   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi32>>>
  95   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xf32>>>
  96   // CHECK: llvm.fptoui %{{.*}} : vector<3xf32> to vector<3xi32>
  97   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi32>>>
  98   %0 = arith.fptoui %arg0: vector<1x2x3xf32> to vector<1x2x3xi32>
  99   return %0 : vector<1x2x3xi32>
 100 }
 101
 102 // CHECK-LABEL: @fpext
 103 func.func @fpext_vector(%arg0 : vector<1x2x3xf16>) -> vector<1x2x3xf64> {
 104   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xf64>>>
 105   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xf16>>>
 106   // CHECK: llvm.fpext %{{.*}} : vector<3xf16> to vector<3xf64>
 107   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xf64>>>
 108   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xf16>>>
 109   // CHECK: llvm.fpext %{{.*}} : vector<3xf16> to vector<3xf64>
 110   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xf64>>>
 111   %0 = arith.extf %arg0: vector<1x2x3xf16> to vector<1x2x3xf64>
 112   return %0 : vector<1x2x3xf64>
 113 }
 114
 115 // CHECK-LABEL: @fptrunc
 116 func.func @fptrunc_vector(%arg0 : vector<1x2x3xf64>) -> vector<1x2x3xf16> {
 117   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xf16>>>
 118   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xf64>>>
 119   // CHECK: llvm.fptrunc %{{.*}} : vector<3xf64> to vector<3xf16>
 120   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xf16>>>
 121   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xf64>>>
 122   // CHECK: llvm.fptrunc %{{.*}} : vector<3xf64> to vector<3xf16>
 123   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xf16>>>
 124   %0 = arith.truncf %arg0: vector<1x2x3xf64> to vector<1x2x3xf16>
 125   return %0 : vector<1x2x3xf16>
 126 }
 127
 128 // CHECK-LABEL: @trunci
 129 func.func @trunci_vector(%arg0 : vector<1x2x3xi64>) -> vector<1x2x3xi16> {
 130   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xi16>>>
 131   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
 132   // CHECK: llvm.trunc %{{.*}} : vector<3xi64> to vector<3xi16>
 133   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi16>>>
 134   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
 135   // CHECK: llvm.trunc %{{.*}} : vector<3xi64> to vector<3xi16>
 136   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi16>>>
 137   %0 = arith.trunci %arg0: vector<1x2x3xi64> to vector<1x2x3xi16>
 138   return %0 : vector<1x2x3xi16>
 139 }
 140
 141 // CHECK-LABEL: @shl
 142 func.func @shl_vector(%arg0 : vector<1x2x3xi64>) -> vector<1x2x3xi64> {
 143   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xi64>>>
 144   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
 145   // CHECK: llvm.shl %{{.*}}, %{{.*}} : vector<3xi64>
 146   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
 147   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
 148   // CHECK: llvm.shl %{{.*}}, %{{.*}} : vector<3xi64>
 149   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
 150   %c1 = arith.constant dense<1> : vector<1x2x3xi64>
 151   %0 = arith.shli %arg0, %c1 : vector<1x2x3xi64>
 152   return %0 : vector<1x2x3xi64>
 153 }
 154
 155 // CHECK-LABEL: @shrs
 156 func.func @shrs_vector(%arg0 : vector<1x2x3xi64>) -> vector<1x2x3xi64> {
 157   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xi64>>>
 158   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
 159   // CHECK: llvm.ashr %{{.*}}, %{{.*}} : vector<3xi64>
 160   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
 161   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
 162   // CHECK: llvm.ashr %{{.*}}, %{{.*}} : vector<3xi64>
 163   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
 164   %c1 = arith.constant dense<1> : vector<1x2x3xi64>
 165   %0 = arith.shrsi %arg0, %c1 : vector<1x2x3xi64>
 166   return %0 : vector<1x2x3xi64>
 167 }
 168
 169 // CHECK-LABEL: @shru
 170 func.func @shru_vector(%arg0 : vector<1x2x3xi64>) -> vector<1x2x3xi64> {
 171   // CHECK: llvm.mlir.undef : !llvm.array<1 x array<2 x vector<3xi64>>>
 172   // CHECK: llvm.extractvalue %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
 173   // CHECK: llvm.lshr %{{.*}}, %{{.*}} : vector<3xi64>
 174   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi64>>>
 175   // CHECK: llvm.extractvalue %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
 176   // CHECK: llvm.lshr %{{.*}}, %{{.*}} : vector<3xi64>
 177   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0, 1] : !llvm.array<1 x array<2 x vector<3xi64>>>
 178   %c1 = arith.constant dense<1> : vector<1x2x3xi64>
 179   %0 = arith.shrui %arg0, %c1 : vector<1x2x3xi64>
 180   return %0 : vector<1x2x3xi64>
 181 }
 182
 183 // -----
 184
 185 // CHECK-LABEL: @bitcast_2d
 186 func.func @bitcast_2d(%arg0: vector<2x4xf32>) {
 187   // CHECK: llvm.mlir.undef
 188   // CHECK: llvm.extractvalue %{{.*}}[0]
 189   // CHECK: llvm.bitcast %{{.*}} : vector<4xf32> to vector<4xi32>
 190   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[0]
 191   // CHECK: llvm.extractvalue %{{.*}}[1]
 192   // CHECK: llvm.bitcast %{{.*}} : vector<4xf32> to vector<4xi32>
 193   // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[1]
 194   arith.bitcast %arg0 : vector<2x4xf32> to vector<2x4xi32>
 195   return
 196 }
 197
 198 // -----
 199
 200 // CHECK-LABEL: func @select_2d(
 201 func.func @select_2d(%arg0 : vector<4x3xi1>, %arg1 : vector<4x3xi32>, %arg2 : vector<4x3xi32>) {
 202   // CHECK-DAG: %[[ARG0:.*]] = builtin.unrealized_conversion_cast %arg0
 203   // CHECK-DAG: %[[ARG1:.*]] = builtin.unrealized_conversion_cast %arg1
 204   // CHECK-DAG: %[[ARG2:.*]] = builtin.unrealized_conversion_cast %arg2
 205   // CHECK: %[[EXTRACT1:.*]] = llvm.extractvalue %[[ARG0]][0] : !llvm.array<4 x vector<3xi1>>
 206   // CHECK: %[[EXTRACT2:.*]] = llvm.extractvalue %[[ARG1]][0] : !llvm.array<4 x vector<3xi32>>
 207   // CHECK: %[[EXTRACT3:.*]] = llvm.extractvalue %[[ARG2]][0] : !llvm.array<4 x vector<3xi32>>
 208   // CHECK: %[[SELECT:.*]] = llvm.select %[[EXTRACT1]], %[[EXTRACT2]], %[[EXTRACT3]] : vector<3xi1>, vector<3xi32>
 209   // CHECK: %[[INSERT:.*]] = llvm.insertvalue %[[SELECT]], %{{.*}}[0] : !llvm.array<4 x vector<3xi32>>
 210   %0 = arith.select %arg0, %arg1, %arg2 : vector<4x3xi1>, vector<4x3xi32>
 211   func.return
 212 }
 213
 214 // CHECK-LABEL: func @index_cast_2d(
 215 // CHECK-SAME:    %[[ARG0:.*]]: vector<1x2x3xi1>)
 216 func.func @index_cast_2d(%arg0: vector<1x2x3xi1>) {
 217   // CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]]
 218   // CHECK: %[[EXTRACT1:.*]] = llvm.extractvalue %[[SRC]][0, 0] : !llvm.array<1 x array<2 x vector<3xi1>>>
 219   // CHECK: %[[SEXT1:.*]] = llvm.sext %[[EXTRACT1]] : vector<3xi1> to vector<3xi{{.*}}>
 220   // CHECK: %[[INSERT1:.*]] = llvm.insertvalue %[[SEXT1]], %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi{{.*}}>>>
 221   // CHECK: %[[EXTRACT2:.*]] = llvm.extractvalue %[[SRC]][0, 1] : !llvm.array<1 x array<2 x vector<3xi1>>>
 222   // CHECK: %[[SEXT2:.*]] = llvm.sext %[[EXTRACT2]] : vector<3xi1> to vector<3xi{{.*}}>
 223   // CHECK: %[[INSERT2:.*]] = llvm.insertvalue %[[SEXT2]], %[[INSERT1]][0, 1] : !llvm.array<1 x array<2 x vector<3xi{{.*}}>>>
 224   %0 = arith.index_cast %arg0: vector<1x2x3xi1> to vector<1x2x3xindex>
 225   // CHECK: %[[EXTRACT3:.*]] = llvm.extractvalue %[[INSERT2]][0, 0] : !llvm.array<1 x array<2 x vector<3xi{{.*}}>>>
 226   // CHECK: %[[TRUNC1:.*]] = llvm.trunc %[[EXTRACT3]] : vector<3xi{{.*}}> to vector<3xi1>
 227   // CHECK: %[[INSERT3:.*]] = llvm.insertvalue %[[TRUNC1]], %{{.*}}[0, 0] : !llvm.array<1 x array<2 x vector<3xi1>>>
 228   // CHECK: %[[EXTRACT4:.*]] = llvm.extractvalue %[[INSERT2]][0, 1] : !llvm.array<1 x array<2 x vector<3xi{{.*}}>>>
 229   // CHECK: %[[TRUNC2:.*]] = llvm.trunc %[[EXTRACT4]] : vector<3xi{{.*}}> to vector<3xi1>
 230   // CHECK: %[[INSERT4:.*]] = llvm.insertvalue %[[TRUNC2]], %[[INSERT3]][0, 1] : !llvm.array<1 x array<2 x vector<3xi1>>>
 231   %1 = arith.index_cast %0: vector<1x2x3xindex> to vector<1x2x3xi1>
 232   return
 233 }