mlir/test/Target/SPIRV/loop.mlir

   1 // RUN: mlir-translate -no-implicit-module -split-input-file -test-spirv-roundtrip %s | FileCheck %s
   2
   3 // Single loop
   4
   5 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
   6   // for (int i = 0; i < count; ++i) {}
   7 // CHECK-LABEL: @loop
   8   spirv.func @loop(%count : i32) -> () "None" {
   9     %zero = spirv.Constant 0: i32
  10     %one = spirv.Constant 1: i32
  11     %var = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
  12
  13 // CHECK:        spirv.Branch ^bb1
  14 // CHECK-NEXT: ^bb1:
  15 // CHECK-NEXT:   spirv.mlir.loop
  16     spirv.mlir.loop {
  17 // CHECK-NEXT:     spirv.Branch ^bb1
  18       spirv.Branch ^header
  19
  20 // CHECK-NEXT:   ^bb1:
  21     ^header:
  22 // CHECK-NEXT:     spirv.Load
  23       %val0 = spirv.Load "Function" %var : i32
  24 // CHECK-NEXT:     spirv.SLessThan
  25       %cmp = spirv.SLessThan %val0, %count : i32
  26 // CHECK-NEXT:     spirv.BranchConditional %{{.*}} [1, 1], ^bb2, ^bb4
  27       spirv.BranchConditional %cmp [1, 1], ^body, ^merge
  28
  29 // CHECK-NEXT:   ^bb2:
  30     ^body:
  31       // Do nothing
  32 // CHECK-NEXT:     spirv.Branch ^bb3
  33       spirv.Branch ^continue
  34
  35 // CHECK-NEXT:   ^bb3:
  36     ^continue:
  37 // CHECK-NEXT:     spirv.Load
  38       %val1 = spirv.Load "Function" %var : i32
  39 // CHECK-NEXT:     spirv.Constant 1
  40 // CHECK-NEXT:     spirv.IAdd
  41       %add = spirv.IAdd %val1, %one : i32
  42 // CHECK-NEXT:     spirv.Store
  43       spirv.Store "Function" %var, %add : i32
  44 // CHECK-NEXT:     spirv.Branch ^bb1
  45       spirv.Branch ^header
  46
  47 // CHECK-NEXT:   ^bb4:
  48 // CHECK-NEXT:     spirv.mlir.merge
  49     ^merge:
  50       spirv.mlir.merge
  51     }
  52     spirv.Return
  53   }
  54
  55   spirv.func @main() -> () "None" {
  56     spirv.Return
  57   }
  58   spirv.EntryPoint "GLCompute" @main
  59 }
  60
  61 // -----
  62
  63 // Single loop with block arguments
  64
  65 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
  66   spirv.GlobalVariable @GV1 bind(0, 0) : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>
  67   spirv.GlobalVariable @GV2 bind(0, 1) : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>
  68 // CHECK-LABEL: @loop_kernel
  69   spirv.func @loop_kernel() "None" {
  70     %0 = spirv.mlir.addressof @GV1 : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>
  71     %1 = spirv.Constant 0 : i32
  72     %2 = spirv.AccessChain %0[%1] : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>, i32 -> !spirv.ptr<!spirv.array<10 x f32, stride=4>, StorageBuffer>
  73     %3 = spirv.mlir.addressof @GV2 : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>
  74     %5 = spirv.AccessChain %3[%1] : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>, i32 -> !spirv.ptr<!spirv.array<10 x f32, stride=4>, StorageBuffer>
  75     %6 = spirv.Constant 4 : i32
  76     %7 = spirv.Constant 42 : i32
  77     %8 = spirv.Constant 2 : i32
  78 // CHECK:        spirv.Branch ^bb1(%{{.*}} : i32)
  79 // CHECK-NEXT: ^bb1(%[[OUTARG:.*]]: i32):
  80 // CHECK-NEXT:   spirv.mlir.loop {
  81     spirv.mlir.loop {
  82 // CHECK-NEXT:     spirv.Branch ^bb1(%[[OUTARG]] : i32)
  83       spirv.Branch ^header(%6 : i32)
  84 // CHECK-NEXT:   ^bb1(%[[HEADARG:.*]]: i32):
  85     ^header(%9: i32):
  86       %10 = spirv.SLessThan %9, %7 : i32
  87 // CHECK:          spirv.BranchConditional %{{.*}}, ^bb2, ^bb3
  88       spirv.BranchConditional %10, ^body, ^merge
  89 // CHECK-NEXT:   ^bb2:     // pred: ^bb1
  90     ^body:
  91       %11 = spirv.AccessChain %2[%9] : !spirv.ptr<!spirv.array<10 x f32, stride=4>, StorageBuffer>, i32 -> !spirv.ptr<f32, StorageBuffer>
  92       %12 = spirv.Load "StorageBuffer" %11 : f32
  93       %13 = spirv.AccessChain %5[%9] : !spirv.ptr<!spirv.array<10 x f32, stride=4>, StorageBuffer>, i32 -> !spirv.ptr<f32, StorageBuffer>
  94       spirv.Store "StorageBuffer" %13, %12 : f32
  95 // CHECK:          %[[ADD:.*]] = spirv.IAdd
  96       %14 = spirv.IAdd %9, %8 : i32
  97 // CHECK-NEXT:     spirv.Branch ^bb1(%[[ADD]] : i32)
  98       spirv.Branch ^header(%14 : i32)
  99 // CHECK-NEXT:   ^bb3:
 100     ^merge:
 101 // CHECK-NEXT:     spirv.mlir.merge
 102       spirv.mlir.merge
 103     }
 104     spirv.Return
 105   }
 106   spirv.EntryPoint "GLCompute" @loop_kernel
 107   spirv.ExecutionMode @loop_kernel "LocalSize", 1, 1, 1
 108 }
 109
 110 // -----
 111
 112 // Nested loop
 113
 114 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
 115   // for (int i = 0; i < count; ++i) {
 116   //   for (int j = 0; j < count; ++j) { }
 117   // }
 118 // CHECK-LABEL: @loop
 119   spirv.func @loop(%count : i32) -> () "None" {
 120     %zero = spirv.Constant 0: i32
 121     %one = spirv.Constant 1: i32
 122     %ivar = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
 123     %jvar = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
 124
 125 // CHECK:        spirv.Branch ^bb1
 126 // CHECK-NEXT: ^bb1:
 127 // CHECK-NEXT:   spirv.mlir.loop control(Unroll)
 128     spirv.mlir.loop control(Unroll) {
 129 // CHECK-NEXT:     spirv.Branch ^bb1
 130       spirv.Branch ^header
 131
 132 // CHECK-NEXT:   ^bb1:
 133     ^header:
 134 // CHECK-NEXT:     spirv.Load
 135       %ival0 = spirv.Load "Function" %ivar : i32
 136 // CHECK-NEXT:     spirv.SLessThan
 137       %icmp = spirv.SLessThan %ival0, %count : i32
 138 // CHECK-NEXT:     spirv.BranchConditional %{{.*}}, ^bb2, ^bb5
 139       spirv.BranchConditional %icmp, ^body, ^merge
 140
 141 // CHECK-NEXT:   ^bb2:
 142     ^body:
 143 // CHECK-NEXT:     spirv.Constant 0
 144 // CHECK-NEXT:           spirv.Store
 145       spirv.Store "Function" %jvar, %zero : i32
 146 // CHECK-NEXT:     spirv.Branch ^bb3
 147 // CHECK-NEXT:   ^bb3:
 148 // CHECK-NEXT:     spirv.mlir.loop control(DontUnroll)
 149       spirv.mlir.loop control(DontUnroll) {
 150 // CHECK-NEXT:       spirv.Branch ^bb1
 151         spirv.Branch ^header
 152
 153 // CHECK-NEXT:     ^bb1:
 154       ^header:
 155 // CHECK-NEXT:       spirv.Load
 156         %jval0 = spirv.Load "Function" %jvar : i32
 157 // CHECK-NEXT:       spirv.SLessThan
 158         %jcmp = spirv.SLessThan %jval0, %count : i32
 159 // CHECK-NEXT:       spirv.BranchConditional %{{.*}}, ^bb2, ^bb4
 160         spirv.BranchConditional %jcmp, ^body, ^merge
 161
 162 // CHECK-NEXT:     ^bb2:
 163       ^body:
 164         // Do nothing
 165 // CHECK-NEXT:       spirv.Branch ^bb3
 166         spirv.Branch ^continue
 167
 168 // CHECK-NEXT:     ^bb3:
 169       ^continue:
 170 // CHECK-NEXT:       spirv.Load
 171         %jval1 = spirv.Load "Function" %jvar : i32
 172 // CHECK-NEXT:       spirv.Constant 1
 173 // CHECK-NEXT:       spirv.IAdd
 174         %add = spirv.IAdd %jval1, %one : i32
 175 // CHECK-NEXT:       spirv.Store
 176         spirv.Store "Function" %jvar, %add : i32
 177 // CHECK-NEXT:       spirv.Branch ^bb1
 178         spirv.Branch ^header
 179
 180 // CHECK-NEXT:     ^bb4:
 181       ^merge:
 182 // CHECK-NEXT:       spirv.mlir.merge
 183         spirv.mlir.merge
 184       } // end inner loop
 185
 186 // CHECK:          spirv.Branch ^bb4
 187       spirv.Branch ^continue
 188
 189 // CHECK-NEXT:   ^bb4:
 190     ^continue:
 191 // CHECK-NEXT:     spirv.Load
 192       %ival1 = spirv.Load "Function" %ivar : i32
 193 // CHECK-NEXT:     spirv.Constant 1
 194 // CHECK-NEXT:     spirv.IAdd
 195       %add = spirv.IAdd %ival1, %one : i32
 196 // CHECK-NEXT:     spirv.Store
 197       spirv.Store "Function" %ivar, %add : i32
 198 // CHECK-NEXT:     spirv.Branch ^bb1
 199       spirv.Branch ^header
 200
 201 // CHECK-NEXT:   ^bb5:
 202 // CHECK-NEXT:     spirv.mlir.merge
 203     ^merge:
 204       spirv.mlir.merge
 205     } // end outer loop
 206     spirv.Return
 207   }
 208
 209   spirv.func @main() -> () "None" {
 210     spirv.Return
 211   }
 212   spirv.EntryPoint "GLCompute" @main
 213 }
 214
 215
 216 // -----
 217
 218 // Loop with selection in its header
 219
 220 spirv.module Physical64 OpenCL requires #spirv.vce<v1.0, [Kernel, Linkage, Addresses, Int64], []> {
 221 // CHECK-LABEL:   @kernel
 222 // CHECK-SAME:    (%[[INPUT0:.+]]: i64)
 223   spirv.func @kernel(%input: i64) "None" {
 224 // CHECK-NEXT:     %[[VAR:.+]] = spirv.Variable : !spirv.ptr<i1, Function>
 225 // CHECK-NEXT:     spirv.Branch ^[[BB0:.+]](%[[INPUT0]] : i64)
 226 // CHECK-NEXT:   ^[[BB0]](%[[INPUT1:.+]]: i64):
 227     %cst0_i64 = spirv.Constant 0 : i64
 228     %true = spirv.Constant true
 229     %false = spirv.Constant false
 230 // CHECK-NEXT:     spirv.mlir.loop {
 231     spirv.mlir.loop {
 232 // CHECK-NEXT:       spirv.Branch ^[[LOOP_HEADER:.+]](%[[INPUT1]] : i64)
 233       spirv.Branch ^loop_header(%input : i64)
 234 // CHECK-NEXT:     ^[[LOOP_HEADER]](%[[ARG1:.+]]: i64):
 235     ^loop_header(%arg1: i64):
 236 // CHECK-NEXT:       spirv.Branch ^[[LOOP_BODY:.+]]
 237 // CHECK-NEXT:     ^[[LOOP_BODY]]:
 238 // CHECK-NEXT:         %[[C0:.+]] = spirv.Constant 0 : i64
 239       %gt = spirv.SGreaterThan %arg1, %cst0_i64 : i64
 240 // CHECK-NEXT:         %[[GT:.+]] = spirv.SGreaterThan %[[ARG1]], %[[C0]] : i64
 241 // CHECK-NEXT:         spirv.Branch ^[[BB1:.+]]
 242 // CHECK-NEXT:     ^[[BB1]]:
 243       %var = spirv.Variable : !spirv.ptr<i1, Function>
 244 // CHECK-NEXT:       spirv.mlir.selection {
 245       spirv.mlir.selection {
 246 // CHECK-NEXT:         spirv.BranchConditional %[[GT]], ^[[THEN:.+]], ^[[ELSE:.+]]
 247         spirv.BranchConditional %gt, ^then, ^else
 248 // CHECK-NEXT:       ^[[THEN]]:
 249       ^then:
 250 // CHECK-NEXT:         %true = spirv.Constant true
 251 // CHECK-NEXT:         spirv.Store "Function" %[[VAR]], %true : i1
 252         spirv.Store "Function" %var, %true : i1
 253 // CHECK-NEXT:         spirv.Branch ^[[SELECTION_MERGE:.+]]
 254         spirv.Branch ^selection_merge
 255 // CHECK-NEXT:       ^[[ELSE]]:
 256       ^else:
 257 // CHECK-NEXT:         %false = spirv.Constant false
 258 // CHECK-NEXT:         spirv.Store "Function" %[[VAR]], %false : i1
 259         spirv.Store "Function" %var, %false : i1
 260 // CHECK-NEXT:         spirv.Branch ^[[SELECTION_MERGE]]
 261         spirv.Branch ^selection_merge
 262 // CHECK-NEXT:       ^[[SELECTION_MERGE]]:
 263       ^selection_merge:
 264 // CHECK-NEXT:         spirv.mlir.merge
 265         spirv.mlir.merge
 266 // CHECK-NEXT:       }
 267       }
 268 // CHECK-NEXT:       %[[LOAD:.+]] = spirv.Load "Function" %[[VAR]] : i1
 269       %load = spirv.Load "Function" %var : i1
 270 // CHECK-NEXT:       spirv.BranchConditional %[[LOAD]], ^[[CONTINUE:.+]](%[[ARG1]] : i64), ^[[LOOP_MERGE:.+]]
 271       spirv.BranchConditional %load, ^continue(%arg1 : i64), ^loop_merge
 272 // CHECK-NEXT:     ^[[CONTINUE]](%[[ARG2:.+]]: i64):
 273     ^continue(%arg2: i64):
 274 // CHECK-NEXT:       %[[C0:.+]] = spirv.Constant 0 : i64
 275 // CHECK-NEXT:       %[[LT:.+]] = spirv.SLessThan %[[ARG2]], %[[C0]] : i64
 276       %lt = spirv.SLessThan %arg2, %cst0_i64 : i64
 277 // CHECK-NEXT:       spirv.Store "Function" %[[VAR]], %[[LT]] : i1
 278       spirv.Store "Function" %var, %lt : i1
 279 // CHECK-NEXT:       spirv.Branch ^[[LOOP_HEADER]](%[[ARG2]] : i64)
 280       spirv.Branch ^loop_header(%arg2 : i64)
 281 // CHECK-NEXT:     ^[[LOOP_MERGE]]:
 282     ^loop_merge:
 283 // CHECK-NEXT:       spirv.mlir.merge
 284       spirv.mlir.merge
 285 // CHECK-NEXT:     }
 286     }
 287 // CHECK-NEXT:     spirv.Return
 288     spirv.Return
 289   }
 290 }