| // RUN: mlir-opt %s -allow-unregistered-dialect -test-linalg-transform-patterns=test-tiled-loop-peeling=2 -split-input-file | FileCheck %s -check-prefix=CHECK-TILE-2 |
| // RUN: mlir-opt %s -allow-unregistered-dialect -test-linalg-transform-patterns=test-tiled-loop-peeling=0,1,2 -split-input-file | FileCheck %s -check-prefix=CHECK-TILE-012 |
| // RUN: mlir-opt %s -allow-unregistered-dialect -test-linalg-transform-patterns="test-tiled-loop-peeling=0,1,2 skip-partial" -split-input-file | FileCheck %s -check-prefix=CHECK-TILE-012-SKIP-PARTIAL |
| |
| // CHECK-TILE-2-LABEL: func @tiled_loop_3d_tensor( |
| // CHECK-TILE-2-SAME: %[[input:.*]]: tensor<?x?x?xf32>, %[[s0:.*]]: index, %[[s1:.*]]: index, %[[s2:.*]]: index |
| // CHECK-TILE-2-DAG: %[[c0:.*]] = arith.constant 0 : index |
| // CHECK-TILE-2-DAG: %[[c1:.*]] = arith.constant 1 : index |
| // CHECK-TILE-2-DAG: %[[c2:.*]] = arith.constant 2 : index |
| // CHECK-TILE-2: %[[dim0:.*]] = tensor.dim %[[input]], %[[c0]] |
| // CHECK-TILE-2: %[[dim1:.*]] = tensor.dim %[[input]], %[[c1]] |
| // CHECK-TILE-2: %[[dim2:.*]] = tensor.dim %[[input]], %[[c2]] |
| // CHECK-TILE-2: %[[init_tensor:.*]] = linalg.init_tensor |
| // CHECK-TILE-2: %[[split_bound:.*]] = affine.apply |
| // CHECK-TILE-2: %[[r1:.*]] = linalg.tiled_loop (%[[iv0:.*]], %[[iv1:.*]], %[[iv2:.*]]) = (%[[c0]], %[[c0]], %[[c0]]) |
| // CHECK-TILE-2-SAME: to (%[[dim0]], %[[dim1]], %[[split_bound]]) |
| // CHECK-TILE-2-SAME: step (%[[s0]], %[[s1]], %[[s2]]) |
| // CHECK-TILE-2-SAME: ins (%[[loop_in1:.*]] = %[[input]]: tensor<?x?x?xf32>) |
| // CHECK-TILE-2-SAME: outs (%[[loop_out1:.*]] = %[[init_tensor]]: tensor<?x?x?xf32>) { |
| // CHECK-TILE-2: %[[min0_1:.*]] = affine.min |
| // CHECK-TILE-2: %[[min1_1:.*]] = affine.min |
| // CHECK-TILE-2: %[[in_slice1:.*]] = tensor.extract_slice %[[loop_in1]][%[[iv0]], %[[iv1]], %[[iv2]]] [%[[min0_1]], %[[min1_1]], %[[s2]]] |
| // CHECK-TILE-2: %[[out_slice1:.*]] = tensor.extract_slice %[[loop_out1]][%[[iv0]], %[[iv1]], %[[iv2]]] [%[[min0_1]], %[[min1_1]], %[[s2]]] |
| // CHECK-TILE-2: %[[mod_slice1:.*]] = tensor.insert_slice %{{.*}} into %[[loop_out1]][%[[iv0]], %[[iv1]], %[[iv2]]] [%[[min0_1]], %[[min1_1]], %[[s2]]] |
| // CHECK-TILE-2: linalg.yield %[[mod_slice1]] |
| // CHECK-TILE-2: %[[r2:.*]] = linalg.tiled_loop (%[[iv0:.*]], %[[iv1:.*]], %[[iv2:.*]]) = (%[[c0]], %[[c0]], %[[split_bound]]) |
| // CHECK-TILE-2-SAME: to (%[[dim0]], %[[dim1]], %[[dim2]]) |
| // CHECK-TILE-2-SAME: step (%[[s0]], %[[s1]], %[[s2]]) |
| // CHECK-TILE-2-SAME: ins (%[[loop_in2:.*]] = %[[input]]: tensor<?x?x?xf32>) |
| // CHECK-TILE-2-SAME: outs (%[[loop_out2:.*]] = %[[r1]]: tensor<?x?x?xf32>) { |
| // CHECK-TILE-2: %[[min0_2:.*]] = affine.min |
| // CHECK-TILE-2: %[[min1_2:.*]] = affine.min |
| // CHECK-TILE-2: %[[apply2:.*]] = affine.apply |
| // CHECK-TILE-2: %[[in_slice2:.*]] = tensor.extract_slice %[[loop_in1]][%[[iv0]], %[[iv1]], %[[iv2]]] [%[[min0_2]], %[[min1_2]], %[[apply2]]] |
| // CHECK-TILE-2: %[[out_slice2:.*]] = tensor.extract_slice %[[loop_out1]][%[[iv0]], %[[iv1]], %[[iv2]]] [%[[min0_2]], %[[min1_2]], %[[apply2]]] |
| // CHECK-TILE-2: %[[mod_slice2:.*]] = tensor.insert_slice %{{.*}} into %[[loop_out1]][%[[iv0]], %[[iv1]], %[[iv2]]] [%[[min0_2]], %[[min1_2]], %[[apply2]]] |
| // CHECK-TILE-2: linalg.yield %[[mod_slice2]] |
| // CHECK-TILE-2: return %[[r2]] |
| |
| // CHECK-TILE-012-LABEL: func @tiled_loop_3d_tensor |
| // CHECK-TILE-012: linalg.tiled_loop {{.*}} { |
| // CHECK-TILE-012: linalg.yield |
| // CHECK-TILE-012: } |
| // CHECK-TILE-012: linalg.tiled_loop {{.*}} { |
| // CHECK-TILE-012: linalg.yield |
| // CHECK-TILE-012: } |
| // CHECK-TILE-012: linalg.tiled_loop {{.*}} { |
| // CHECK-TILE-012: linalg.yield |
| // CHECK-TILE-012: } |
| // CHECK-TILE-012: linalg.tiled_loop {{.*}} { |
| // CHECK-TILE-012: linalg.yield |
| // CHECK-TILE-012: } |
| // CHECK-TILE-012: linalg.tiled_loop {{.*}} { |
| // CHECK-TILE-012: linalg.yield |
| // CHECK-TILE-012: } |
| // CHECK-TILE-012: linalg.tiled_loop {{.*}} { |
| // CHECK-TILE-012: linalg.yield |
| // CHECK-TILE-012: } |
| // CHECK-TILE-012: linalg.tiled_loop {{.*}} { |
| // CHECK-TILE-012: linalg.yield |
| // CHECK-TILE-012: } |
| // CHECK-TILE-012: linalg.tiled_loop {{.*}} { |
| // CHECK-TILE-012: linalg.yield |
| // CHECK-TILE-012: } |
| // CHECK-TILE-012-NOT: linalg.tiled_loop |
| |
| // CHECK-TILE-012-SKIP-PARTIAL: func @tiled_loop_3d_tensor( |
| // CHECK-TILE-012-SKIP-PARTIAL-SAME: %[[input:.*]]: tensor<?x?x?xf32> |
| // CHECK-TILE-012-SKIP-PARTIAL-DAG: %[[c0:.*]] = arith.constant 0 : index |
| // CHECK-TILE-012-SKIP-PARTIAL-DAG: %[[c1:.*]] = arith.constant 1 : index |
| // CHECK-TILE-012-SKIP-PARTIAL-DAG: %[[c2:.*]] = arith.constant 2 : index |
| // CHECK-TILE-012-SKIP-PARTIAL-DAG: %[[dim0:.*]] = tensor.dim %[[input]], %[[c0]] |
| // CHECK-TILE-012-SKIP-PARTIAL-DAG: %[[dim1:.*]] = tensor.dim %[[input]], %[[c1]] |
| // CHECK-TILE-012-SKIP-PARTIAL-DAG: %[[dim2:.*]] = tensor.dim %[[input]], %[[c2]] |
| // CHECK-TILE-012-SKIP-PARTIAL: %[[p0:.*]] = affine.apply #{{.*}}()[%[[dim0]] |
| // CHECK-TILE-012-SKIP-PARTIAL: %[[p1:.*]] = affine.apply #{{.*}}()[%[[dim1]] |
| // CHECK-TILE-012-SKIP-PARTIAL: %[[p2:.*]] = affine.apply #{{.*}}()[%[[dim2]] |
| // CHECK-TILE-012-SKIP-PARTIAL: linalg.tiled_loop {{.*}} = (%[[c0]], %[[c0]], %[[c0]]) to (%[[p0]], %[[p1]], %[[p2]]) |
| // CHECK-TILE-012-SKIP-PARTIAL: linalg.tiled_loop {{.*}} = (%[[c0]], %[[c0]], %[[p2]]) to (%[[p0]], %[[p1]], %[[dim2]]) |
| // CHECK-TILE-012-SKIP-PARTIAL: linalg.tiled_loop {{.*}} = (%[[c0]], %[[p1]], %[[c0]]) to (%[[p0]], %[[dim1]], %[[dim2]]) |
| // CHECK-TILE-012-SKIP-PARTIAL: linalg.tiled_loop {{.*}} = (%[[p0]], %[[c0]], %[[c0]]) to (%[[dim0]], %[[dim1]], %[[dim2]]) |
| func @tiled_loop_3d_tensor(%arg0: tensor<?x?x?xf32>, %s0: index, %s1: index, |
| %s2: index) -> tensor<?x?x?xf32> { |
| %cst = arith.constant 0.000000e+00 : f32 |
| %c0 = arith.constant 0 : index |
| %c1 = arith.constant 1 : index |
| %c2 = arith.constant 2 : index |
| %c8 = arith.constant 8 : index |
| %dim0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> |
| %dim1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> |
| %dim2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> |
| %output = linalg.init_tensor [%dim0, %dim1, %dim2] : tensor<?x?x?xf32> |
| %result = linalg.tiled_loop |
| (%arg1, %arg2, %arg3) = (%c0, %c0, %c0) to (%dim0, %dim1, %dim2) |
| step (%s0, %s1, %s2) ins (%arg4 = %arg0: tensor<?x?x?xf32>) |
| outs (%arg5 = %output: tensor<?x?x?xf32>) { |
| %min0 = affine.min affine_map<(d0, d1)[s0] -> (d1, -d0 + s0)>(%arg1, %s0)[%dim0] |
| %min1 = affine.min affine_map<(d0, d1)[s0] -> (d1, -d0 + s0)>(%arg2, %s1)[%dim1] |
| %min2 = affine.min affine_map<(d0, d1)[s0] -> (d1, -d0 + s0)>(%arg3, %s2)[%dim2] |
| %in_slice = tensor.extract_slice %arg4[%arg1, %arg2, %arg3] [%min0, %min1, %min2] [1, 1, 1]: tensor<?x?x?xf32> to tensor<?x?x?xf32> |
| %out_slice = tensor.extract_slice %arg5[%arg1, %arg2, %arg3] [%min0, %min1, %min2] [1, 1, 1] : tensor<?x?x?xf32> to tensor<?x?x?xf32> |
| %comp = "computation"(%in_slice, %out_slice) : (tensor<?x?x?xf32>, tensor<?x?x?xf32>) -> tensor<?x?x?xf32> |
| %updated_slice = tensor.insert_slice %comp into %arg5[%arg1, %arg2, %arg3] [%min0, %min1, %min2] [1, 1, 1] : tensor<?x?x?xf32> into tensor<?x?x?xf32> |
| linalg.yield %updated_slice : tensor<?x?x?xf32> |
| } |
| return %result : tensor<?x?x?xf32> |
| } |
| |
| // ----- |
| |
| // CHECK-TILE-2-LABEL: func @tiled_loop_3d_memref( |
| // CHECK-TILE-2-SAME: %[[input:.*]]: memref<?x?x?xf32>, %[[output:.*]]: memref<?x?x?xf32>, %[[s0:.*]]: index, %[[s1:.*]]: index, %[[s2:.*]]: index |
| // CHECK-TILE-2-DAG: %[[c0:.*]] = arith.constant 0 : index |
| // CHECK-TILE-2-DAG: %[[c1:.*]] = arith.constant 1 : index |
| // CHECK-TILE-2-DAG: %[[c2:.*]] = arith.constant 2 : index |
| // CHECK-TILE-2: %[[dim0:.*]] = memref.dim %[[input]], %[[c0]] |
| // CHECK-TILE-2: %[[dim1:.*]] = memref.dim %[[input]], %[[c1]] |
| // CHECK-TILE-2: %[[dim2:.*]] = memref.dim %[[input]], %[[c2]] |
| // CHECK-TILE-2: %[[split_bound:.*]] = affine.apply |
| // CHECK-TILE-2: linalg.tiled_loop (%[[iv0:.*]], %[[iv1:.*]], %[[iv2:.*]]) = (%[[c0]], %[[c0]], %[[c0]]) |
| // CHECK-TILE-2-SAME: to (%[[dim0]], %[[dim1]], %[[split_bound]]) |
| // CHECK-TILE-2-SAME: step (%[[s0]], %[[s1]], %[[s2]]) |
| // CHECK-TILE-2-SAME: ins (%[[loop_in1:.*]] = %[[input]]: memref<?x?x?xf32>) |
| // CHECK-TILE-2-SAME: outs (%[[loop_out1:.*]] = %[[output]]: memref<?x?x?xf32>) { |
| // CHECK-TILE-2: %[[min0_1:.*]] = affine.min |
| // CHECK-TILE-2: %[[min1_1:.*]] = affine.min |
| // CHECK-TILE-2: memref.subview %[[loop_in1]][%[[iv0]], %[[iv1]], %[[iv2]]] [%[[min0_1]], %[[min1_1]], %[[s2]]] |
| // CHECK-TILE-2: linalg.yield |
| // CHECK-TILE-2: linalg.tiled_loop (%[[iv0:.*]], %[[iv1:.*]], %[[iv2:.*]]) = (%[[c0]], %[[c0]], %[[split_bound]]) |
| // CHECK-TILE-2-SAME: to (%[[dim0]], %[[dim1]], %[[dim2]]) |
| // CHECK-TILE-2-SAME: step (%[[s0]], %[[s1]], %[[s2]]) |
| // CHECK-TILE-2-SAME: ins (%[[loop_in2:.*]] = %[[input]]: memref<?x?x?xf32>) |
| // CHECK-TILE-2-SAME: outs (%[[loop_out2:.*]] = %[[output]]: memref<?x?x?xf32>) { |
| // CHECK-TILE-2: %[[min0_2:.*]] = affine.min |
| // CHECK-TILE-2: %[[min1_2:.*]] = affine.min |
| // CHECK-TILE-2: %[[apply2:.*]] = affine.apply |
| // CHECK-TILE-2: memref.subview %[[loop_in1]][%[[iv0]], %[[iv1]], %[[iv2]]] [%[[min0_2]], %[[min1_2]], %[[apply2]]] |
| // CHECK-TILE-2: linalg.yield |
| // CHECK-TILE-2: return |
| |
| // CHECK-TILE-012-LABEL: func @tiled_loop_3d_memref |
| |
| !memref_subview_type = type memref<?x?x?xf32, affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>> |
| |
| func @tiled_loop_3d_memref(%arg0: memref<?x?x?xf32>, %output: memref<?x?x?xf32>, |
| %s0: index, %s1: index, %s2: index) { |
| %cst = arith.constant 0.000000e+00 : f32 |
| %c0 = arith.constant 0 : index |
| %c1 = arith.constant 1 : index |
| %c2 = arith.constant 2 : index |
| %c8 = arith.constant 8 : index |
| %dim0 = memref.dim %arg0, %c0 : memref<?x?x?xf32> |
| %dim1 = memref.dim %arg0, %c1 : memref<?x?x?xf32> |
| %dim2 = memref.dim %arg0, %c2 : memref<?x?x?xf32> |
| linalg.tiled_loop |
| (%arg1, %arg2, %arg3) = (%c0, %c0, %c0) to (%dim0, %dim1, %dim2) |
| step (%s0, %s1, %s2) ins (%arg4 = %arg0: memref<?x?x?xf32>) |
| outs (%arg5 = %output : memref<?x?x?xf32>) { |
| %min0 = affine.min affine_map<(d0, d1)[s0] -> (d1, -d0 + s0)>(%arg1, %s0)[%dim0] |
| %min1 = affine.min affine_map<(d0, d1)[s0] -> (d1, -d0 + s0)>(%arg2, %s1)[%dim1] |
| %min2 = affine.min affine_map<(d0, d1)[s0] -> (d1, -d0 + s0)>(%arg3, %s2)[%dim2] |
| %in_slice = memref.subview %arg4[%arg1, %arg2, %arg3] [%min0, %min1, %min2] [1, 1, 1]: memref<?x?x?xf32> to !memref_subview_type |
| "computation"(%in_slice) : (!memref_subview_type) -> memref<?x?x?xf32> |
| linalg.yield |
| } |
| return |
| } |
| |
| // ----- |
| |
| // CHECK-TILE-2-LABEL: func @step_1_do_not_peel |
| // CHECK-TILE-2: linalg.tiled_loop |
| // CHECK-TILE-2-NOT: linalg.tiled_loop |
| |
| // CHECK-TILE-012-LABEL: func @step_1_do_not_peel |
| |
| func @step_1_do_not_peel(%arg0: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> { |
| %cst = arith.constant 0.000000e+00 : f32 |
| %c0 = arith.constant 0 : index |
| %c1 = arith.constant 1 : index |
| %c2 = arith.constant 2 : index |
| %c8 = arith.constant 8 : index |
| %dim0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> |
| %dim1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> |
| %dim2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> |
| %output = linalg.init_tensor [%dim0, %dim1, %dim2] : tensor<?x?x?xf32> |
| %result = linalg.tiled_loop |
| (%arg1, %arg2, %arg3) = (%c0, %c0, %c0) to (%dim0, %dim1, %dim2) |
| step (%c1, %c1, %c1) ins (%arg4 = %arg0: tensor<?x?x?xf32>) |
| outs (%arg5 = %output: tensor<?x?x?xf32>) { |
| %in_slice = tensor.extract_slice %arg4[%arg1, %arg2, %arg3] [%c1, %c1, %c1] [1, 1, 1]: tensor<?x?x?xf32> to tensor<?x?x?xf32> |
| %out_slice = tensor.extract_slice %arg5[%arg1, %arg2, %arg3] [%c1, %c1, %c1] [1, 1, 1] : tensor<?x?x?xf32> to tensor<?x?x?xf32> |
| %comp = "computation"(%in_slice, %out_slice) : (tensor<?x?x?xf32>, tensor<?x?x?xf32>) -> tensor<?x?x?xf32> |
| %updated_slice = tensor.insert_slice %comp into %arg5[%arg1, %arg2, %arg3] [%c1, %c1, %c1] [1, 1, 1] : tensor<?x?x?xf32> into tensor<?x?x?xf32> |
| linalg.yield %updated_slice : tensor<?x?x?xf32> |
| } |
| return %result : tensor<?x?x?xf32> |
| } |
| |
| // ----- |
| |
| // CHECK-TILE-2-LABEL: func @divides_evenly_do_not_peel |
| // CHECK-TILE-2: linalg.tiled_loop |
| // CHECK-TILE-2-NOT: linalg.tiled_loop |
| |
| // CHECK-TILE-012-LABEL: func @divides_evenly_do_not_peel |
| |
| func @divides_evenly_do_not_peel(%arg0: tensor<?x?x?xf32>, %s: index) |
| -> tensor<?x?x?xf32> { |
| %cst = arith.constant 0.000000e+00 : f32 |
| %c0 = arith.constant 0 : index |
| %c1 = arith.constant 1 : index |
| %c2 = arith.constant 2 : index |
| %c8 = arith.constant 8 : index |
| %c64 = arith.constant 64 : index |
| %dim0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> |
| %dim1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> |
| %dim2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> |
| %output = linalg.init_tensor [%dim0, %dim1, %dim2] : tensor<?x?x?xf32> |
| %result = linalg.tiled_loop |
| (%arg1, %arg2, %arg3) = (%c0, %c0, %c0) to (%dim0, %dim1, %c64) |
| step (%s, %s, %c8) ins (%arg4 = %arg0: tensor<?x?x?xf32>) |
| outs (%arg5 = %output: tensor<?x?x?xf32>) { |
| %in_slice = tensor.extract_slice %arg4[%arg1, %arg2, %arg3] [%c1, %c1, %c1] [1, 1, 1]: tensor<?x?x?xf32> to tensor<?x?x?xf32> |
| %out_slice = tensor.extract_slice %arg5[%arg1, %arg2, %arg3] [%c1, %c1, %c1] [1, 1, 1] : tensor<?x?x?xf32> to tensor<?x?x?xf32> |
| %comp = "computation"(%in_slice, %out_slice) : (tensor<?x?x?xf32>, tensor<?x?x?xf32>) -> tensor<?x?x?xf32> |
| %updated_slice = tensor.insert_slice %comp into %arg5[%arg1, %arg2, %arg3] [%c1, %c1, %c1] [1, 1, 1] : tensor<?x?x?xf32> into tensor<?x?x?xf32> |
| linalg.yield %updated_slice : tensor<?x?x?xf32> |
| } |
| return %result : tensor<?x?x?xf32> |
| } |
