171 lines · plain
1// RUN: mlir-opt -convert-scf-to-spirv %s -o - | FileCheck %s2 3module attributes {4 spirv.target_env = #spirv.target_env<5 #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>6} {7 8// CHECK-LABEL: @kernel_simple_selection9func.func @kernel_simple_selection(%arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>>, %arg3 : i1) {10 %value = arith.constant 0.0 : f3211 %i = arith.constant 0 : index12 13 // CHECK: spirv.mlir.selection {14 // CHECK-NEXT: spirv.BranchConditional {{%.*}}, [[TRUE:\^.*]], [[MERGE:\^.*]]15 // CHECK-NEXT: [[TRUE]]:16 // CHECK: spirv.Branch [[MERGE]]17 // CHECK-NEXT: [[MERGE]]:18 // CHECK-NEXT: spirv.mlir.merge19 // CHECK-NEXT: }20 // CHECK-NEXT: spirv.Return21 22 scf.if %arg3 {23 memref.store %value, %arg2[%i] : memref<10xf32, #spirv.storage_class<StorageBuffer>>24 }25 return26}27 28// CHECK-LABEL: @kernel_nested_selection29func.func @kernel_nested_selection(%arg3 : memref<10xf32, #spirv.storage_class<StorageBuffer>>, %arg4 : memref<10xf32, #spirv.storage_class<StorageBuffer>>, %arg5 : i1, %arg6 : i1) {30 %i = arith.constant 0 : index31 %j = arith.constant 9 : index32 33 // CHECK: spirv.mlir.selection {34 // CHECK-NEXT: spirv.BranchConditional {{%.*}}, [[TRUE_TOP:\^.*]], [[FALSE_TOP:\^.*]]35 // CHECK-NEXT: [[TRUE_TOP]]:36 // CHECK-NEXT: spirv.mlir.selection {37 // CHECK-NEXT: spirv.BranchConditional {{%.*}}, [[TRUE_NESTED_TRUE_PATH:\^.*]], [[FALSE_NESTED_TRUE_PATH:\^.*]]38 // CHECK-NEXT: [[TRUE_NESTED_TRUE_PATH]]:39 // CHECK: spirv.Branch [[MERGE_NESTED_TRUE_PATH:\^.*]]40 // CHECK-NEXT: [[FALSE_NESTED_TRUE_PATH]]:41 // CHECK: spirv.Branch [[MERGE_NESTED_TRUE_PATH]]42 // CHECK-NEXT: [[MERGE_NESTED_TRUE_PATH]]:43 // CHECK-NEXT: spirv.mlir.merge44 // CHECK-NEXT: }45 // CHECK-NEXT: spirv.Branch [[MERGE_TOP:\^.*]]46 // CHECK-NEXT: [[FALSE_TOP]]:47 // CHECK-NEXT: spirv.mlir.selection {48 // CHECK-NEXT: spirv.BranchConditional {{%.*}}, [[TRUE_NESTED_FALSE_PATH:\^.*]], [[FALSE_NESTED_FALSE_PATH:\^.*]]49 // CHECK-NEXT: [[TRUE_NESTED_FALSE_PATH]]:50 // CHECK: spirv.Branch [[MERGE_NESTED_FALSE_PATH:\^.*]]51 // CHECK-NEXT: [[FALSE_NESTED_FALSE_PATH]]:52 // CHECK: spirv.Branch [[MERGE_NESTED_FALSE_PATH]]53 // CHECK: [[MERGE_NESTED_FALSE_PATH]]:54 // CHECK-NEXT: spirv.mlir.merge55 // CHECK-NEXT: }56 // CHECK-NEXT: spirv.Branch [[MERGE_TOP]]57 // CHECK-NEXT: [[MERGE_TOP]]:58 // CHECK-NEXT: spirv.mlir.merge59 // CHECK-NEXT: }60 // CHECK-NEXT: spirv.Return61 62 scf.if %arg5 {63 scf.if %arg6 {64 %value = memref.load %arg3[%i] : memref<10xf32, #spirv.storage_class<StorageBuffer>>65 memref.store %value, %arg4[%i] : memref<10xf32, #spirv.storage_class<StorageBuffer>>66 } else {67 %value = memref.load %arg4[%i] : memref<10xf32, #spirv.storage_class<StorageBuffer>>68 memref.store %value, %arg3[%i] : memref<10xf32, #spirv.storage_class<StorageBuffer>>69 }70 } else {71 scf.if %arg6 {72 %value = memref.load %arg3[%j] : memref<10xf32, #spirv.storage_class<StorageBuffer>>73 memref.store %value, %arg4[%j] : memref<10xf32, #spirv.storage_class<StorageBuffer>>74 } else {75 %value = memref.load %arg4[%j] : memref<10xf32, #spirv.storage_class<StorageBuffer>>76 memref.store %value, %arg3[%j] : memref<10xf32, #spirv.storage_class<StorageBuffer>>77 }78 }79 return80}81 82// CHECK-LABEL: @simple_if_yield83func.func @simple_if_yield(%arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>>, %arg3 : i1) {84 // CHECK: %[[VAR1:.*]] = spirv.Variable : !spirv.ptr<f32, Function>85 // CHECK: %[[VAR2:.*]] = spirv.Variable : !spirv.ptr<f32, Function>86 // CHECK: spirv.mlir.selection {87 // CHECK-NEXT: spirv.BranchConditional {{%.*}}, [[TRUE:\^.*]], [[FALSE:\^.*]]88 // CHECK-NEXT: [[TRUE]]:89 // CHECK: %[[RET1TRUE:.*]] = spirv.Constant 0.000000e+00 : f3290 // CHECK: %[[RET2TRUE:.*]] = spirv.Constant 1.000000e+00 : f3291 // CHECK-DAG: spirv.Store "Function" %[[VAR1]], %[[RET1TRUE]] : f3292 // CHECK-DAG: spirv.Store "Function" %[[VAR2]], %[[RET2TRUE]] : f3293 // CHECK: spirv.Branch ^[[MERGE:.*]]94 // CHECK-NEXT: [[FALSE]]:95 // CHECK: %[[RET2FALSE:.*]] = spirv.Constant 2.000000e+00 : f3296 // CHECK: %[[RET1FALSE:.*]] = spirv.Constant 3.000000e+00 : f3297 // CHECK-DAG: spirv.Store "Function" %[[VAR1]], %[[RET1FALSE]] : f3298 // CHECK-DAG: spirv.Store "Function" %[[VAR2]], %[[RET2FALSE]] : f3299 // CHECK: spirv.Branch ^[[MERGE]]100 // CHECK-NEXT: ^[[MERGE]]:101 // CHECK: spirv.mlir.merge102 // CHECK-NEXT: }103 // CHECK-DAG: %[[OUT1:.*]] = spirv.Load "Function" %[[VAR1]] : f32104 // CHECK-DAG: %[[OUT2:.*]] = spirv.Load "Function" %[[VAR2]] : f32105 // CHECK: spirv.Store "StorageBuffer" {{%.*}}, %[[OUT1]] : f32106 // CHECK: spirv.Store "StorageBuffer" {{%.*}}, %[[OUT2]] : f32107 // CHECK: spirv.Return108 %0:2 = scf.if %arg3 -> (f32, f32) {109 %c0 = arith.constant 0.0 : f32110 %c1 = arith.constant 1.0 : f32111 scf.yield %c0, %c1 : f32, f32112 } else {113 %c0 = arith.constant 2.0 : f32114 %c1 = arith.constant 3.0 : f32115 scf.yield %c1, %c0 : f32, f32116 }117 %i = arith.constant 0 : index118 %j = arith.constant 1 : index119 memref.store %0#0, %arg2[%i] : memref<10xf32, #spirv.storage_class<StorageBuffer>>120 memref.store %0#1, %arg2[%j] : memref<10xf32, #spirv.storage_class<StorageBuffer>>121 return122}123 124// TODO: The transformation should only be legal if VariablePointer capability125// is supported. This test is still useful to make sure we can handle scf op126// result with type change.127func.func @simple_if_yield_type_change(%arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>>, %arg3 : memref<10xf32, #spirv.storage_class<StorageBuffer>>, %arg4 : i1) {128 // CHECK-LABEL: @simple_if_yield_type_change129 // CHECK: %[[VAR:.*]] = spirv.Variable : !spirv.ptr<!spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>, Function>130 // CHECK: spirv.mlir.selection {131 // CHECK-NEXT: spirv.BranchConditional {{%.*}}, [[TRUE:\^.*]], [[FALSE:\^.*]]132 // CHECK-NEXT: [[TRUE]]:133 // CHECK: spirv.Store "Function" %[[VAR]], {{%.*}} : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>134 // CHECK: spirv.Branch ^[[MERGE:.*]]135 // CHECK-NEXT: [[FALSE]]:136 // CHECK: spirv.Store "Function" %[[VAR]], {{%.*}} : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>137 // CHECK: spirv.Branch ^[[MERGE]]138 // CHECK-NEXT: ^[[MERGE]]:139 // CHECK: spirv.mlir.merge140 // CHECK-NEXT: }141 // CHECK: %[[OUT:.*]] = spirv.Load "Function" %[[VAR]] : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>142 // CHECK: %[[ADD:.*]] = spirv.AccessChain %[[OUT]][{{%.*}}, {{%.*}}] : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>143 // CHECK: spirv.Store "StorageBuffer" %[[ADD]], {{%.*}} : f32144 // CHECK: spirv.Return145 %i = arith.constant 0 : index146 %value = arith.constant 0.0 : f32147 %0 = scf.if %arg4 -> (memref<10xf32, #spirv.storage_class<StorageBuffer>>) {148 scf.yield %arg2 : memref<10xf32, #spirv.storage_class<StorageBuffer>>149 } else {150 scf.yield %arg3 : memref<10xf32, #spirv.storage_class<StorageBuffer>>151 }152 memref.store %value, %0[%i] : memref<10xf32, #spirv.storage_class<StorageBuffer>>153 return154}155 156// Memrefs without a spirv storage class are not supported. The conversion157// should preserve the `scf.if` and not crash.158func.func @unsupported_yield_type(%arg0 : memref<8xi32>, %arg1 : memref<8xi32>, %c : i1) {159// CHECK-LABEL: @unsupported_yield_type160// CHECK-NEXT: scf.if161// CHECK: spirv.Return162 %r = scf.if %c -> (memref<8xi32>) {163 scf.yield %arg0 : memref<8xi32>164 } else {165 scf.yield %arg1 : memref<8xi32>166 }167 return168}169 170} // end module171