100 lines · plain
1// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries test-analysis-only" -split-input-file | FileCheck %s2 3func.func @not_elementwise(%a: tensor<5x6xf32>) -> tensor<5x6xf32> {4 %cst = arith.constant 5.0 : f325 // CHECK: tensor.extract_slice6 // CHECK-SAME: {__inplace_operands_attr__ = ["false"]}7 %b = tensor.extract_slice %a[0, 0] [1, 6] [1, 1]8 : tensor<5x6xf32> to tensor<6xf32>9 // CHECK: linalg.generic10 // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true"]}11 %0 = linalg.generic 12 { iterator_types = ["parallel", "parallel"],13 indexing_maps = [ affine_map<(d0, d1) -> (d1)>,14 affine_map<(d0, d1) -> (d0, d1)>] }15 ins(%b: tensor<6xf32>) outs(%a: tensor<5x6xf32>) {16 ^bb0(%arg0: f32, %arg1: f32):17 %r = arith.addf %arg0, %arg1 : f3218 linalg.yield %r : f3219 } -> tensor<5x6xf32>20 return %0 : tensor<5x6xf32>21}22 23// -----24 25#map = affine_map<(d0, d1) -> (d0, d1)>26#map1 = affine_map<(d0, d1) -> (d1)>27 28// CHECK-LABEL: @elementwise_no_conflict_429func.func @elementwise_no_conflict_4(%arg0: tensor<8x32x32x32xf32>, %arg1: tensor<32x32x32xf32>) -> tensor<8x32x32x32xf32> {30 %cst = arith.constant dense<3.000000e-02> : tensor<32x32x32xf32>31 %cst_0 = arith.constant dense<6.000000e-01> : tensor<32xf32>32 %cst_1 = arith.constant 0.000000e+00 : f3233 %r = scf.forall (%arg2, %arg3) in (8, 32) shared_outs(%arg4 = %arg0) -> (tensor<8x32x32x32xf32>) {34 // CHECK: tensor.extract_slice35 // CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]}36 %extracted_slice = tensor.extract_slice %arg4[%arg2, %arg3, 0, 0] [1, 1, 32, 32] [1, 1, 1, 1] : tensor<8x32x32x32xf32> to tensor<32x32xf32>37 38 // CHECK: linalg.fill39 // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}40 %4 = linalg.fill ins(%cst_1 : f32) outs(%extracted_slice : tensor<32x32xf32>) -> tensor<32x32xf32>41 42 // CHECK: linalg.batch_reduce_matmul43 // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}44 %5 = linalg.batch_reduce_matmul ins(%arg1, %cst : tensor<32x32x32xf32>, tensor<32x32x32xf32>) outs(%4 : tensor<32x32xf32>) -> tensor<32x32xf32>45 46 // CHECK: linalg.generic47 // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}48 // %cst_0 has a non-identity layout may, but %5 and %extracted_slice still49 // bufferize to element-wise access.50 %6 = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel"]} ins(%5, %cst_0 : tensor<32x32xf32>, tensor<32xf32>) outs(%extracted_slice : tensor<32x32xf32>) {51 ^bb0(%in: f32, %in_4: f32, %out: f32):52 %8 = arith.addf %in, %in_4 : f3253 linalg.yield %8 : f3254 } -> tensor<32x32xf32>55 56 // CHECK: linalg.generic57 // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true"]}58 // They are different SSA values, but %6 and %extract_slice are equivalent.59 %7 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%6 : tensor<32x32xf32>) outs(%extracted_slice : tensor<32x32xf32>) {60 ^bb0(%in: f32, %out: f32):61 %8 = arith.maximumf %in, %cst_1 : f3262 linalg.yield %8 : f3263 } -> tensor<32x32xf32>64 scf.forall.in_parallel {65 // CHECK: tensor.parallel_insert_slice66 // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none"]}67 tensor.parallel_insert_slice %7 into %arg4[%arg2, %arg3, 0, 0] [1, 1, 32, 32] [1, 1, 1, 1] : tensor<32x32xf32> into tensor<8x32x32x32xf32>68 }69 }70 return %r : tensor<8x32x32x32xf32>71}72 73// -----74 75// CHECK-LABEL: func @elementwise_access_regression(76// CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "false"]}77// CHECK: linalg.map78// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}79// CHECK: linalg.map80// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}81func.func private @f(%arg: tensor<32x1xf32>) -> ()82func.func @elementwise_access_regression(%arg0: i32, %arg2: tensor<32x1xf32>, %arg3: tensor<32x1xf32>) {83 %cst_0 = arith.constant 0.000000e+00 : f3284 %c0_i32 = arith.constant 0 : i3285 %c1_i32 = arith.constant 1 : i3286 %0 = tensor.empty() : tensor<32x1xf32>87 88 // This op must bufferize out-of-place so that the filled tensor is not89 // overwritten by the ops inside of the loop.90 %1 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<32x1xf32>) -> tensor<32x1xf32>91 92 scf.for %arg1 = %c0_i32 to %arg0 step %c1_i32 : i32 {93 %2 = linalg.map { arith.subf } ins(%1, %arg2 : tensor<32x1xf32>, tensor<32x1xf32>) outs(%0 : tensor<32x1xf32>)94 %3 = tensor.empty() : tensor<32x1xf32>95 %4 = linalg.map { arith.subf } ins(%2, %arg3 : tensor<32x1xf32>, tensor<32x1xf32>) outs(%3 : tensor<32x1xf32>)96 func.call @f(%4) : (tensor<32x1xf32>) -> ()97 }98 return99}100