162 lines · plain
1// RUN: mlir-opt %s -allow-unregistered-dialect \2// RUN: -transform-interpreter -canonicalize \3// RUN: -split-input-file -verify-diagnostics | FileCheck %s4 5// This is a test case where "high" padding depends on the IV.6 7// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 - s1)>8// CHECK: #[[$map1:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 + 5)>9// CHECK-LABEL: func @make_pad_loop_independent_1(10// CHECK-SAME: %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index,11// CHECK-SAME: %[[t:.*]]: tensor<?xf32>12func.func @make_pad_loop_independent_1(%lb: index, %ub: index, %step: index,13 %t: tensor<?xf32>, %f: f32) {14 // CHECK: scf.for %[[iv:.*]] = %[[lb]] to %[[ub]]15 scf.for %i = %lb to %ub step %step {16 // CHECK: %[[high:.*]] = affine.apply #[[$map]]()[%[[ub]], %[[lb]]]17 // CHECK: %[[padded:.*]] = tensor.pad %[[t]] low[5] high[%[[high]]]18 // CHECK: %[[dim:.*]] = tensor.dim %[[t]]19 // CHECK: %[[size:.*]] = affine.apply #[[$map1]](%[[iv]])[%[[ub]], %[[dim]]]20 // CHECK: %[[replacement:.*]] = tensor.extract_slice %[[padded]][0] [%[[size]]] [1]21 %high = affine.apply affine_map<(d0)[s0] -> (s0 - d0)> (%i)[%ub]22 %p = tensor.pad %t low[5] high[%high] {23 ^bb0(%arg1: index):24 tensor.yield %f : f3225 } : tensor<?xf32> to tensor<?xf32>26 // CHECK: "dummy.some_use"(%[[replacement]])27 "dummy.some_use"(%p) : (tensor<?xf32>) -> ()28 }29 return30}31 32module attributes {transform.with_named_sequence} {33 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {34 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op35 %1 = transform.tensor.make_loop_independent %0 {num_loops = 1} : (!transform.any_op) -> !transform.any_op36 transform.yield37 }38}39 40// -----41 42// This is a test case where "low" padding depends on the IV.43 44// CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 - s1)>45// CHECK: #[[$map1:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 + 5)>46// CHECK: #[[$map2:.*]] = affine_map<(d0)[s0] -> (d0 - s0)>47// CHECK-LABEL: func @make_pad_loop_independent_1(48// CHECK-SAME: %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index,49// CHECK-SAME: %[[t:.*]]: tensor<?xf32>50func.func @make_pad_loop_independent_1(%lb: index, %ub: index, %step: index,51 %t: tensor<?xf32>, %f: f32) {52 // CHECK: scf.for %[[iv:.*]] = %[[lb]] to %[[ub]]53 scf.for %i = %lb to %ub step %step {54 // CHECK: %[[low:.*]] = affine.apply #[[$map]]()[%[[ub]], %[[lb]]]55 // CHECK: %[[padded:.*]] = tensor.pad %[[t]] low[%[[low]]] high[5]56 // CHECK: %[[dim:.*]] = tensor.dim %[[t]]57 // CHECK: %[[size:.*]] = affine.apply #[[$map1]](%[[iv]])[%[[ub]], %[[dim]]]58 // CHECK: %[[offset:.*]] = affine.apply #[[$map2]](%[[iv]])[%[[lb]]]59 // CHECK: %[[replacement:.*]] = tensor.extract_slice %[[padded]][%[[offset]]] [%[[size]]] [1]60 %low = affine.apply affine_map<(d0)[s0] -> (s0 - d0)> (%i)[%ub]61 %p = tensor.pad %t low[%low] high[5] {62 ^bb0(%arg1: index):63 tensor.yield %f : f3264 } : tensor<?xf32> to tensor<?xf32>65 // CHECK: "dummy.some_use"(%[[replacement]])66 "dummy.some_use"(%p) : (tensor<?xf32>) -> ()67 }68 return69}70 71module attributes {transform.with_named_sequence} {72 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {73 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op74 %1 = transform.tensor.make_loop_independent %0 {num_loops = 1} : (!transform.any_op) -> !transform.any_op75 transform.yield76 }77}78 79// -----80 81// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 * 2 - 2)>82// CHECK-LABEL: func @two_loops(83func.func @two_loops(%lb: index, %ub: index, %step: index,84 %t: tensor<?xf32>, %f: f32) {85 scf.for %i = %lb to %ub step %step {86 scf.for %j = %lb to %ub step %step {87 // CHECK: affine.apply #map()[%{{.*}}]88 %low = affine.apply affine_map<(d0, d1)[] -> (d0 + d1)> (%i, %j)[]89 %p = tensor.pad %t low[%low] high[5] {90 ^bb0(%arg1: index):91 tensor.yield %f : f3292 } : tensor<?xf32> to tensor<?xf32>93 "dummy.some_use"(%p) : (tensor<?xf32>) -> ()94 }95 }96 return97}98 99module attributes {transform.with_named_sequence} {100 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {101 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op102 %1 = transform.tensor.make_loop_independent %0 {num_loops = 2} : (!transform.any_op) -> !transform.any_op103 transform.yield104 }105}106 107// -----108 109func.func @not_enough_loops(%lb: index, %ub: index, %step: index,110 %t: tensor<?xf32>, %f: f32) {111 scf.for %i = %lb to %ub step %step {112 scf.for %j = %lb to %ub step %step {113 %low = affine.apply affine_map<(d0, d1)[] -> (d0 + d1)> (%i, %j)[]114 // expected-note@below {{target op}}115 %p = tensor.pad %t low[%low] high[5] {116 ^bb0(%arg1: index):117 tensor.yield %f : f32118 } : tensor<?xf32> to tensor<?xf32>119 "dummy.some_use"(%p) : (tensor<?xf32>) -> ()120 }121 }122 return123}124 125module attributes {transform.with_named_sequence} {126 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {127 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op128 // expected-error@below {{could not find 2-th enclosing loop}}129 %1 = transform.tensor.make_loop_independent %0 {num_loops = 3} : (!transform.any_op) -> !transform.any_op130 transform.yield131 }132}133 134// -----135 136// CHECK: #[[$map:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)>137// CHECK: #[[$map1:.*]] = affine_map<()[s0, s1] -> (s0 - s1)>138// CHECK-LABEL: func @make_empty_loop_independent(139// CHECK-SAME: %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index)140func.func @make_empty_loop_independent(%lb: index, %ub: index, %step: index) {141 // CHECK: scf.for %[[iv:.*]] = %[[lb]] to %[[ub]]142 scf.for %i = %lb to %ub step %step {143 // CHECK: %[[slice_sz:.*]] = affine.apply #[[$map]](%[[iv]])[%[[ub]]]144 // CHECK: %[[empty_sz:.*]] = affine.apply #[[$map1]]()[%[[ub]], %[[lb]]]145 // CHECK: %[[empty:.*]] = tensor.empty(%[[empty_sz]]) : tensor<?xf32>146 // CHECK: %[[replacement:.*]] = tensor.extract_slice %[[empty]][0] [%[[slice_sz]]] [1]147 %sz = affine.apply affine_map<(d0)[s0] -> (s0 - d0)> (%i)[%ub]148 %empty = tensor.empty(%sz) : tensor<?xf32>149 // CHECK: "dummy.some_use"(%[[replacement]])150 "dummy.some_use"(%empty) : (tensor<?xf32>) -> ()151 }152 return153}154 155module attributes {transform.with_named_sequence} {156 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {157 %0 = transform.structured.match ops{["tensor.empty"]} in %arg1 : (!transform.any_op) -> !transform.any_op158 %1 = transform.tensor.make_loop_independent %0 {num_loops = 1} : (!transform.any_op) -> !transform.any_op159 transform.yield160 }161}162