184 lines · plain
1// RUN: mlir-opt %s -split-input-file --test-vector-scan-lowering | FileCheck %s2 3// CHECK-LABEL: func @scan1d_inc4// CHECK-SAME: %[[ARG0:.*]]: vector<2xi32>,5// CHECK-SAME: %[[ARG1:.*]]: vector<i32>6// CHECK: %[[A:.*]] = arith.constant dense<0> : vector<2xi32>7// CHECK: %[[B:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xi32> to vector<1xi32>8// CHECK: %[[C:.*]] = vector.insert_strided_slice %[[B]], %[[A]] {offsets = [0], strides = [1]} : vector<1xi32> into vector<2xi32>9// CHECK: %[[D:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [1], sizes = [1], strides = [1]} : vector<2xi32> to vector<1xi32>10// CHECK: %[[E:.*]] = arith.addi %[[B]], %[[D]] : vector<1xi32>11// CHECK: %[[F:.*]] = vector.insert_strided_slice %[[E]], %[[C]] {offsets = [1], strides = [1]} : vector<1xi32> into vector<2xi32>12// CHECK: %[[G:.*]] = vector.extract %[[E]][0] : i32 from vector<1xi32>13// CHECK: %[[H:.*]] = vector.broadcast %[[G]] : i32 to vector<i32>14// CHECK: return %[[F]], %[[H]] : vector<2xi32>, vector<i32>15func.func @scan1d_inc(%arg0 : vector<2xi32>, %arg1 : vector<i32>) -> (vector<2xi32>, vector<i32>) {16 %0:2 = vector.scan <add>, %arg0, %arg1 {inclusive = true, reduction_dim = 0} :17 vector<2xi32>, vector<i32>18 return %0#0, %0#1 : vector<2xi32>, vector<i32>19}20 21// -----22 23// Reducing scalable dims is not yet supported!24 25// CHECK-LABEL: func @scan1d_inc_scalable26// CHECK: vector.scan27func.func @scan1d_inc_scalable(%arg0 : vector<[2]xi32>, %arg1 : vector<i32>) -> (vector<[2]xi32>, vector<i32>) {28 %0:2 = vector.scan <add>, %arg0, %arg1 {inclusive = true, reduction_dim = 0} :29 vector<[2]xi32>, vector<i32>30 return %0#0, %0#1 : vector<[2]xi32>, vector<i32>31}32 33// -----34 35// CHECK-LABEL: func @scan1d_exc36// CHECK-SAME: %[[ARG0:.*]]: vector<2xi32>,37// CHECK-SAME: %[[ARG1:.*]]: vector<i32>38// CHECK: %[[A:.*]] = arith.constant dense<0> : vector<2xi32>39// CHECK: %[[B:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xi32> to vector<1xi32>40// CHECK: %[[C:.*]] = vector.broadcast %[[ARG1]] : vector<i32> to vector<1xi32>41// CHECK: %[[D:.*]] = vector.insert_strided_slice %[[C]], %[[A]] {offsets = [0], strides = [1]} : vector<1xi32> into vector<2xi32>42// CHECK: %[[E:.*]] = arith.addi %[[C]], %[[B]] : vector<1xi32>43// CHECK: %[[F:.*]] = vector.insert_strided_slice %[[E]], %[[D]] {offsets = [1], strides = [1]} : vector<1xi32> into vector<2xi32>44// CHECK: %[[G:.*]] = vector.extract %[[E]][0] : i32 from vector<1xi32>45// CHECK: %[[H:.*]] = vector.broadcast %[[G]] : i32 to vector<i32>46// CHECK: return %[[F]], %[[H]] : vector<2xi32>, vector<i32>47func.func @scan1d_exc(%arg0 : vector<2xi32>, %arg1 : vector<i32>) -> (vector<2xi32>, vector<i32>) {48 %0:2 = vector.scan <add>, %arg0, %arg1 {inclusive = false, reduction_dim = 0} :49 vector<2xi32>, vector<i32>50 return %0#0, %0#1 : vector<2xi32>, vector<i32>51}52 53// -----54 55// Rducing scalable dims is not yet supported!56 57// CHECK-LABEL: func @scan1d_exc_scalable58// CHECK: vector.scan59func.func @scan1d_exc_scalable(%arg0 : vector<[2]xi32>, %arg1 : vector<i32>) -> (vector<[2]xi32>, vector<i32>) {60 %0:2 = vector.scan <add>, %arg0, %arg1 {inclusive = false, reduction_dim = 0} :61 vector<[2]xi32>, vector<i32>62 return %0#0, %0#1 : vector<[2]xi32>, vector<i32>63}64 65// -----66 67// CHECK-LABEL: func @scan2d_mul_dim068// CHECK-SAME: %[[ARG0:.*]]: vector<2x3xi32>,69// CHECK-SAME: %[[ARG1:.*]]: vector<3xi32>70// CHECK: %[[A:.*]] = arith.constant dense<0> : vector<2x3xi32>71// CHECK: %[[B:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0, 0], sizes = [1, 3], strides = [1, 1]} : vector<2x3xi32> to vector<1x3xi32>72// CHECK: %[[C:.*]] = vector.insert_strided_slice %[[B]], %[[A]] {offsets = [0, 0], strides = [1, 1]} : vector<1x3xi32> into vector<2x3xi32>73// CHECK: %[[D:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [1, 0], sizes = [1, 3], strides = [1, 1]} : vector<2x3xi32> to vector<1x3xi32>74// CHECK: %[[E:.*]] = arith.muli %[[B]], %[[D]] : vector<1x3xi32>75// CHECK: %[[F:.*]] = vector.insert_strided_slice %[[E]], %[[C]] {offsets = [1, 0], strides = [1, 1]} : vector<1x3xi32> into vector<2x3xi32>76// CHECK: %[[G:.*]] = vector.shape_cast %[[E]] : vector<1x3xi32> to vector<3xi32>77// CHECK: return %[[F]], %[[G]] : vector<2x3xi32>, vector<3xi32>78func.func @scan2d_mul_dim0(%arg0 : vector<2x3xi32>, %arg1 : vector<3xi32>) -> (vector<2x3xi32>, vector<3xi32>) {79 %0:2 = vector.scan <mul>, %arg0, %arg1 {inclusive = true, reduction_dim = 0} :80 vector<2x3xi32>, vector<3xi32>81 return %0#0, %0#1 : vector<2x3xi32>, vector<3xi32>82}83 84// -----85 86// CHECK-LABEL: func @scan2d_mul_dim0_scalable87// CHECK-SAME: %[[ARG0:.*]]: vector<2x[3]xi32>,88// CHECK-SAME: %[[ARG1:.*]]: vector<[3]xi32>89// CHECK: %[[A:.*]] = arith.constant dense<0> : vector<2x[3]xi32>90// CHECK: %[[B:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0, 0], sizes = [1, 3], strides = [1, 1]} : vector<2x[3]xi32> to vector<1x[3]xi32>91// CHECK: %[[C:.*]] = vector.insert_strided_slice %[[B]], %[[A]] {offsets = [0, 0], strides = [1, 1]} : vector<1x[3]xi32> into vector<2x[3]xi32>92// CHECK: %[[D:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [1, 0], sizes = [1, 3], strides = [1, 1]} : vector<2x[3]xi32> to vector<1x[3]xi32>93// CHECK: %[[E:.*]] = arith.muli %[[B]], %[[D]] : vector<1x[3]xi32>94// CHECK: %[[F:.*]] = vector.insert_strided_slice %[[E]], %[[C]] {offsets = [1, 0], strides = [1, 1]} : vector<1x[3]xi32> into vector<2x[3]xi32>95// CHECK: %[[G:.*]] = vector.shape_cast %[[E]] : vector<1x[3]xi32> to vector<[3]xi32>96// CHECK: return %[[F]], %[[G]] : vector<2x[3]xi32>, vector<[3]xi32>97func.func @scan2d_mul_dim0_scalable(%arg0 : vector<2x[3]xi32>, %arg1 : vector<[3]xi32>) -> (vector<2x[3]xi32>, vector<[3]xi32>) {98 %0:2 = vector.scan <mul>, %arg0, %arg1 {inclusive = true, reduction_dim = 0} :99 vector<2x[3]xi32>, vector<[3]xi32>100 return %0#0, %0#1 : vector<2x[3]xi32>, vector<[3]xi32>101}102 103// -----104 105// CHECK-LABEL: func @scan2d_mul_dim1106// CHECK-SAME: %[[ARG0:.*]]: vector<2x3xi32>,107// CHECK-SAME: %[[ARG1:.*]]: vector<2xi32>108// CHECK: %[[A:.*]] = arith.constant dense<0> : vector<2x3xi32>109// CHECK: %[[B:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0, 0], sizes = [2, 1], strides = [1, 1]} : vector<2x3xi32> to vector<2x1xi32>110// CHECK: %[[C:.*]] = vector.insert_strided_slice %[[B]], %[[A]] {offsets = [0, 0], strides = [1, 1]} : vector<2x1xi32> into vector<2x3xi32>111// CHECK: %[[D:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0, 1], sizes = [2, 1], strides = [1, 1]} : vector<2x3xi32> to vector<2x1xi32>112// CHECK: %[[E:.*]] = arith.muli %[[B]], %[[D]] : vector<2x1xi32>113// CHECK: %[[F:.*]] = vector.insert_strided_slice %[[E]], %[[C]] {offsets = [0, 1], strides = [1, 1]} : vector<2x1xi32> into vector<2x3xi32>114// CHECK: %[[G:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0, 2], sizes = [2, 1], strides = [1, 1]} : vector<2x3xi32> to vector<2x1xi32>115// CHECK: %[[H:.*]] = arith.muli %[[E]], %[[G]] : vector<2x1xi32>116// CHECK: %[[I:.*]] = vector.insert_strided_slice %[[H]], %[[F]] {offsets = [0, 2], strides = [1, 1]} : vector<2x1xi32> into vector<2x3xi32>117// CHECK: %[[J:.*]] = vector.shape_cast %[[H]] : vector<2x1xi32> to vector<2xi32>118// CHECK: return %[[I]], %[[J]] : vector<2x3xi32>, vector<2xi32>119func.func @scan2d_mul_dim1(%arg0 : vector<2x3xi32>, %arg1 : vector<2xi32>) -> (vector<2x3xi32>, vector<2xi32>) {120 %0:2 = vector.scan <mul>, %arg0, %arg1 {inclusive = true, reduction_dim = 1} :121 vector<2x3xi32>, vector<2xi32>122 return %0#0, %0#1 : vector<2x3xi32>, vector<2xi32>123}124 125// -----126 127// CHECK-LABEL: func @scan2d_mul_dim1_scalable128// CHECK-SAME: %[[ARG0:.*]]: vector<[2]x3xi32>,129// CHECK-SAME: %[[ARG1:.*]]: vector<[2]xi32>130// CHECK: %[[A:.*]] = arith.constant dense<0> : vector<[2]x3xi32>131// CHECK: %[[B:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0, 0], sizes = [2, 1], strides = [1, 1]} : vector<[2]x3xi32> to vector<[2]x1xi32>132// CHECK: %[[C:.*]] = vector.insert_strided_slice %[[B]], %[[A]] {offsets = [0, 0], strides = [1, 1]} : vector<[2]x1xi32> into vector<[2]x3xi32>133// CHECK: %[[D:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0, 1], sizes = [2, 1], strides = [1, 1]} : vector<[2]x3xi32> to vector<[2]x1xi32>134// CHECK: %[[E:.*]] = arith.muli %[[B]], %[[D]] : vector<[2]x1xi32>135// CHECK: %[[F:.*]] = vector.insert_strided_slice %[[E]], %[[C]] {offsets = [0, 1], strides = [1, 1]} : vector<[2]x1xi32> into vector<[2]x3xi32>136// CHECK: %[[G:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0, 2], sizes = [2, 1], strides = [1, 1]} : vector<[2]x3xi32> to vector<[2]x1xi32>137// CHECK: %[[H:.*]] = arith.muli %[[E]], %[[G]] : vector<[2]x1xi32>138// CHECK: %[[I:.*]] = vector.insert_strided_slice %[[H]], %[[F]] {offsets = [0, 2], strides = [1, 1]} : vector<[2]x1xi32> into vector<[2]x3xi32>139// CHECK: %[[J:.*]] = vector.shape_cast %[[H]] : vector<[2]x1xi32> to vector<[2]xi32>140// CHECK: return %[[I]], %[[J]] : vector<[2]x3xi32>, vector<[2]xi32>141func.func @scan2d_mul_dim1_scalable(%arg0 : vector<[2]x3xi32>, %arg1 : vector<[2]xi32>) -> (vector<[2]x3xi32>, vector<[2]xi32>) {142 %0:2 = vector.scan <mul>, %arg0, %arg1 {inclusive = true, reduction_dim = 1} :143 vector<[2]x3xi32>, vector<[2]xi32>144 return %0#0, %0#1 : vector<[2]x3xi32>, vector<[2]xi32>145}146 147// -----148 149// CHECK-LABEL: func @scan3d_mul_dim1150// CHECK-SAME: %[[ARG0:.*]]: vector<4x2x3xf32>,151// CHECK-SAME: %[[ARG1:.*]]: vector<4x3xf32>152// CHECK: %[[A:.*]] = arith.constant dense<0.000000e+00> : vector<4x2x3xf32>153// CHECK: %[[B:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0, 0, 0], sizes = [4, 1, 3], strides = [1, 1, 1]} : vector<4x2x3xf32> to vector<4x1x3xf32>154// CHECK: %[[C:.*]] = vector.shape_cast %[[ARG1]] : vector<4x3xf32> to vector<4x1x3xf32>155// CHECK: %[[D:.*]] = vector.insert_strided_slice %[[C]], %[[A]] {offsets = [0, 0, 0], strides = [1, 1, 1]} : vector<4x1x3xf32> into vector<4x2x3xf32>156// CHECK: %[[E:.*]] = arith.mulf %[[C]], %[[B]] : vector<4x1x3xf32>157// CHECK: %[[F:.*]] = vector.insert_strided_slice %[[E]], %[[D]] {offsets = [0, 1, 0], strides = [1, 1, 1]} : vector<4x1x3xf32> into vector<4x2x3xf32>158// CHECK: %[[G:.*]] = vector.shape_cast %[[E]] : vector<4x1x3xf32> to vector<4x3xf32>159// CHECK: return %[[F]], %[[G]] : vector<4x2x3xf32>, vector<4x3xf32>160func.func @scan3d_mul_dim1(%arg0 : vector<4x2x3xf32>, %arg1 : vector<4x3xf32>) -> (vector<4x2x3xf32>, vector<4x3xf32>) {161 %0:2 = vector.scan <mul>, %arg0, %arg1 {inclusive = false, reduction_dim = 1} :162 vector<4x2x3xf32>, vector<4x3xf32>163 return %0#0, %0#1 : vector<4x2x3xf32>, vector<4x3xf32>164}165 166// -----167 168// CHECK-LABEL: func @scan3d_mul_dim1_scalable169// CHECK-SAME: %[[ARG0:.*]]: vector<4x2x[3]xf32>,170// CHECK-SAME: %[[ARG1:.*]]: vector<4x[3]xf32>171// CHECK: %[[A:.*]] = arith.constant dense<0.000000e+00> : vector<4x2x[3]xf32>172// CHECK: %[[B:.*]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0, 0, 0], sizes = [4, 1, 3], strides = [1, 1, 1]} : vector<4x2x[3]xf32> to vector<4x1x[3]xf32>173// CHECK: %[[C:.*]] = vector.shape_cast %[[ARG1]] : vector<4x[3]xf32> to vector<4x1x[3]xf32>174// CHECK: %[[D:.*]] = vector.insert_strided_slice %[[C]], %[[A]] {offsets = [0, 0, 0], strides = [1, 1, 1]} : vector<4x1x[3]xf32> into vector<4x2x[3]xf32>175// CHECK: %[[E:.*]] = arith.mulf %[[C]], %[[B]] : vector<4x1x[3]xf32>176// CHECK: %[[F:.*]] = vector.insert_strided_slice %[[E]], %[[D]] {offsets = [0, 1, 0], strides = [1, 1, 1]} : vector<4x1x[3]xf32> into vector<4x2x[3]xf32>177// CHECK: %[[G:.*]] = vector.shape_cast %[[E]] : vector<4x1x[3]xf32> to vector<4x[3]xf32>178// CHECK: return %[[F]], %[[G]] : vector<4x2x[3]xf32>, vector<4x[3]xf32>179func.func @scan3d_mul_dim1_scalable(%arg0 : vector<4x2x[3]xf32>, %arg1 : vector<4x[3]xf32>) -> (vector<4x2x[3]xf32>, vector<4x[3]xf32>) {180 %0:2 = vector.scan <mul>, %arg0, %arg1 {inclusive = false, reduction_dim = 1} :181 vector<4x2x[3]xf32>, vector<4x[3]xf32>182 return %0#0, %0#1 : vector<4x2x[3]xf32>, vector<4x[3]xf32>183}184