b153c05cba
Add uplifting from `scf.while` to `scf.for`. This uplifting expects a very specific ops pattern: * `before` block consisting of single `arith.cmp` op * `after` block containing `arith.addi` We also have a set of patterns to cleanup `scf.while` loops to get them close to the desired form, they will be added in separate PRs. This is part of upstreaming `numba-mlir` scf uplifting pipeline: `cf -> scf.while -> scf.for -> scf.parallel` Original code: https://github.com/numba/numba-mlir/blob/main/mlir/lib/Transforms/PromoteToParallel.cpp
158 lines
7.0 KiB
MLIR
158 lines
7.0 KiB
MLIR
// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-scf-uplift-while-to-for))' -split-input-file -allow-unregistered-dialect | FileCheck %s
|
|
|
|
func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {
|
|
%0 = scf.while (%arg3 = %arg0) : (index) -> (index) {
|
|
%1 = arith.cmpi slt, %arg3, %arg1 : index
|
|
scf.condition(%1) %arg3 : index
|
|
} do {
|
|
^bb0(%arg3: index):
|
|
"test.test1"(%arg3) : (index) -> ()
|
|
%added = arith.addi %arg3, %arg2 : index
|
|
"test.test2"(%added) : (index) -> ()
|
|
scf.yield %added : index
|
|
}
|
|
return %0 : index
|
|
}
|
|
|
|
// CHECK-LABEL: func @uplift_while
|
|
// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index
|
|
// CHECK: %[[C1:.*]] = arith.constant 1 : index
|
|
// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] {
|
|
// CHECK: "test.test1"(%[[I]]) : (index) -> ()
|
|
// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : index
|
|
// CHECK: "test.test2"(%[[INC]]) : (index) -> ()
|
|
// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : index
|
|
// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index
|
|
// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index
|
|
// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index
|
|
// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index
|
|
// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index
|
|
// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index
|
|
// CHECK: return %[[R7]] : index
|
|
|
|
// -----
|
|
|
|
func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {
|
|
%0 = scf.while (%arg3 = %arg0) : (index) -> (index) {
|
|
%1 = arith.cmpi sgt, %arg1, %arg3 : index
|
|
scf.condition(%1) %arg3 : index
|
|
} do {
|
|
^bb0(%arg3: index):
|
|
"test.test1"(%arg3) : (index) -> ()
|
|
%added = arith.addi %arg3, %arg2 : index
|
|
"test.test2"(%added) : (index) -> ()
|
|
scf.yield %added : index
|
|
}
|
|
return %0 : index
|
|
}
|
|
|
|
// CHECK-LABEL: func @uplift_while
|
|
// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index
|
|
// CHECK: %[[C1:.*]] = arith.constant 1 : index
|
|
// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] {
|
|
// CHECK: "test.test1"(%[[I]]) : (index) -> ()
|
|
// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : index
|
|
// CHECK: "test.test2"(%[[INC]]) : (index) -> ()
|
|
// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : index
|
|
// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index
|
|
// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index
|
|
// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index
|
|
// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index
|
|
// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index
|
|
// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index
|
|
// CHECK: return %[[R7]] : index
|
|
|
|
// -----
|
|
|
|
func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> index {
|
|
%0 = scf.while (%arg3 = %arg0) : (index) -> (index) {
|
|
%1 = arith.cmpi slt, %arg3, %arg1 : index
|
|
scf.condition(%1) %arg3 : index
|
|
} do {
|
|
^bb0(%arg3: index):
|
|
"test.test1"(%arg3) : (index) -> ()
|
|
%added = arith.addi %arg2, %arg3 : index
|
|
"test.test2"(%added) : (index) -> ()
|
|
scf.yield %added : index
|
|
}
|
|
return %0 : index
|
|
}
|
|
|
|
// CHECK-LABEL: func @uplift_while
|
|
// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> index
|
|
// CHECK: %[[C1:.*]] = arith.constant 1 : index
|
|
// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] {
|
|
// CHECK: "test.test1"(%[[I]]) : (index) -> ()
|
|
// CHECK: %[[INC:.*]] = arith.addi %[[STEP]], %[[I]] : index
|
|
// CHECK: "test.test2"(%[[INC]]) : (index) -> ()
|
|
// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : index
|
|
// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : index
|
|
// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : index
|
|
// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : index
|
|
// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : index
|
|
// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : index
|
|
// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : index
|
|
// CHECK: return %[[R7]] : index
|
|
|
|
|
|
// -----
|
|
|
|
func.func @uplift_while(%arg0: index, %arg1: index, %arg2: index) -> (i32, f32) {
|
|
%c1 = arith.constant 1 : i32
|
|
%c2 = arith.constant 2.0 : f32
|
|
%0:3 = scf.while (%arg4 = %c1, %arg3 = %arg0, %arg5 = %c2) : (i32, index, f32) -> (i32, index, f32) {
|
|
%1 = arith.cmpi slt, %arg3, %arg1 : index
|
|
scf.condition(%1) %arg4, %arg3, %arg5 : i32, index, f32
|
|
} do {
|
|
^bb0(%arg4: i32, %arg3: index, %arg5: f32):
|
|
%1 = "test.test1"(%arg4) : (i32) -> i32
|
|
%added = arith.addi %arg3, %arg2 : index
|
|
%2 = "test.test2"(%arg5) : (f32) -> f32
|
|
scf.yield %1, %added, %2 : i32, index, f32
|
|
}
|
|
return %0#0, %0#2 : i32, f32
|
|
}
|
|
|
|
// CHECK-LABEL: func @uplift_while
|
|
// CHECK-SAME: (%[[BEGIN:.*]]: index, %[[END:.*]]: index, %[[STEP:.*]]: index) -> (i32, f32)
|
|
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32
|
|
// CHECK-DAG: %[[C2:.*]] = arith.constant 2.000000e+00 : f32
|
|
// CHECK: %[[RES:.*]]:2 = scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]]
|
|
// CHECK-SAME: iter_args(%[[ARG1:.*]] = %[[C1]], %[[ARG2:.*]] = %[[C2]]) -> (i32, f32) {
|
|
// CHECK: %[[T1:.*]] = "test.test1"(%[[ARG1]]) : (i32) -> i32
|
|
// CHECK: %[[T2:.*]] = "test.test2"(%[[ARG2]]) : (f32) -> f32
|
|
// CHECK: scf.yield %[[T1]], %[[T2]] : i32, f32
|
|
// CHECK: return %[[RES]]#0, %[[RES]]#1 : i32, f32
|
|
|
|
// -----
|
|
|
|
func.func @uplift_while(%arg0: i64, %arg1: i64, %arg2: i64) -> i64 {
|
|
%0 = scf.while (%arg3 = %arg0) : (i64) -> (i64) {
|
|
%1 = arith.cmpi slt, %arg3, %arg1 : i64
|
|
scf.condition(%1) %arg3 : i64
|
|
} do {
|
|
^bb0(%arg3: i64):
|
|
"test.test1"(%arg3) : (i64) -> ()
|
|
%added = arith.addi %arg3, %arg2 : i64
|
|
"test.test2"(%added) : (i64) -> ()
|
|
scf.yield %added : i64
|
|
}
|
|
return %0 : i64
|
|
}
|
|
|
|
// CHECK-LABEL: func @uplift_while
|
|
// CHECK-SAME: (%[[BEGIN:.*]]: i64, %[[END:.*]]: i64, %[[STEP:.*]]: i64) -> i64
|
|
// CHECK: %[[C1:.*]] = arith.constant 1 : i64
|
|
// CHECK: scf.for %[[I:.*]] = %[[BEGIN]] to %[[END]] step %[[STEP]] : i64 {
|
|
// CHECK: "test.test1"(%[[I]]) : (i64) -> ()
|
|
// CHECK: %[[INC:.*]] = arith.addi %[[I]], %[[STEP]] : i64
|
|
// CHECK: "test.test2"(%[[INC]]) : (i64) -> ()
|
|
// CHECK: %[[R1:.*]] = arith.subi %[[STEP]], %[[C1]] : i64
|
|
// CHECK: %[[R2:.*]] = arith.subi %[[END]], %[[BEGIN]] : i64
|
|
// CHECK: %[[R3:.*]] = arith.addi %[[R2]], %[[R1]] : i64
|
|
// CHECK: %[[R4:.*]] = arith.divsi %[[R3]], %[[STEP]] : i64
|
|
// CHECK: %[[R5:.*]] = arith.subi %[[R4]], %[[C1]] : i64
|
|
// CHECK: %[[R6:.*]] = arith.muli %[[R5]], %[[STEP]] : i64
|
|
// CHECK: %[[R7:.*]] = arith.addi %[[BEGIN]], %[[R6]] : i64
|
|
// CHECK: return %[[R7]] : i64
|