[mlir][spirv] Add tests for scf.while and scf.for in convert-to-spirv pass

[angelz913]

This PR adds lit tests that check for scf.while and scf.for conversions for the convert-to-spirv pass, introduced in #95942.

[llvmbot]

@llvm/pr-subscribers-mlir-spirv

@llvm/pr-subscribers-mlir

Author: Angel Zhang (angelz913)

Changes

This PR adds lit tests that check for scf.while and scf.for conversions for the convert-to-spirv pass, introduced in #95942.

---

Full diff: https://github.com/llvm/llvm-project/pull/102528.diff

1 Files Affected:

• (modified) mlir/test/Conversion/ConvertToSPIRV/scf.mlir (+54)

diff --git a/mlir/test/Conversion/ConvertToSPIRV/scf.mlir b/mlir/test/Conversion/ConvertToSPIRV/scf.mlir
index 246464928b81c0..9ff5ae7a91b1df 100644
--- a/mlir/test/Conversion/ConvertToSPIRV/scf.mlir
+++ b/mlir/test/Conversion/ConvertToSPIRV/scf.mlir
@@ -33,6 +33,23 @@ func.func @if_yield(%arg0: i1) -> f32 {
 }
 
 // CHECK-LABEL: @while
+// CHECK-SAME: (%[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
+// CHECK:       %[[INITVAR:.*]] = spirv.Constant 2 : i32
+// CHECK:       %[[VAR1:.*]] = spirv.Variable : !spirv.ptr<i32, Function>
+// CHECK:       spirv.mlir.loop {
+// CHECK:         spirv.Branch ^[[HEADER:.*]](%[[ARG1]] : i32)
+// CHECK:       ^[[HEADER]](%[[INDVAR1:.*]]: i32):
+// CHECK:         %[[CMP:.*]] = spirv.SLessThan %[[INDVAR1]], %[[ARG2]] : i32
+// CHECK:         spirv.Store "Function" %[[VAR1]], %[[INDVAR1]] : i32
+// CHECK:         spirv.BranchConditional %[[CMP]], ^[[BODY:.*]](%[[INDVAR1]] : i32), ^[[MERGE:.*]]
+// CHECK:       ^[[BODY]](%[[INDVAR2:.*]]: i32):
+// CHECK:         %[[UPDATED:.*]] = spirv.IMul %[[INDVAR2]], %[[INITVAR]] : i32
+// CHECK:       spirv.Branch ^[[HEADER]](%[[UPDATED]] : i32)
+// CHECK:       ^[[MERGE]]:
+// CHECK:         spirv.mlir.merge
+// CHECK:       }
+// CHECK:       %[[OUT:.*]] = spirv.Load "Function" %[[VAR1]] : i32
+// CHECK:       spirv.ReturnValue %[[OUT]] : i32
 func.func @while(%arg0: i32, %arg1: i32) -> i32 {
   %c2_i32 = arith.constant 2 : i32
   %0 = scf.while (%arg3 = %arg0) : (i32) -> (i32) {
@@ -45,3 +62,40 @@ func.func @while(%arg0: i32, %arg1: i32) -> i32 {
   }
   return %0 : i32
 }
+
+// CHECK-LABEL: @for
+// CHECK:       %[[LB:.*]] = spirv.Constant 4 : i32
+// CHECK:       %[[UB:.*]] = spirv.Constant 42 : i32
+// CHECK:       %[[STEP:.*]] = spirv.Constant 2 : i32
+// CHECK:       %[[INITVAR1:.*]] = spirv.Constant 0.000000e+00 : f32
+// CHECK:       %[[INITVAR2:.*]] = spirv.Constant 1.000000e+00 : f32
+// CHECK:       %[[VAR1:.*]] = spirv.Variable : !spirv.ptr<f32, Function>
+// CHECK:       %[[VAR2:.*]] = spirv.Variable : !spirv.ptr<f32, Function>
+// CHECK:       spirv.mlir.loop {
+// CHECK:         spirv.Branch ^[[HEADER:.*]](%[[LB]], %[[INITVAR1]], %[[INITVAR2]] : i32, f32, f32)
+// CHECK:       ^[[HEADER]](%[[INDVAR:.*]]: i32, %[[CARRIED1:.*]]: f32, %[[CARRIED2:.*]]: f32):
+// CHECK:         %[[CMP:.*]] = spirv.SLessThan %[[INDVAR]], %[[UB]] : i32
+// CHECK:         spirv.BranchConditional %[[CMP]], ^[[BODY:.*]], ^[[MERGE:.*]]
+// CHECK:       ^[[BODY]]:
+// CHECK:         %[[UPDATED:.*]] = spirv.FAdd %[[CARRIED1]], %[[CARRIED1]] : f32
+// CHECK-DAG:     %[[INCREMENT:.*]] = spirv.IAdd %[[INDVAR]], %[[STEP]] : i32
+// CHECK-DAG:     spirv.Store "Function" %[[VAR1]], %[[UPDATED]] : f32
+// CHECK-DAG:     spirv.Store "Function" %[[VAR2]], %[[UPDATED]] : f32
+// CHECK:       spirv.Branch ^[[HEADER]](%[[INCREMENT]], %[[UPDATED]], %[[UPDATED]] : i32, f32, f32)
+// CHECK:       ^[[MERGE]]:
+// CHECK:         spirv.mlir.merge
+// CHECK:       }
+// CHECK-DAG:  %[[OUT1:.*]] = spirv.Load "Function" %[[VAR1]] : f32
+// CHECK-DAG:  %[[OUT2:.*]] = spirv.Load "Function" %[[VAR2]] : f32
+func.func @for() {
+  %lb = arith.constant 4 : index
+  %ub = arith.constant 42 : index
+  %step = arith.constant 2 : index
+  %s0 = arith.constant 0.0 : f32
+  %s1 = arith.constant 1.0 : f32
+  %result:2 = scf.for %i0 = %lb to %ub step %step iter_args(%si = %s0, %sj = %s1) -> (f32, f32) {
+    %sn = arith.addf %si, %si : f32
+    scf.yield %sn, %sn: f32, f32
+  }
+  return
+}

[kuhar]

Are these tested by the scf to spirv conversion? If yes, could we trim these tests down so that we don't duplicate tests?

[angelz913]

Are these tested by the scf to spirv conversion? If yes, could we trim these tests down so that we don't duplicate tests?

Yes, these tests were mostly from the scf-to-spirv conversion. I just realized that I forgot to include them in the initial version PR, which made me think that scf.while and scf.for were not covered. To prevent duplicate tests, I can try to include tests that combine with memref (with storage class conversion that is not covered in scf-to-spirv).

[kuhar]

I think we can check that they produce spirv.mlir.loop, spirv.Branch, etc., but not the exact instruction sequence. This way we will know that the conversion produced something reasonable, end the exact expansion is tested in the scf to spirv test.

[kuhar]

LGTM