mlir/test/Conversion/IndexToSPIRV/index-to-spirv.mlir

   1 // RUN: mlir-opt %s -convert-index-to-spirv | FileCheck %s
   2 // RUN: mlir-opt %s -convert-index-to-spirv=use-64bit-index=false | FileCheck %s --check-prefix=INDEX32
   3 // RUN: mlir-opt %s -convert-index-to-spirv=use-64bit-index=true | FileCheck %s --check-prefix=INDEX64
   4
   5 module attributes {
   6   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Int64], []>, #spirv.resource_limits<>>
   7 } {
   8 // CHECK-LABEL: @trivial_ops
   9 func.func @trivial_ops(%a: index, %b: index) {
  10   // CHECK: spirv.IAdd
  11   %0 = index.add %a, %b
  12   // CHECK: spirv.ISub
  13   %1 = index.sub %a, %b
  14   // CHECK: spirv.IMul
  15   %2 = index.mul %a, %b
  16   // CHECK: spirv.SDiv
  17   %3 = index.divs %a, %b
  18   // CHECK: spirv.UDiv
  19   %4 = index.divu %a, %b
  20   // CHECK: spirv.SRem
  21   %5 = index.rems %a, %b
  22   // CHECK: spirv.UMod
  23   %6 = index.remu %a, %b
  24   // CHECK: spirv.GL.SMax
  25   %7 = index.maxs %a, %b
  26   // CHECK: spirv.GL.UMax
  27   %8 = index.maxu %a, %b
  28   // CHECK: spirv.GL.SMin
  29   %9 = index.mins %a, %b
  30   // CHECK: spirv.GL.UMin
  31   %10 = index.minu %a, %b
  32   // CHECK: spirv.ShiftLeftLogical
  33   %11 = index.shl %a, %b
  34   // CHECK: spirv.ShiftRightArithmetic
  35   %12 = index.shrs %a, %b
  36   // CHECK: spirv.ShiftRightLogical
  37   %13 = index.shru %a, %b
  38   return
  39 }
  40
  41 // CHECK-LABEL: @bitwise_ops
  42 func.func @bitwise_ops(%a: index, %b: index) {
  43   // CHECK: spirv.BitwiseAnd
  44   %0 = index.and %a, %b
  45   // CHECK: spirv.BitwiseOr
  46   %1 = index.or %a, %b
  47   // CHECK: spirv.BitwiseXor
  48   %2 = index.xor %a, %b
  49   return
  50 }
  51
  52 // INDEX32-LABEL: @constant_ops
  53 // INDEX64-LABEL: @constant_ops
  54 func.func @constant_ops() {
  55   // INDEX32: spirv.Constant 42 : i32
  56   // INDEX64: spirv.Constant 42 : i64
  57   %0 = index.constant 42
  58   // INDEX32: spirv.Constant true
  59   // INDEX64: spirv.Constant true
  60   %1 = index.bool.constant true
  61   // INDEX32: spirv.Constant false
  62   // INDEX64: spirv.Constant false
  63   %2 = index.bool.constant false
  64   return
  65 }
  66
  67 // CHECK-LABEL: @ceildivs
  68 // CHECK-SAME: %[[NI:.*]]: index, %[[MI:.*]]: index
  69 func.func @ceildivs(%n: index, %m: index) -> index {
  70   // CHECK-DAG: %[[N:.*]] = builtin.unrealized_conversion_cast %[[NI]]
  71   // CHECK-DAG: %[[M:.*]] = builtin.unrealized_conversion_cast %[[MI]]
  72   // CHECK: %[[ZERO:.*]] = spirv.Constant 0
  73   // CHECK: %[[POS_ONE:.*]] = spirv.Constant 1
  74   // CHECK: %[[NEG_ONE:.*]] = spirv.Constant -1
  75
  76   // CHECK: %[[M_POS:.*]] = spirv.SGreaterThan %[[M]], %[[ZERO]]
  77   // CHECK: %[[X:.*]] = spirv.Select %[[M_POS]], %[[NEG_ONE]], %[[POS_ONE]]
  78
  79   // CHECK: %[[N_PLUS_X:.*]] = spirv.IAdd %[[N]], %[[X]]
  80   // CHECK: %[[N_PLUS_X_DIV_M:.*]] = spirv.SDiv %[[N_PLUS_X]], %[[M]]
  81   // CHECK: %[[POS_RES:.*]] = spirv.IAdd %[[N_PLUS_X_DIV_M]], %[[POS_ONE]]
  82
  83   // CHECK: %[[NEG_N:.*]] = spirv.ISub %[[ZERO]], %[[N]]
  84   // CHECK: %[[NEG_N_DIV_M:.*]] = spirv.SDiv %[[NEG_N]], %[[M]]
  85   // CHECK: %[[NEG_RES:.*]] = spirv.ISub %[[ZERO]], %[[NEG_N_DIV_M]]
  86
  87   // CHECK: %[[N_POS:.*]] = spirv.SGreaterThan %[[N]], %[[ZERO]]
  88   // CHECK: %[[SAME_SIGN:.*]] = spirv.LogicalEqual %[[N_POS]], %[[M_POS]]
  89   // CHECK: %[[N_NON_ZERO:.*]] = spirv.INotEqual %[[N]], %[[ZERO]]
  90   // CHECK: %[[CMP:.*]] = spirv.LogicalAnd %[[SAME_SIGN]], %[[N_NON_ZERO]]
  91   // CHECK: %[[RESULT:.*]] = spirv.Select %[[CMP]], %[[POS_RES]], %[[NEG_RES]]
  92   %result = index.ceildivs %n, %m
  93
  94   // %[[RESULTI:.*] = builtin.unrealized_conversion_cast %[[RESULT]]
  95   // return %[[RESULTI]]
  96   return %result : index
  97 }
  98
  99 // CHECK-LABEL: @ceildivu
 100 // CHECK-SAME: %[[NI:.*]]: index, %[[MI:.*]]: index
 101 func.func @ceildivu(%n: index, %m: index) -> index {
 102   // CHECK-DAG: %[[N:.*]] = builtin.unrealized_conversion_cast %[[NI]]
 103   // CHECK-DAG: %[[M:.*]] = builtin.unrealized_conversion_cast %[[MI]]
 104   // CHECK: %[[ZERO:.*]] = spirv.Constant 0
 105   // CHECK: %[[ONE:.*]] = spirv.Constant 1
 106
 107   // CHECK: %[[N_MINUS_ONE:.*]] = spirv.ISub %[[N]], %[[ONE]]
 108   // CHECK: %[[N_MINUS_ONE_DIV_M:.*]] = spirv.UDiv %[[N_MINUS_ONE]], %[[M]]
 109   // CHECK: %[[N_MINUS_ONE_DIV_M_PLUS_ONE:.*]] = spirv.IAdd %[[N_MINUS_ONE_DIV_M]], %[[ONE]]
 110
 111   // CHECK: %[[CMP:.*]] = spirv.IEqual %[[N]], %[[ZERO]]
 112   // CHECK: %[[RESULT:.*]] = spirv.Select %[[CMP]], %[[ZERO]], %[[N_MINUS_ONE_DIV_M_PLUS_ONE]]
 113   %result = index.ceildivu %n, %m
 114
 115   // %[[RESULTI:.*] = builtin.unrealized_conversion_cast %[[RESULT]]
 116   // return %[[RESULTI]]
 117   return %result : index
 118 }
 119
 120 // CHECK-LABEL: @floordivs
 121 // CHECK-SAME: %[[NI:.*]]: index, %[[MI:.*]]: index
 122 func.func @floordivs(%n: index, %m: index) -> index {
 123   // CHECK-DAG: %[[N:.*]] = builtin.unrealized_conversion_cast %[[NI]]
 124   // CHECK-DAG: %[[M:.*]] = builtin.unrealized_conversion_cast %[[MI]]
 125   // CHECK: %[[ZERO:.*]] = spirv.Constant 0
 126   // CHECK: %[[POS_ONE:.*]] = spirv.Constant 1
 127   // CHECK: %[[NEG_ONE:.*]] = spirv.Constant -1
 128
 129   // CHECK: %[[M_NEG:.*]] = spirv.SLessThan %[[M]], %[[ZERO]]
 130   // CHECK: %[[X:.*]] = spirv.Select %[[M_NEG]], %[[POS_ONE]], %[[NEG_ONE]]
 131
 132   // CHECK: %[[X_MINUS_N:.*]] = spirv.ISub %[[X]], %[[N]]
 133   // CHECK: %[[X_MINUS_N_DIV_M:.*]] = spirv.SDiv %[[X_MINUS_N]], %[[M]]
 134   // CHECK: %[[NEG_RES:.*]] = spirv.ISub %[[NEG_ONE]], %[[X_MINUS_N_DIV_M]]
 135
 136   // CHECK: %[[POS_RES:.*]] = spirv.SDiv %[[N]], %[[M]]
 137
 138   // CHECK: %[[N_NEG:.*]] = spirv.SLessThan %[[N]], %[[ZERO]]
 139   // CHECK: %[[DIFF_SIGN:.*]] = spirv.LogicalNotEqual %[[N_NEG]], %[[M_NEG]]
 140   // CHECK: %[[N_NON_ZERO:.*]] = spirv.INotEqual %[[N]], %[[ZERO]]
 141
 142   // CHECK: %[[CMP:.*]] = spirv.LogicalAnd %[[DIFF_SIGN]], %[[N_NON_ZERO]]
 143   // CHECK: %[[RESULT:.*]] = spirv.Select %[[CMP]], %[[POS_RES]], %[[NEG_RES]]
 144   %result = index.floordivs %n, %m
 145
 146   // %[[RESULTI:.*] = builtin.unrealized_conversion_cast %[[RESULT]]
 147   // return %[[RESULTI]]
 148   return %result : index
 149 }
 150
 151 // CHECK-LABEL: @index_cmp
 152 func.func @index_cmp(%a : index, %b : index) {
 153   // CHECK: spirv.IEqual
 154   %0 = index.cmp eq(%a, %b)
 155   // CHECK: spirv.INotEqual
 156   %1 = index.cmp ne(%a, %b)
 157
 158   // CHECK: spirv.SLessThan
 159   %2 = index.cmp slt(%a, %b)
 160   // CHECK: spirv.SLessThanEqual
 161   %3 = index.cmp sle(%a, %b)
 162   // CHECK: spirv.SGreaterThan
 163   %4 = index.cmp sgt(%a, %b)
 164   // CHECK: spirv.SGreaterThanEqual
 165   %5 = index.cmp sge(%a, %b)
 166
 167   // CHECK: spirv.ULessThan
 168   %6 = index.cmp ult(%a, %b)
 169   // CHECK: spirv.ULessThanEqual
 170   %7 = index.cmp ule(%a, %b)
 171   // CHECK: spirv.UGreaterThan
 172   %8 = index.cmp ugt(%a, %b)
 173   // CHECK: spirv.UGreaterThanEqual
 174   %9 = index.cmp uge(%a, %b)
 175   return
 176 }
 177
 178 // CHECK-LABEL: @index_sizeof
 179 func.func @index_sizeof() {
 180   // CHECK: spirv.Constant 32 : i32
 181   %0 = index.sizeof
 182   return
 183 }
 184
 185 // INDEX32-LABEL: @index_cast_from
 186 // INDEX64-LABEL: @index_cast_from
 187 // INDEX32-SAME: %[[AI:.*]]: index
 188 // INDEX64-SAME: %[[AI:.*]]: index
 189 func.func @index_cast_from(%a: index) -> (i64, i32, i64, i32) {
 190   // INDEX32: %[[A:.*]] = builtin.unrealized_conversion_cast %[[AI]] : index to i32
 191   // INDEX64: %[[A:.*]] = builtin.unrealized_conversion_cast %[[AI]] : index to i64
 192
 193   // INDEX32: %[[V0:.*]] = spirv.SConvert %[[A]] : i32 to i64
 194   %0 = index.casts %a : index to i64
 195   // INDEX64: %[[V1:.*]] = spirv.SConvert %[[A]] : i64 to i32
 196   %1 = index.casts %a : index to i32
 197   // INDEX32: %[[V2:.*]] = spirv.UConvert %[[A]] : i32 to i64
 198   %2 = index.castu %a : index to i64
 199   // INDEX64: %[[V3:.*]] = spirv.UConvert %[[A]] : i64 to i32
 200   %3 = index.castu %a : index to i32
 201
 202   // INDEX32: return %[[V0]], %[[A]], %[[V2]], %[[A]]
 203   // INDEX64: return %[[A]], %[[V1]], %[[A]], %[[V3]]
 204   return %0, %1, %2, %3 : i64, i32, i64, i32
 205 }
 206
 207 // INDEX32-LABEL: @index_cast_to
 208 // INDEX64-LABEL: @index_cast_to
 209 // INDEX32-SAME: %[[A:.*]]: i32, %[[B:.*]]: i64
 210 // INDEX64-SAME: %[[A:.*]]: i32, %[[B:.*]]: i64
 211 func.func @index_cast_to(%a: i32, %b: i64) -> (index, index, index, index) {
 212   // INDEX64: %[[V0:.*]] = spirv.SConvert %[[A]] : i32 to i64
 213   %0 = index.casts %a : i32 to index
 214   // INDEX32: %[[V1:.*]] = spirv.SConvert %[[B]] : i64 to i32
 215   %1 = index.casts %b : i64 to index
 216   // INDEX64: %[[V2:.*]] = spirv.UConvert %[[A]] : i32 to i64
 217   %2 = index.castu %a : i32 to index
 218   // INDEX32: %[[V3:.*]] = spirv.UConvert %[[B]] : i64 to i32
 219   %3 = index.castu %b : i64 to index
 220   return %0, %1, %2, %3 : index, index, index, index
 221 }
 222 }