Skip to content

Commit e42687b

Browse files
htyuhtyu
committed
improve test
1 parent e71cb8a commit e42687b

File tree

1 file changed

+52
-11
lines changed

1 file changed

+52
-11
lines changed

‎mlir/test/Dialect/SCF/loop-unroll.mlir

Lines changed: 52 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -311,10 +311,10 @@ func.func @static_loop_unroll_up_to_factor(%arg0 : memref<?xf32>) {
311311
// Test that epilogue's arguments are correctly renamed.
312312
func.func @static_loop_unroll_by_3_rename_epilogue_arguments() -> (f32, f32) {
313313
%0 = arith.constant 7.0 : f32
314-
%lb = arith.constant 0 : i32
315-
%ub = arith.constant 20 : i32
316-
%step = arith.constant 1 : i32
317-
%result:2 = scf.for %i0 = %lb to %ub step %step iter_args(%arg0 = %0, %arg1 = %0) -> (f32, f32) : i32{
314+
%lb = arith.constant 0 : index
315+
%ub = arith.constant 20 : index
316+
%step = arith.constant 1 : index
317+
%result:2 = scf.for %i0 = %lb to %ub step %step iter_args(%arg0 = %0, %arg1 = %0) -> (f32, f32) {
318318
%add = arith.addf %arg0, %arg1 : f32
319319
%mul = arith.mulf %arg0, %arg1 : f32
320320
scf.yield %add, %mul : f32, f32
@@ -324,13 +324,13 @@ func.func @static_loop_unroll_by_3_rename_epilogue_arguments() -> (f32, f32) {
324324
// UNROLL-BY-3-LABEL: func @static_loop_unroll_by_3_rename_epilogue_arguments
325325
//
326326
// UNROLL-BY-3-DAG: %[[CST:.*]] = arith.constant {{.*}} : f32
327-
// UNROLL-BY-3-DAG: %[[C0:.*]] = arith.constant 0 : i32
328-
// UNROLL-BY-3-DAG: %[[C1:.*]] = arith.constant 1 : i32
329-
// UNROLL-BY-3-DAG: %[[C20:.*]] = arith.constant 20 : i32
330-
// UNROLL-BY-3-DAG: %[[C18:.*]] = arith.constant 18 : i32
331-
// UNROLL-BY-3-DAG: %[[C3:.*]] = arith.constant 3 : i32
327+
// UNROLL-BY-3-DAG: %[[C0:.*]] = arith.constant 0 : index
328+
// UNROLL-BY-3-DAG: %[[C1:.*]] = arith.constant 1 : index
329+
// UNROLL-BY-3-DAG: %[[C20:.*]] = arith.constant 20 : index
330+
// UNROLL-BY-3-DAG: %[[C18:.*]] = arith.constant 18 : index
331+
// UNROLL-BY-3-DAG: %[[C3:.*]] = arith.constant 3 : index
332332
// UNROLL-BY-3: %[[FOR:.*]]:2 = scf.for %[[IV:.*]] = %[[C0]] to %[[C18]] step %[[C3]]
333-
// UNROLL-BY-3-SAME: iter_args(%[[ARG0:.*]] = %[[CST]], %[[ARG1:.*]] = %[[CST]]) -> (f32, f32) : i32 {
333+
// UNROLL-BY-3-SAME: iter_args(%[[ARG0:.*]] = %[[CST]], %[[ARG1:.*]] = %[[CST]]) -> (f32, f32) {
334334
// UNROLL-BY-3-NEXT: %[[ADD0:.*]] = arith.addf %[[ARG0]], %[[ARG1]] : f32
335335
// UNROLL-BY-3-NEXT: %[[MUL0:.*]] = arith.mulf %[[ARG0]], %[[ARG1]] : f32
336336
// UNROLL-BY-3-NEXT: %[[ADD1:.*]] = arith.addf %[[ADD0]], %[[MUL0]] : f32
@@ -340,7 +340,7 @@ func.func @static_loop_unroll_by_3_rename_epilogue_arguments() -> (f32, f32) {
340340
// UNROLL-BY-3-NEXT: scf.yield %[[ADD2]], %[[MUL2]] : f32, f32
341341
// UNROLL-BY-3-NEXT: }
342342
// UNROLL-BY-3: %[[EFOR:.*]]:2 = scf.for %[[EIV:.*]] = %[[C18]] to %[[C20]] step %[[C1]]
343-
// UNROLL-BY-3-SAME: iter_args(%[[EARG0:.*]] = %[[FOR]]#0, %[[EARG1:.*]] = %[[FOR]]#1) -> (f32, f32) : i32 {
343+
// UNROLL-BY-3-SAME: iter_args(%[[EARG0:.*]] = %[[FOR]]#0, %[[EARG1:.*]] = %[[FOR]]#1) -> (f32, f32) {
344344
// UNROLL-BY-3-NEXT: %[[EADD:.*]] = arith.addf %[[EARG0]], %[[EARG1]] : f32
345345
// UNROLL-BY-3-NEXT: %[[EMUL:.*]] = arith.mulf %[[EARG0]], %[[EARG1]] : f32
346346
// UNROLL-BY-3-NEXT: scf.yield %[[EADD]], %[[EMUL]] : f32, f32
@@ -448,3 +448,44 @@ func.func @loop_unroll_yield_iter_arg() {
448448
// CHECK-NEXT: affine.yield %[[ITER_ARG]] : index
449449
// CHECK-NEXT: }
450450
// CHECK-NEXT: return
451+
452+
// -----
453+
454+
// Test the loop unroller works with integer IV type.
455+
func.func @static_loop_unroll_with_integer_iv() -> (f32, f32) {
456+
%0 = arith.constant 7.0 : f32
457+
%lb = arith.constant 0 : i32
458+
%ub = arith.constant 20 : i32
459+
%step = arith.constant 1 : i32
460+
%result:2 = scf.for %i0 = %lb to %ub step %step iter_args(%arg0 = %0, %arg1 = %0) -> (f32, f32) : i32{
461+
%add = arith.addf %arg0, %arg1 : f32
462+
%mul = arith.mulf %arg0, %arg1 : f32
463+
scf.yield %add, %mul : f32, f32
464+
}
465+
return %result#0, %result#1 : f32, f32
466+
}
467+
// UNROLL-BY-3-LABEL: func @static_loop_unroll_with_integer_iv
468+
//
469+
// UNROLL-BY-3-DAG: %[[CST:.*]] = arith.constant {{.*}} : f32
470+
// UNROLL-BY-3-DAG: %[[C0:.*]] = arith.constant 0 : i32
471+
// UNROLL-BY-3-DAG: %[[C1:.*]] = arith.constant 1 : i32
472+
// UNROLL-BY-3-DAG: %[[C20:.*]] = arith.constant 20 : i32
473+
// UNROLL-BY-3-DAG: %[[C18:.*]] = arith.constant 18 : i32
474+
// UNROLL-BY-3-DAG: %[[C3:.*]] = arith.constant 3 : i32
475+
// UNROLL-BY-3: %[[FOR:.*]]:2 = scf.for %[[IV:.*]] = %[[C0]] to %[[C18]] step %[[C3]]
476+
// UNROLL-BY-3-SAME: iter_args(%[[ARG0:.*]] = %[[CST]], %[[ARG1:.*]] = %[[CST]]) -> (f32, f32) : i32 {
477+
// UNROLL-BY-3-NEXT: %[[ADD0:.*]] = arith.addf %[[ARG0]], %[[ARG1]] : f32
478+
// UNROLL-BY-3-NEXT: %[[MUL0:.*]] = arith.mulf %[[ARG0]], %[[ARG1]] : f32
479+
// UNROLL-BY-3-NEXT: %[[ADD1:.*]] = arith.addf %[[ADD0]], %[[MUL0]] : f32
480+
// UNROLL-BY-3-NEXT: %[[MUL1:.*]] = arith.mulf %[[ADD0]], %[[MUL0]] : f32
481+
// UNROLL-BY-3-NEXT: %[[ADD2:.*]] = arith.addf %[[ADD1]], %[[MUL1]] : f32
482+
// UNROLL-BY-3-NEXT: %[[MUL2:.*]] = arith.mulf %[[ADD1]], %[[MUL1]] : f32
483+
// UNROLL-BY-3-NEXT: scf.yield %[[ADD2]], %[[MUL2]] : f32, f32
484+
// UNROLL-BY-3-NEXT: }
485+
// UNROLL-BY-3: %[[EFOR:.*]]:2 = scf.for %[[EIV:.*]] = %[[C18]] to %[[C20]] step %[[C1]]
486+
// UNROLL-BY-3-SAME: iter_args(%[[EARG0:.*]] = %[[FOR]]#0, %[[EARG1:.*]] = %[[FOR]]#1) -> (f32, f32) : i32 {
487+
// UNROLL-BY-3-NEXT: %[[EADD:.*]] = arith.addf %[[EARG0]], %[[EARG1]] : f32
488+
// UNROLL-BY-3-NEXT: %[[EMUL:.*]] = arith.mulf %[[EARG0]], %[[EARG1]] : f32
489+
// UNROLL-BY-3-NEXT: scf.yield %[[EADD]], %[[EMUL]] : f32, f32
490+
// UNROLL-BY-3-NEXT: }
491+
// UNROLL-BY-3-NEXT: return %[[EFOR]]#0, %[[EFOR]]#1 : f32, f32

0 commit comments

Comments
 (0)