132 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -split-input-file | FileCheck %s2 3///----------------------------------------------------------------------------------------4/// Tests for tensor.pad5///----------------------------------------------------------------------------------------6 7// CHECK-LABEL: func @test_masked_vectorize_pad8func.func @test_masked_vectorize_pad(9 %0 : tensor<?x?xf32>, %h0 : index, %h1 : index)10 -> tensor<2x4xf32>11{12 // CHECK-DAG: %[[c42:.*]] = arith.constant 4.243000e+01 : f3213 // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index14 // CHECK-DAG: %[[c0_0:.*]] = arith.constant 0 : index15 // CHECK: %[[d0:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>16 // CHECK: %[[d1:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>17 // CHECK: %[[mask:.*]] = vector.create_mask %[[d0]], %[[d1]] : vector<2x4xi1>18 // CHECK: %[[masked_read:.*]] = vector.mask %[[mask]] {19 // CHECK-SAME: vector.transfer_read %{{.*}}[%[[c0_0]], %[[c0_0]]], %[[c42]]20 // CHECK-SAME: {in_bounds = [true, true]} : tensor<?x?xf32>, vector<2x4xf32>21 // CHECK-SAME: } : vector<2x4xi1> -> vector<2x4xf32>22 // CHECK-DAG: %[[c0_1:.*]] = arith.constant 0 : index23 // CHECK-DAG: %[[empty:.*]] = tensor.empty() : tensor<2x4xf32>24 // CHECK: vector.transfer_write %[[masked_read]], %[[empty]][%[[c0_1]], %[[c0_1]]]25 // CHECK-SAME: {in_bounds = [true, true]} : vector<2x4xf32>, tensor<2x4xf32>26 %cst = arith.constant 42.43 : f3227 %c0 = arith.constant 0 : index28 %1 = tensor.pad %0 low[0, %c0] high[%h0, %h1] {29 ^bb0(%hh1: index, %hh2: index):30 tensor.yield %cst : f3231 } : tensor<?x?xf32> to tensor<2x4xf32>32 return %1: tensor<2x4xf32>33}34 35module attributes {transform.with_named_sequence} {36 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {37 %0 = transform.structured.match ops{["tensor.pad"]} in %arg138 : (!transform.any_op) -> !transform.any_op39 transform.structured.vectorize %0 vector_sizes [2, 4] : !transform.any_op40 transform.yield41 }42}43 44// -----45 46// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>47// CHECK: func @test_masked_vectorize_dynamic_pad48func.func @test_masked_vectorize_dynamic_pad(49 %0 : tensor<?x?xf32>, %h0 : index, %h1 : index)50 -> tensor<?x?xf32>51{52 // CHECK-DAG: %[[c42:.*]] = arith.constant 4.243000e+01 : f3253 // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index54 // CHECK-DAG: %[[res_d0:.+]] = affine.apply #[[MAP]]()55 // CHECK-DAG: %[[res_d1:.+]] = affine.apply #[[MAP]]()56 // CHECK: %[[c0_2:.*]] = arith.constant 0 : index57 // CHECK: %[[d0:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>58 // CHECK: %[[d1:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>59 // CHECK: %[[mask:.*]] = vector.create_mask %[[d0]], %[[d1]] : vector<2x4xi1>60 // CHECK: %[[masked_read:.*]] = vector.mask %[[mask]] {61 // CHECK-SAME: vector.transfer_read %{{.*}}[%[[c0_2]], %[[c0_2]]], %[[c42]]62 // CHECK-SAME: {in_bounds = [true, true]} : tensor<?x?xf32>, vector<2x4xf32>63 // CHECK-SAME: } : vector<2x4xi1> -> vector<2x4xf32>64 // CHECK-DAG: %[[empty:.*]] = tensor.empty(%[[res_d0]], %[[res_d1]]) : tensor<?x?xf32>65 // CHECK-DAG: %[[c0_3:.*]] = arith.constant 0 : index66 // CHECK-DAG: %[[d2:.*]] = tensor.dim %[[empty]], {{.*}} : tensor<?x?xf32>67 // CHECK-DAG: %[[d3:.*]] = tensor.dim %[[empty]], {{.*}} : tensor<?x?xf32>68 // CHECK: %[[mask_2:.*]] = vector.create_mask %[[d2]], %[[d3]] : vector<2x4xi1>69 // CHECK: %[[masked_write:.*]] = vector.mask %[[mask_2]] {70 // CHECK-SAME: vector.transfer_write %[[masked_read]], %[[empty]][%[[c0_3]], %[[c0_3]]]71 // CHECK-SAME: {in_bounds = [true, true]} : vector<2x4xf32>, tensor<?x?xf32>72 // CHECK: return %[[masked_write]] : tensor<?x?xf32>73 %cst = arith.constant 42.43 : f3274 %c0 = arith.constant 0 : index75 %1 = tensor.pad %0 low[0, %c0] high[%h0, %h1] {76 ^bb0(%hh1: index, %hh2: index):77 tensor.yield %cst : f3278 } : tensor<?x?xf32> to tensor<?x?xf32>79 return %1: tensor<?x?xf32>80}81 82module attributes {transform.with_named_sequence} {83 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {84 %0 = transform.structured.match ops{["tensor.pad"]} in %arg185 : (!transform.any_op) -> !transform.any_op86 transform.structured.vectorize %0 vector_sizes [2, 4] : !transform.any_op87 transform.yield88 }89}90 91// -----92// This case is supported because low padding `%l0` is applied on93// a unit dimension which is supported, non unit result dimension low94// padding is currently unsupported.95// CHECK-LABEL: func @test_masked_vectorize_non_zero_low_pad_unit_res_dim96func.func @test_masked_vectorize_non_zero_low_pad_unit_res_dim(97 %0 : tensor<?x?xf32>, %h0 : index, %h1 : index, %l0 : index)98 -> tensor<1x4xf32>99{100 // CHECK-DAG: %[[C42:.*]] = arith.constant 4.243000e+01 : f32101 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index102 // CHECK: %[[C0_1:.*]] = arith.constant 0 : index103 // CHECK-DAG: %[[D0:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>104 // CHECK-DAG: %[[D1:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>105 // CHECK: %[[MASK:.*]] = vector.create_mask %[[D0]], %[[D1]] : vector<1x4xi1>106 // CHECK: %[[MASKED_READ:.*]] = vector.mask %[[MASK]] {107 // CHECK-SAME: vector.transfer_read %{{.*}}[%[[C0_1]], %[[C0_1]]], %[[C42]]108 // CHECK-SAME: {in_bounds = [true, true]} : tensor<?x?xf32>, vector<1x4xf32>109 // CHECK-SAME: } : vector<1x4xi1> -> vector<1x4xf32>110 // CHECK-DAG: %[[EMPTY:.*]] = tensor.empty() : tensor<1x4xf32>111 // CHECK-DAG: %[[C0_2:.*]] = arith.constant 0 : index112 // CHECK: %[[MASKED_WRITE:.*]] = vector.transfer_write %[[MASKED_READ]], %[[EMPTY]][%[[C0_2]], %[[C0_2]]]113 // CHECK-SAME: {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>114 // CHECK: return %[[MASKED_WRITE]] : tensor<1x4xf32>115 %cst = arith.constant 42.43 : f32116 %c0 = arith.constant 0 : index117 %1 = tensor.pad %0 low[%l0, %c0] high[%h0, %h1] {118 ^bb0(%hh1: index, %hh2: index):119 tensor.yield %cst : f32120 } : tensor<?x?xf32> to tensor<1x4xf32>121 return %1: tensor<1x4xf32>122}123 124module attributes {transform.with_named_sequence} {125 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {126 %0 = transform.structured.match ops{["tensor.pad"]} in %arg1127 : (!transform.any_op) -> !transform.any_op128 transform.structured.vectorize %0 vector_sizes [1, 4] : !transform.any_op129 transform.yield130 }131}132