| // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-op=linalg.conv_2d fuse tile-sizes=4,4,0,0 tile-interchange=0,1,2,3 run-enable-pass=false" -split-input-file | FileCheck --check-prefix=CONV %s |
| // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-op=linalg.matmul fuse tile-sizes=4,4,0 tile-interchange=0,1,2 run-enable-pass=false" -split-input-file | FileCheck --check-prefix=MATMUL %s |
| |
| // CONV: fuse_conv_chain |
| // CONV-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<2x2xf32> |
| // CONV-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<11x11xf32> |
| // CONV-SAME: %[[ARG2:[0-9a-zA-Z]*]]: tensor<10x10xf32> |
| // CONV-SAME: %[[ARG3:[0-9a-zA-Z]*]]: tensor<9x9xf32> |
| // CONV-SAME: %[[ARG4:[0-9a-zA-Z]*]]: tensor<8x8xf32> |
| builtin.func @fuse_conv_chain(%arg0: tensor<2x2xf32>, |
| %arg1: tensor<11x11xf32>, |
| %arg2: tensor<10x10xf32>, |
| %arg3: tensor<9x9xf32>, |
| %arg4: tensor<8x8xf32>) -> tensor<8x8xf32> { |
| %cst = arith.constant 1.0 : f32 |
| |
| // Do not tile the filter fill since the filter dimensions are not tiled. |
| // CONV: %[[T0:.*]] = linalg.fill(%{{.*}}, %[[ARG0]]) |
| %0 = linalg.fill(%cst, %arg0) : f32, tensor<2x2xf32> -> tensor<2x2xf32> |
| |
| // Fuse all other operations. |
| // CONV: scf.for %[[IV0:.*]] = {{.*}} iter_args(%[[ARG5:.*]] = %[[ARG4]] |
| // CONV: scf.for %[[IV1:.*]] = {{.*}} iter_args(%[[ARG6:.*]] = %[[ARG5]] |
| |
| // CONV: %[[T1:.*]] = tensor.extract_slice %[[ARG1]] |
| // CONV-SAME: %[[IV0]], %[[IV1]] |
| // CONV: %[[T2:.*]] = tensor.extract_slice %[[ARG2]] |
| // CONV-SAME: %[[IV0]], %[[IV1]] |
| // CONV: %[[T3:.*]] = linalg.fill(%{{.*}}, %[[T2]]) |
| // CONV: %[[T4:.*]] = linalg.conv_2d ins(%[[T1]], %[[T0]] : {{.*}} outs(%[[T3]] |
| %1 = linalg.fill(%cst, %arg2) : f32, tensor<10x10xf32> -> tensor<10x10xf32> |
| %2 = linalg.conv_2d ins(%arg1, %0 : tensor<11x11xf32>, tensor<2x2xf32>) outs(%1 : tensor<10x10xf32>) -> tensor<10x10xf32> |
| |
| // CONV: %[[T5:.*]] = tensor.extract_slice %[[ARG3]] |
| // CONV-SAME: %[[IV0]], %[[IV1]] |
| // CONV: %[[T6:.*]] = linalg.fill(%{{.*}}, %[[T5]]) |
| // CONV: %[[T7:.*]] = linalg.conv_2d ins(%[[T4]], %[[T0]] : {{.*}} outs(%[[T6]] |
| %3 = linalg.fill(%cst, %arg3) : f32, tensor<9x9xf32> -> tensor<9x9xf32> |
| %4 = linalg.conv_2d ins(%2, %0 : tensor<10x10xf32>, tensor<2x2xf32>) outs(%3 : tensor<9x9xf32>) -> tensor<9x9xf32> |
| |
| // Use the argument passed in by iteration argument. |
| // CONV: %[[T8:.*]] = tensor.extract_slice %[[ARG6]] |
| // CONV-SAME: %[[IV0]], %[[IV1]] |
| // CONV: %[[T9:.*]] = linalg.fill(%{{.*}}, %[[T8]]) |
| // CONV: %[[T5:.*]] = linalg.conv_2d ins(%[[T7]], %[[T0]] {{.*}} outs(%[[T9]] |
| %5 = linalg.fill(%cst, %arg4) : f32, tensor<8x8xf32> -> tensor<8x8xf32> |
| %6 = linalg.conv_2d ins(%4, %0 : tensor<9x9xf32>, tensor<2x2xf32>) outs(%5 : tensor<8x8xf32>) -> tensor<8x8xf32> |
| return %6 : tensor<8x8xf32> |
| } |
| |
| // ----- |
| |
| // MATMUL: fuse_matmul_chain |
| // MATMUL-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<8x8xf32> |
| builtin.func @fuse_matmul_chain(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> { |
| %c0 = arith.constant 0 : index |
| %c12 = arith.constant 12 : index |
| %c25 = arith.constant 25 : index |
| %c24 = arith.constant 24 : index |
| %c4 = arith.constant 4 : index |
| %cst = arith.constant 0.000000e+00 : f32 |
| |
| // Do not tile rhs fill of the producer matmul since none of its loop dimension is tiled. |
| // MATMUL: %[[T0:.*]] = linalg.fill(%{{.*}}, %[[ARG0]]) |
| %0 = linalg.fill(%cst, %arg0) : f32, tensor<8x8xf32> -> tensor<8x8xf32> |
| |
| // MATMUL: scf.for %[[IV0:.*]] = {{.*}} iter_args(%[[ARG1:.*]] = %[[ARG0]] |
| // MATMUL: scf.for %[[IV1:.*]] = {{.*}} iter_args(%[[ARG2:.*]] = %[[ARG1]] |
| |
| // Only the outermost loop of the producer matmul is tiled. |
| // MATMUL: %[[T1:.*]] = tensor.extract_slice %[[ARG0]] |
| // MATMUL-SAME: %[[IV0]], 0 |
| // MATMUL: %[[T2:.*]] = linalg.fill(%{{.*}}, %[[T1]]) |
| // MATMUL: %[[T3:.*]] = linalg.matmul ins(%[[T2]], %[[T0]] {{.*}} |
| %1 = linalg.matmul ins(%0, %0 : tensor<8x8xf32>, tensor<8x8xf32>) outs(%0 : tensor<8x8xf32>) -> tensor<8x8xf32> |
| |
| // Use the argument passed in by iteration argument. |
| // MATMUL: %[[T4:.*]] = tensor.extract_slice %[[ARG2]] |
| // MATMUL-SAME: %[[IV0]], %[[IV1]] |
| // MATMUL: %[[T5:.*]] = linalg.fill(%{{.*}}, %[[T4]]) |
| // MATMUL: %{{.*}} = linalg.matmul ins(%[[T3]], {{.*}} outs(%[[T5]] |
| %2 = linalg.matmul ins(%1, %0 : tensor<8x8xf32>, tensor<8x8xf32>) outs(%0 : tensor<8x8xf32>) -> tensor<8x8xf32> |
| return %2 : tensor<8x8xf32> |
| } |
