mlir/test/Dialect/Tensor/fold-constant-extract-slice.mlir

   1 // RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-fold-constant-extract-slice %s | FileCheck %s
   2
   3 // CHECK-LABEL: func @slice_constant
   4 //   CHECK-NOT:   tensor.extract_slice
   5 //       CHECK:   %[[CONST:.+]] = arith.constant dense<1.000000e+01> : tensor<1x1xf32>
   6 //       CHECK:   return %[[CONST]] :  tensor<1x1xf32>
   7 func.func @slice_constant(%arg0 : tensor<2x1xf32>) -> tensor<1x1xf32>
   8 {
   9   %cst = arith.constant dense<[[10.0], [11.0]]> : tensor<2x1xf32>
  10   %slice = tensor.extract_slice %cst[0, 0] [1, 1] [1, 1] : tensor<2x1xf32> to tensor<1x1xf32>
  11   return %slice : tensor<1x1xf32>
  12 }
  13
  14 // -----
  15
  16 // CHECK-LABEL: func @slice_constant_3x4
  17 //   CHECK-NOT:   tensor.extract_slice
  18 //       CHECK:   %[[CONST:.+]] = arith.constant dense<{{\[}}[1.000000e+01, 9.000000e+00], [1.100000e+01, 1.200000e+01]]> : tensor<2x2xf32>
  19 //       CHECK:   return %[[CONST]] :  tensor<2x2xf32>
  20 func.func @slice_constant_3x4(%arg0 : tensor<3x4xf32>) -> tensor<2x2xf32>
  21 {
  22   %cst = arith.constant dense<[[10.0, 9.0, 8.0, 7.0], [11.0, 12.0, 13.0, 14.0], [1.0, 3.0, 5.0, 7.0]]> : tensor<3x4xf32>
  23   %slice = tensor.extract_slice %cst[0, 0] [2, 2] [1, 1] : tensor<3x4xf32> to tensor<2x2xf32>
  24   return %slice : tensor<2x2xf32>
  25 }
  26
  27 // -----
  28
  29 // CHECK-LABEL: func @slice_constant_3x4_offsets
  30 //   CHECK-NOT:   tensor.extract_slice
  31 //       CHECK:   %[[CONST:.+]] = arith.constant dense<{{\[}}[1.200000e+01, 1.300000e+01], [3.000000e+00, 5.000000e+00]]> : tensor<2x2xf32>
  32 //       CHECK:   return %[[CONST]] :  tensor<2x2xf32>
  33 func.func @slice_constant_3x4_offsets(%arg0 : tensor<3x4xf32>) -> tensor<2x2xf32>
  34 {
  35   %cst = arith.constant dense<[[10.0, 9.0, 8.0, 7.0], [11.0, 12.0, 13.0, 14.0], [1.0, 3.0, 5.0, 7.0]]> : tensor<3x4xf32>
  36   %slice = tensor.extract_slice %cst[1, 1] [2, 2] [1, 1] : tensor<3x4xf32> to tensor<2x2xf32>
  37   return %slice : tensor<2x2xf32>
  38 }
  39