1// RUN: mlir-opt %s -allow-unregistered-dialect \ 2// RUN: -transform-interpreter -canonicalize \ 3// RUN: -split-input-file -verify-diagnostics | FileCheck %s 4 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 : f32 25 } : tensor<?xf32> to tensor<?xf32> 26 // CHECK: "dummy.some_use"(%[[replacement]]) 27 "dummy.some_use"(%p) : (tensor<?xf32>) -> () 28 } 29 return 30} 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_op 35 %1 = transform.tensor.make_loop_independent %0 {num_loops = 1} : (!transform.any_op) -> !transform.any_op 36 transform.yield 37 } 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 : f32 64 } : tensor<?xf32> to tensor<?xf32> 65 // CHECK: "dummy.some_use"(%[[replacement]]) 66 "dummy.some_use"(%p) : (tensor<?xf32>) -> () 67 } 68 return 69} 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_op 74 %1 = transform.tensor.make_loop_independent %0 {num_loops = 1} : (!transform.any_op) -> !transform.any_op 75 transform.yield 76 } 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 : f32 92 } : tensor<?xf32> to tensor<?xf32> 93 "dummy.some_use"(%p) : (tensor<?xf32>) -> () 94 } 95 } 96 return 97} 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_op 102 %1 = transform.tensor.make_loop_independent %0 {num_loops = 2} : (!transform.any_op) -> !transform.any_op 103 transform.yield 104 } 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 : f32 118 } : tensor<?xf32> to tensor<?xf32> 119 "dummy.some_use"(%p) : (tensor<?xf32>) -> () 120 } 121 } 122 return 123} 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_op 128 // 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_op 130 transform.yield 131 } 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 return 153} 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_op 158 %1 = transform.tensor.make_loop_independent %0 {num_loops = 1} : (!transform.any_op) -> !transform.any_op 159 transform.yield 160 } 161} 162