efc0ba0275
This patch made a minor refactor of LoopCoalescing.cpp's walkLoops templated method and placed it in Affine's LoopUtils.cpp/h. This method is also renamed as coalescePerfectlyNestedLoops method. This minor change enables this method to be invoked by both the original LoopCoalescing pass as well as the newly introduced loop.coalesce transform op. The loop.coalesce transform op has the ability to coalesce affine, and scf loop nests, leveraging existing LoopCoalescing mechanism. I have created it inside the SCFTransformOps.td instead of AffineTransformOps.td as it feels to be similar in spirit as the loop.unroll op that can handle both scf and affine loops. Please let me know if you feel that this op should be moved into AffineTransformOps.td instead. The testcase added illustrates loop.coalesce transform op working for scf, affine loops (inner, outer) as well as coalesced loop can be further unrolled (achieving composibility). Reviewed By: ftynse Differential Revision: https://reviews.llvm.org/D141202
261 lines
8.6 KiB
MLIR
261 lines
8.6 KiB
MLIR
// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s
|
|
|
|
// CHECK-LABEL: @get_parent_for_op
|
|
func.func @get_parent_for_op(%arg0: index, %arg1: index, %arg2: index) {
|
|
// expected-remark @below {{first loop}}
|
|
scf.for %i = %arg0 to %arg1 step %arg2 {
|
|
// expected-remark @below {{second loop}}
|
|
scf.for %j = %arg0 to %arg1 step %arg2 {
|
|
// expected-remark @below {{third loop}}
|
|
scf.for %k = %arg0 to %arg1 step %arg2 {
|
|
arith.addi %i, %j : index
|
|
}
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
transform.sequence failures(propagate) {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["arith.addi"]} in %arg1
|
|
// CHECK: = transform.loop.get_parent_for
|
|
%1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
|
|
%2 = transform.loop.get_parent_for %0 { num_loops = 2 } : (!pdl.operation) -> !transform.op<"scf.for">
|
|
%3 = transform.loop.get_parent_for %0 { num_loops = 3 } : (!pdl.operation) -> !transform.op<"scf.for">
|
|
transform.test_print_remark_at_operand %1, "third loop" : !transform.op<"scf.for">
|
|
transform.test_print_remark_at_operand %2, "second loop" : !transform.op<"scf.for">
|
|
transform.test_print_remark_at_operand %3, "first loop" : !transform.op<"scf.for">
|
|
}
|
|
|
|
// -----
|
|
|
|
func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
|
|
// expected-note @below {{target op}}
|
|
arith.addi %arg0, %arg1 : index
|
|
return
|
|
}
|
|
|
|
transform.sequence failures(propagate) {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["arith.addi"]} in %arg1
|
|
// expected-error @below {{could not find an 'scf.for' parent}}
|
|
%1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
|
|
}
|
|
|
|
// -----
|
|
|
|
// Outlined functions:
|
|
//
|
|
// CHECK: func @foo(%{{.+}}, %{{.+}}, %{{.+}}, %{{.+}})
|
|
// CHECK: scf.for
|
|
// CHECK: arith.addi
|
|
//
|
|
// CHECK: func @foo[[SUFFIX:.+]](%{{.+}}, %{{.+}}, %{{.+}})
|
|
// CHECK: scf.for
|
|
// CHECK: arith.addi
|
|
//
|
|
// CHECK-LABEL @loop_outline_op
|
|
func.func @loop_outline_op(%arg0: index, %arg1: index, %arg2: index) {
|
|
// CHECK: scf.for
|
|
// CHECK-NOT: scf.for
|
|
// CHECK: scf.execute_region
|
|
// CHECK: func.call @foo
|
|
scf.for %i = %arg0 to %arg1 step %arg2 {
|
|
scf.for %j = %arg0 to %arg1 step %arg2 {
|
|
arith.addi %i, %j : index
|
|
}
|
|
}
|
|
// CHECK: scf.execute_region
|
|
// CHECK-NOT: scf.for
|
|
// CHECK: func.call @foo[[SUFFIX]]
|
|
scf.for %j = %arg0 to %arg1 step %arg2 {
|
|
arith.addi %j, %j : index
|
|
}
|
|
return
|
|
}
|
|
|
|
transform.sequence failures(propagate) {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["arith.addi"]} in %arg1
|
|
%1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
|
|
// CHECK: = transform.loop.outline %{{.*}}
|
|
transform.loop.outline %1 {func_name = "foo"} : (!transform.op<"scf.for">) -> !pdl.operation
|
|
}
|
|
|
|
// -----
|
|
|
|
// CHECK-LABEL: @loop_peel_op
|
|
func.func @loop_peel_op() {
|
|
// CHECK: %[[C0:.+]] = arith.constant 0
|
|
// CHECK: %[[C42:.+]] = arith.constant 42
|
|
// CHECK: %[[C5:.+]] = arith.constant 5
|
|
// CHECK: %[[C40:.+]] = arith.constant 40
|
|
// CHECK: scf.for %{{.+}} = %[[C0]] to %[[C40]] step %[[C5]]
|
|
// CHECK: arith.addi
|
|
// CHECK: scf.for %{{.+}} = %[[C40]] to %[[C42]] step %[[C5]]
|
|
// CHECK: arith.addi
|
|
%0 = arith.constant 0 : index
|
|
%1 = arith.constant 42 : index
|
|
%2 = arith.constant 5 : index
|
|
scf.for %i = %0 to %1 step %2 {
|
|
arith.addi %i, %i : index
|
|
}
|
|
return
|
|
}
|
|
|
|
transform.sequence failures(propagate) {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["arith.addi"]} in %arg1
|
|
%1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
|
|
transform.loop.peel %1 : (!transform.op<"scf.for">) -> !pdl.operation
|
|
}
|
|
|
|
// -----
|
|
|
|
func.func @loop_pipeline_op(%A: memref<?xf32>, %result: memref<?xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
%c1 = arith.constant 1 : index
|
|
%c4 = arith.constant 4 : index
|
|
%cf = arith.constant 1.0 : f32
|
|
// CHECK: memref.load %[[MEMREF:.+]][%{{.+}}]
|
|
// CHECK: memref.load %[[MEMREF]]
|
|
// CHECK: arith.addf
|
|
// CHECK: scf.for
|
|
// CHECK: memref.load
|
|
// CHECK: arith.addf
|
|
// CHECK: memref.store
|
|
// CHECK: arith.addf
|
|
// CHECK: memref.store
|
|
// CHECK: memref.store
|
|
// expected-remark @below {{transformed}}
|
|
scf.for %i0 = %c0 to %c4 step %c1 {
|
|
%A_elem = memref.load %A[%i0] : memref<?xf32>
|
|
%A1_elem = arith.addf %A_elem, %cf : f32
|
|
memref.store %A1_elem, %result[%i0] : memref<?xf32>
|
|
}
|
|
return
|
|
}
|
|
|
|
transform.sequence failures(propagate) {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["arith.addf"]} in %arg1
|
|
%1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
|
|
%2 = transform.loop.pipeline %1 : (!transform.op<"scf.for">) -> !pdl.operation
|
|
// Verify that the returned handle is usable.
|
|
transform.test_print_remark_at_operand %2, "transformed" : !pdl.operation
|
|
}
|
|
|
|
// -----
|
|
|
|
// CHECK-LABEL: @loop_unroll_op
|
|
func.func @loop_unroll_op() {
|
|
%c0 = arith.constant 0 : index
|
|
%c42 = arith.constant 42 : index
|
|
%c5 = arith.constant 5 : index
|
|
// CHECK: scf.for %[[I:.+]] =
|
|
scf.for %i = %c0 to %c42 step %c5 {
|
|
// CHECK-COUNT-4: arith.addi %[[I]]
|
|
arith.addi %i, %i : index
|
|
}
|
|
return
|
|
}
|
|
|
|
transform.sequence failures(propagate) {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["arith.addi"]} in %arg1
|
|
%1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
|
|
transform.loop.unroll %1 { factor = 4 } : !transform.op<"scf.for">
|
|
}
|
|
|
|
// -----
|
|
|
|
// CHECK-LABEL: @get_parent_for_op
|
|
func.func @get_parent_for_op(%arg0: index, %arg1: index, %arg2: index) {
|
|
// expected-remark @below {{first loop}}
|
|
affine.for %i = %arg0 to %arg1 {
|
|
// expected-remark @below {{second loop}}
|
|
affine.for %j = %arg0 to %arg1 {
|
|
// expected-remark @below {{third loop}}
|
|
affine.for %k = %arg0 to %arg1 {
|
|
arith.addi %i, %j : index
|
|
}
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
transform.sequence failures(propagate) {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["arith.addi"]} in %arg1
|
|
// CHECK: = transform.loop.get_parent_for
|
|
%1 = transform.loop.get_parent_for %0 { affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
|
|
%2 = transform.loop.get_parent_for %0 { num_loops = 2, affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
|
|
%3 = transform.loop.get_parent_for %0 { num_loops = 3, affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
|
|
transform.test_print_remark_at_operand %1, "third loop" : !transform.op<"affine.for">
|
|
transform.test_print_remark_at_operand %2, "second loop" : !transform.op<"affine.for">
|
|
transform.test_print_remark_at_operand %3, "first loop" : !transform.op<"affine.for">
|
|
}
|
|
|
|
// -----
|
|
|
|
func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
|
|
// expected-note @below {{target op}}
|
|
arith.addi %arg0, %arg1 : index
|
|
return
|
|
}
|
|
|
|
transform.sequence failures(propagate) {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["arith.addi"]} in %arg1
|
|
// expected-error @below {{could not find an 'affine.for' parent}}
|
|
%1 = transform.loop.get_parent_for %0 { affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
|
|
}
|
|
|
|
// -----
|
|
|
|
func.func @loop_unroll_op() {
|
|
%c0 = arith.constant 0 : index
|
|
%c42 = arith.constant 42 : index
|
|
%c5 = arith.constant 5 : index
|
|
// CHECK: affine.for %[[I:.+]] =
|
|
// expected-remark @below {{affine for loop}}
|
|
affine.for %i = %c0 to %c42 {
|
|
// CHECK-COUNT-4: arith.addi
|
|
arith.addi %i, %i : index
|
|
}
|
|
return
|
|
}
|
|
|
|
transform.sequence failures(propagate) {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["arith.addi"]} in %arg1
|
|
%1 = transform.loop.get_parent_for %0 { affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
|
|
transform.test_print_remark_at_operand %1, "affine for loop" : !transform.op<"affine.for">
|
|
transform.loop.unroll %1 { factor = 4, affine = true } : !transform.op<"affine.for">
|
|
}
|
|
|
|
// -----
|
|
|
|
func.func @test_mixed_loops() {
|
|
%c0 = arith.constant 0 : index
|
|
%c42 = arith.constant 42 : index
|
|
%c5 = arith.constant 5 : index
|
|
scf.for %j = %c0 to %c42 step %c5 {
|
|
// CHECK: affine.for %[[I:.+]] =
|
|
// expected-remark @below {{affine for loop}}
|
|
affine.for %i = %c0 to %c42 {
|
|
// CHECK-COUNT-4: arith.addi
|
|
arith.addi %i, %i : index
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
transform.sequence failures(propagate) {
|
|
^bb1(%arg1: !pdl.operation):
|
|
%0 = transform.structured.match ops{["arith.addi"]} in %arg1
|
|
%1 = transform.loop.get_parent_for %0 { num_loops = 1, affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
|
|
transform.test_print_remark_at_operand %1, "affine for loop" : !transform.op<"affine.for">
|
|
transform.loop.unroll %1 { factor = 4 } : !transform.op<"affine.for">
|
|
}
|