derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
417 lines
21 KiB
MLIR
417 lines
21 KiB
MLIR
// RUN: mlir-opt %s -transform-interpreter -split-input-file -verify-diagnostics -allow-unregistered-dialect --cse --mlir-print-local-scope | FileCheck %s
|
|
|
|
func.func @coalesce_inner() {
|
|
%c0 = arith.constant 0 : index
|
|
%c1 = arith.constant 1 : index
|
|
%c10 = arith.constant 10 : index
|
|
|
|
// CHECK: scf.for %[[IV0:.+]]
|
|
// CHECK: scf.for %[[IV1:.+]]
|
|
// CHECK: scf.for %[[IV2:.+]]
|
|
// CHECK-NOT: scf.for %[[IV3:.+]]
|
|
scf.for %i = %c0 to %c10 step %c1 {
|
|
scf.for %j = %c0 to %c10 step %c1 {
|
|
scf.for %k = %i to %j step %c1 {
|
|
// Inner loop must have been removed.
|
|
scf.for %l = %i to %j step %c1 {
|
|
"use"(%i, %j) : (index, index) -> ()
|
|
}
|
|
} {coalesce}
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
|
|
%2 = transform.loop.coalesce %1: (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
func.func @coalesce_outer(%arg1: memref<64x64xf32, 1>, %arg2: memref<64x64xf32, 1>, %arg3: memref<64x64xf32, 1>) attributes {} {
|
|
// CHECK: %[[T0:.+]] = affine.apply affine_map<() -> (64)>()
|
|
// CHECK: %[[UB:.+]] = affine.apply affine_map<(d0)[s0] -> (d0 * s0)>(%[[T0]])[%[[T0]]]
|
|
// CHECK: affine.for %[[IV1:.+]] = 0 to %[[UB:.+]] {
|
|
// CHECK-NOT: affine.for %[[IV2:.+]]
|
|
affine.for %arg4 = 0 to 64 {
|
|
affine.for %arg5 = 0 to 64 {
|
|
// CHECK: %[[IDX0:.+]] = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)>(%[[IV1]])[%{{.+}}]
|
|
// CHECK: %[[IDX1:.+]] = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)>(%[[IV1]])[%{{.+}}]
|
|
// CHECK-NEXT: %{{.+}} = affine.load %{{.+}}[%[[IDX1]], %[[IDX0]]] : memref<64x64xf32, 1>
|
|
%0 = affine.load %arg1[%arg4, %arg5] : memref<64x64xf32, 1>
|
|
%1 = affine.load %arg2[%arg4, %arg5] : memref<64x64xf32, 1>
|
|
%2 = arith.addf %0, %1 : f32
|
|
affine.store %2, %arg3[%arg4, %arg5] : memref<64x64xf32, 1>
|
|
}
|
|
} {coalesce}
|
|
return
|
|
}
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["affine.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.cast %0 : !transform.any_op to !transform.op<"affine.for">
|
|
%2 = transform.loop.coalesce %1 : (!transform.op<"affine.for">) -> (!transform.op<"affine.for">)
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
func.func @coalesce_and_unroll(%arg1: memref<64x64xf32, 1>, %arg2: memref<64x64xf32, 1>, %arg3: memref<64x64xf32, 1>) attributes {} {
|
|
// CHECK: scf.for %[[IV1:.+]] =
|
|
%c0 = arith.constant 0 : index
|
|
%c1 = arith.constant 1 : index
|
|
%c64 = arith.constant 64 : index
|
|
|
|
scf.for %arg4 = %c0 to %c64 step %c1 {
|
|
// CHECK-NOT: scf.for
|
|
scf.for %arg5 = %c0 to %c64 step %c1 {
|
|
// CHECK: %[[IDX:.+]]:2 = affine.delinearize_index
|
|
// CHECK-NEXT: %{{.+}} = memref.load %{{.+}}[%[[IDX]]#0, %[[IDX]]#1] : memref<64x64xf32, 1>
|
|
%0 = memref.load %arg1[%arg4, %arg5] : memref<64x64xf32, 1>
|
|
%1 = memref.load %arg2[%arg4, %arg5] : memref<64x64xf32, 1>
|
|
%2 = arith.addf %0, %1 : f32
|
|
// CHECK: memref.store
|
|
// CHECK: memref.store
|
|
// CHECK: memref.store
|
|
// Residual loop must have a single store.
|
|
// CHECK: memref.store
|
|
memref.store %2, %arg3[%arg4, %arg5] : memref<64x64xf32, 1>
|
|
}
|
|
} {coalesce}
|
|
return
|
|
}
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
|
|
%2 = transform.loop.coalesce %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
|
|
transform.loop.unroll %2 {factor = 3} : !transform.op<"scf.for">
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
func.func @tensor_loops(%arg0 : tensor<?x?xf32>, %lb0 : index, %ub0 : index, %step0 : index,
|
|
%lb1 : index, %ub1 : index, %step1 : index, %lb2 : index, %ub2 : index, %step2 : index) -> tensor<?x?xf32> {
|
|
%0 = scf.for %i = %lb0 to %ub0 step %step0 iter_args(%arg1 = %arg0) -> tensor<?x?xf32> {
|
|
%1 = scf.for %j = %lb1 to %ub1 step %step1 iter_args(%arg2 = %arg1) -> tensor<?x?xf32> {
|
|
%2 = scf.for %k = %lb2 to %ub2 step %step2 iter_args(%arg3 = %arg2) -> tensor<?x?xf32> {
|
|
%3 = "use"(%arg3, %i, %j, %k) : (tensor<?x?xf32>, index, index, index) -> (tensor<?x?xf32>)
|
|
scf.yield %3 : tensor<?x?xf32>
|
|
}
|
|
scf.yield %2 : tensor<?x?xf32>
|
|
}
|
|
scf.yield %1 : tensor<?x?xf32>
|
|
} {coalesce}
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
|
|
%2 = transform.loop.coalesce %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
|
|
transform.yield
|
|
}
|
|
}
|
|
// CHECK: func.func @tensor_loops(
|
|
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>
|
|
// CHECK-SAME: %[[LB0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[LB1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[LB2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK: %[[NITERS0:.+]] = affine.apply
|
|
// CHECK-SAME: affine_map<()[s0, s1, s2] -> ((-s0 + s1) ceildiv s2)>()[%[[LB0]], %[[UB0]], %[[STEP0]]]
|
|
// CHECK: %[[C0:.+]] = arith.constant 0 : index
|
|
// CHECK: %[[C1:.+]] = arith.constant 1 : index
|
|
// CHECK: %[[NITERS1:.+]] = affine.apply
|
|
// CHECK-SAME: affine_map<()[s0, s1, s2] -> ((-s0 + s1) ceildiv s2)>()[%[[LB1]], %[[UB1]], %[[STEP1]]]
|
|
// CHECK: %[[NITERS2:.+]] = affine.apply
|
|
// CHECK-SAME: affine_map<()[s0, s1, s2] -> ((-s0 + s1) ceildiv s2)>()[%[[LB2]], %[[UB2]], %[[STEP2]]]
|
|
// CHECK: %[[NEWUB:.+]] = affine.apply affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8] ->
|
|
// CHECK-SAME: ((((-s0 + s1) ceildiv s2) * ((-s3 + s4) ceildiv s5)) * ((-s6 + s7) ceildiv s8))
|
|
// CHECK-SAME: [%[[LB0]], %[[UB0]], %[[STEP0]], %[[LB1]], %[[UB1]], %[[STEP1]], %[[LB2]], %[[UB2]], %[[STEP2]]]
|
|
// CHECK: %[[RESULT:.+]] = scf.for %[[IV:[a-zA-Z0-9]+]] = %[[C0]] to %[[NEWUB]] step %[[C1]] iter_args(%[[ITER_ARG:.+]] = %[[ARG0]])
|
|
// CHECK: %[[DELINEARIZE:.+]]:3 = affine.delinearize_index %[[IV]] into (%[[NITERS0]], %[[NITERS1]], %[[NITERS2]])
|
|
// CHECK-DAG: %[[K:.+]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>(%[[DELINEARIZE]]#2)[%[[LB2]], %[[STEP2]]]
|
|
// CHECK-DAG: %[[J:.+]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>(%[[DELINEARIZE]]#1)[%[[LB1]], %[[STEP1]]]
|
|
// CHECK-DAG: %[[I:.+]] = affine.apply affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>(%[[DELINEARIZE]]#0)[%[[LB0]], %[[STEP0]]]
|
|
// CHECK: %[[USE:.+]] = "use"(%[[ITER_ARG]], %[[I]], %[[J]], %[[K]])
|
|
// CHECK: scf.yield %[[USE]]
|
|
// CHECK: return %[[RESULT]]
|
|
|
|
// -----
|
|
|
|
// Coalesce only first two loops, but not the last since the iter_args dont line up
|
|
func.func @tensor_loops_first_two(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %lb0 : index, %ub0 : index, %step0 : index,
|
|
%lb1 : index, %ub1 : index, %step1 : index, %lb2 : index, %ub2 : index, %step2 : index) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%0:2 = scf.for %i = %lb0 to %ub0 step %step0 iter_args(%arg2 = %arg0, %arg3 = %arg1) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%1:2 = scf.for %j = %lb1 to %ub1 step %step1 iter_args(%arg4 = %arg2, %arg5 = %arg3) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%2:2 = scf.for %k = %lb2 to %ub2 step %step2 iter_args(%arg6 = %arg5, %arg7 = %arg4) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%3:2 = "use"(%arg3, %i, %j, %k) : (tensor<?x?xf32>, index, index, index) -> (tensor<?x?xf32>, tensor<?x?xf32>)
|
|
scf.yield %3#0, %3#1 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
}
|
|
scf.yield %2#0, %2#1 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
}
|
|
scf.yield %1#0, %1#1 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
} {coalesce}
|
|
return %0#0, %0#1 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
}
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
|
|
%2 = transform.loop.coalesce %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
|
|
transform.yield
|
|
}
|
|
}
|
|
// CHECK: func.func @tensor_loops_first_two(
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
|
|
// CHECK-SAME: %[[LB0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[LB1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[LB2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK: scf.for
|
|
// CHECK: affine.delinearize_index
|
|
// CHECK: scf.for %{{[a-zA-Z0-9]+}} = %[[LB2]] to %[[UB2]] step %[[STEP2]]
|
|
// CHECK-NOT: scf.for
|
|
// CHECK: transform.named_sequence
|
|
|
|
// -----
|
|
|
|
// Coalesce only first two loops, but not the last since the yields dont match up
|
|
func.func @tensor_loops_first_two_2(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %lb0 : index, %ub0 : index, %step0 : index,
|
|
%lb1 : index, %ub1 : index, %step1 : index, %lb2 : index, %ub2 : index, %step2 : index) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%0:2 = scf.for %i = %lb0 to %ub0 step %step0 iter_args(%arg2 = %arg0, %arg3 = %arg1) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%1:2 = scf.for %j = %lb1 to %ub1 step %step1 iter_args(%arg4 = %arg2, %arg5 = %arg3) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%2:2 = scf.for %k = %lb2 to %ub2 step %step2 iter_args(%arg6 = %arg4, %arg7 = %arg5) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%3:2 = "use"(%arg3, %i, %j, %k) : (tensor<?x?xf32>, index, index, index) -> (tensor<?x?xf32>, tensor<?x?xf32>)
|
|
scf.yield %3#0, %3#1 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
}
|
|
scf.yield %2#1, %2#0 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
}
|
|
scf.yield %1#0, %1#1 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
} {coalesce}
|
|
return %0#0, %0#1 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
}
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
|
|
%2 = transform.loop.coalesce %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
|
|
transform.yield
|
|
}
|
|
}
|
|
// CHECK: func.func @tensor_loops_first_two_2(
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
|
|
// CHECK-SAME: %[[LB0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[LB1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[LB2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK: scf.for
|
|
// CHECK: affine.delinearize_index
|
|
// CHECK: scf.for %{{[a-zA-Z0-9]+}} = %[[LB2]] to %[[UB2]] step %[[STEP2]]
|
|
// CHECK-NOT: scf.for
|
|
// CHECK: transform.named_sequence
|
|
|
|
// -----
|
|
|
|
// Coalesce only last two loops, but not the first since the yields dont match up
|
|
func.func @tensor_loops_last_two(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %lb0 : index, %ub0 : index, %step0 : index,
|
|
%lb1 : index, %ub1 : index, %step1 : index, %lb2 : index, %ub2 : index, %step2 : index) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%0:2 = scf.for %i = %lb0 to %ub0 step %step0 iter_args(%arg2 = %arg0, %arg3 = %arg1) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%1:2 = scf.for %j = %lb1 to %ub1 step %step1 iter_args(%arg4 = %arg2, %arg5 = %arg3) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%2:2 = scf.for %k = %lb2 to %ub2 step %step2 iter_args(%arg6 = %arg4, %arg7 = %arg5) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
|
|
%3:2 = "use"(%arg3, %i, %j, %k) : (tensor<?x?xf32>, index, index, index) -> (tensor<?x?xf32>, tensor<?x?xf32>)
|
|
scf.yield %3#0, %3#1 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
}
|
|
scf.yield %2#0, %2#1 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
}
|
|
scf.yield %1#1, %1#0 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
} {coalesce}
|
|
return %0#0, %0#1 : tensor<?x?xf32>, tensor<?x?xf32>
|
|
}
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
|
|
%2 = transform.loop.coalesce %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
|
|
transform.yield
|
|
}
|
|
}
|
|
// CHECK: func.func @tensor_loops_last_two(
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
|
|
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
|
|
// CHECK-SAME: %[[LB0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP0:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[LB1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP1:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[LB2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[UB2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK-SAME: %[[STEP2:[a-zA-Z0-9_]+]]: index
|
|
// CHECK: scf.for %{{[a-zA-Z0-9]+}} = %[[LB0]] to %[[UB0]] step %[[STEP0]]
|
|
// CHECK-NOT: affine.delinearize_index
|
|
// CHECK: scf.for
|
|
// CHECK: affine.delinearize_index
|
|
// CHECK-NOT: scf.for
|
|
// CHECK: transform.named_sequence
|
|
|
|
// -----
|
|
|
|
// Check avoiding generating unnecessary operations while collapsing trip-1 loops.
|
|
func.func @trip_one_loops(%arg0 : tensor<?x?xf32>, %arg1 : index, %arg2 : index) -> tensor<?x?xf32> {
|
|
%c0 = arith.constant 0 : index
|
|
%c1 = arith.constant 1 : index
|
|
%0 = scf.for %iv0 = %c0 to %c1 step %c1 iter_args(%iter0 = %arg0) -> tensor<?x?xf32> {
|
|
%1 = scf.for %iv1 = %c0 to %c1 step %c1 iter_args(%iter1 = %iter0) -> tensor<?x?xf32> {
|
|
%2 = scf.for %iv2 = %c0 to %arg1 step %c1 iter_args(%iter2 = %iter1) -> tensor<?x?xf32> {
|
|
%3 = scf.for %iv3 = %c0 to %c1 step %c1 iter_args(%iter3 = %iter2) -> tensor<?x?xf32> {
|
|
%4 = scf.for %iv4 = %c0 to %arg2 step %c1 iter_args(%iter4 = %iter3) -> tensor<?x?xf32> {
|
|
%5 = "some_use"(%iter4, %iv0, %iv1, %iv2, %iv3, %iv4)
|
|
: (tensor<?x?xf32>, index, index, index, index, index) -> (tensor<?x?xf32>)
|
|
scf.yield %5 : tensor<?x?xf32>
|
|
}
|
|
scf.yield %4 : tensor<?x?xf32>
|
|
}
|
|
scf.yield %3 : tensor<?x?xf32>
|
|
}
|
|
scf.yield %2 : tensor<?x?xf32>
|
|
}
|
|
scf.yield %1 : tensor<?x?xf32>
|
|
} {coalesce}
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
|
|
%2 = transform.loop.coalesce %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
|
|
transform.apply_patterns to %2 {
|
|
transform.apply_patterns.canonicalization
|
|
} : !transform.op<"scf.for">
|
|
transform.yield
|
|
}
|
|
}
|
|
// CHECK-LABEL: func @trip_one_loops
|
|
// CHECK-SAME: , %[[ARG1:.+]]: index,
|
|
// CHECK-SAME: %[[ARG2:.+]]: index)
|
|
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
|
|
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
|
|
// CHECK: %[[UB:.+]] = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%[[ARG1]], %[[ARG2]]]
|
|
// CHECK: scf.for %[[IV:.+]] = %[[C0]] to %[[UB]] step %[[C1]]
|
|
// CHECK: %[[DELINEARIZE:.+]]:2 = affine.delinearize_index %[[IV]] into (%[[ARG1]], %[[ARG2]])
|
|
// CHECK: "some_use"(%{{[a-zA-Z0-9]+}}, %[[C0]], %[[C0]], %[[DELINEARIZE]]#0, %[[C0]], %[[DELINEARIZE]]#1)
|
|
|
|
// -----
|
|
|
|
// Check generating no instructions when all except one loops is non unit-trip.
|
|
func.func @all_outer_trip_one(%arg0 : tensor<?x?xf32>, %arg1 : index) -> tensor<?x?xf32> {
|
|
%c0 = arith.constant 0 : index
|
|
%c1 = arith.constant 1 : index
|
|
%0 = scf.for %iv0 = %c0 to %c1 step %c1 iter_args(%iter0 = %arg0) -> tensor<?x?xf32> {
|
|
%1 = scf.for %iv1 = %c0 to %c1 step %c1 iter_args(%iter1 = %iter0) -> tensor<?x?xf32> {
|
|
%2 = scf.for %iv2 = %c0 to %arg1 step %c1 iter_args(%iter2 = %iter1) -> tensor<?x?xf32> {
|
|
%3 = "some_use"(%iter2, %iv0, %iv1, %iv2)
|
|
: (tensor<?x?xf32>, index, index, index) -> (tensor<?x?xf32>)
|
|
scf.yield %3 : tensor<?x?xf32>
|
|
}
|
|
scf.yield %2 : tensor<?x?xf32>
|
|
}
|
|
scf.yield %1 : tensor<?x?xf32>
|
|
} {coalesce}
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["scf.for"]} attributes {coalesce} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
|
|
%2 = transform.loop.coalesce %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
|
|
transform.apply_patterns to %2 {
|
|
transform.apply_patterns.canonicalization
|
|
} : !transform.op<"scf.for">
|
|
transform.yield
|
|
}
|
|
}
|
|
// CHECK-LABEL: func @all_outer_trip_one
|
|
// CHECK-SAME: , %[[ARG1:.+]]: index)
|
|
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
|
|
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
|
|
// CHECK: scf.for %[[IV:.+]] = %[[C0]] to %[[ARG1]] step %[[C1]]
|
|
// CHECK: "some_use"(%{{[a-zA-Z0-9]+}}, %[[C0]], %[[C0]], %[[IV]])
|
|
|
|
// -----
|
|
|
|
// Loops with operations between them that use the outer loop's induction variable.
|
|
// The delinearization should be inserted early enough that these operations can
|
|
// use the delinearized IV.
|
|
|
|
func.func @opsinbetween_nested_loops(%init: f32) -> f32 {
|
|
%c0 = arith.constant 0 : index
|
|
%c1 = arith.constant 1 : index
|
|
%c3 = arith.constant 3 : index
|
|
%c5 = arith.constant 5 : index
|
|
|
|
%result = scf.for %i = %c0 to %c3 step %c1 iter_args(%outer = %init) -> (f32) {
|
|
%computed = arith.addi %i, %c5 : index
|
|
%used = arith.muli %computed, %c3 : index
|
|
|
|
%inner_result = scf.for %j = %c0 to %c3 step %c1 iter_args(%inner = %outer) -> (f32) {
|
|
%updated = "use"(%inner, %i, %j, %computed, %used) : (f32, index, index, index, index) -> f32
|
|
scf.yield %updated : f32
|
|
}
|
|
scf.yield %inner_result : f32
|
|
} {coalesce_nested}
|
|
return %result : f32
|
|
}
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%0 = transform.structured.match ops{["scf.for"]} attributes {coalesce_nested} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
%1 = transform.cast %0 : !transform.any_op to !transform.op<"scf.for">
|
|
%2 = transform.loop.coalesce_nested %1 : (!transform.op<"scf.for">) -> (!transform.op<"scf.for">)
|
|
transform.yield
|
|
}
|
|
}
|
|
// CHECK-LABEL: func @opsinbetween_nested_loops
|
|
// CHECK-SAME: %[[INIT:.+]]: f32
|
|
// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
|
|
// CHECK-DAG: %[[C5:.+]] = arith.constant 5 : index
|
|
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
|
|
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
|
|
// CHECK-DAG: %[[C9:.+]] = arith.constant 9 : index
|
|
// CHECK: %[[RESULT:.+]] = scf.for %[[IV:.+]] = %[[C0]] to %[[C9]] step %[[C1]] iter_args(%[[ARG:.+]] = %[[INIT]]) -> (f32) {
|
|
// CHECK: %[[DELIN:.+]]:2 = affine.delinearize_index %[[IV]] into (3, 3)
|
|
// CHECK: %[[COMPUTED:.+]] = arith.addi %[[DELIN]]#0, %[[C5]]
|
|
// CHECK: %[[USED:.+]] = arith.muli %[[COMPUTED]], %[[C3]]
|
|
// CHECK: %[[UPDATED:.+]] = "use"(%[[ARG]], %[[DELIN]]#0, %[[DELIN]]#1, %[[COMPUTED]], %[[USED]])
|
|
// CHECK: scf.yield %[[UPDATED]]
|
|
// CHECK: }
|
|
// CHECK: return %[[RESULT]]
|