vector handling

Author: Hardcode84

mlir/lib/Dialect/Vector/Transforms/SLPVectorizer.cpp

@@ -1092,8 +1092,14 @@ SLPGraph::vectorize(IRRewriter &rewriter,
           Value elem;
 
           if (auto vecType = dyn_cast<VectorType>(originalResultType)) {
-            elem = rewriter.create<vector::ExtractStridedSliceOp>(
-                loc, newResult, offset, vecType.getNumElements(), 1);
+            assert(vecType.getRank() <= 1);
+            if (vecType.getRank() == 0) {
+              elem = rewriter.create<vector::ExtractOp>(loc, newResult, offset);
+              elem = rewriter.create<vector::SplatOp>(loc, vecType, elem);
+            } else {
+              elem = rewriter.create<vector::ExtractStridedSliceOp>(
+                  loc, newResult, offset, vecType.getNumElements(), 1);
+            }
           } else {
             elem = rewriter.create<vector::ExtractOp>(loc, newResult, offset);
           }

mlir/test/Dialect/Vector/slp-vectorize.mlir

@@ -699,6 +699,62 @@ func.func @read_read_add_add_vec(%arg0: memref<8xi32>, %arg1: memref<8xi32>) ->
 }
 
 
+// CHECK-LABEL: func @read_read_add_add_vec1
+//  CHECK-SAME: (%[[ARG0:.*]]: memref<8xi32>, %[[ARG1:.*]]: memref<8xi32>)
+func.func @read_read_add_add_vec1(%arg0: memref<8xi32>, %arg1: memref<8xi32>) ->
+                                   (vector<1xi32>, vector<1xi32>){
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK: %[[V0:.*]] = vector.load %[[ARG0]][%[[C0]]] : memref<8xi32>, vector<2xi32>
+  // CHECK: %[[V1:.*]] = vector.load %[[ARG1]][%[[C0]]] : memref<8xi32>, vector<2xi32>
+  // CHECK: %[[V2:.*]] = arith.addi %[[V0]], %[[V1]] : vector<2xi32>
+  // CHECK: %[[V3:.*]] = vector.extract_strided_slice %[[V2]] {offsets = [0], sizes = [1], strides = [1]} : vector<2xi32> to vector<1xi32>
+  // CHECK: %[[V4:.*]] = vector.extract_strided_slice %[[V2]] {offsets = [1], sizes = [1], strides = [1]} : vector<2xi32> to vector<1xi32>
+  // CHECK: return %[[V3]], %[[V4]] : vector<1xi32>, vector<1xi32>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+
+  %0 = vector.load %arg0[%c0] : memref<8xi32>, vector<1xi32>
+  %2 = vector.load %arg0[%c1] : memref<8xi32>, vector<1xi32>
+
+  %4 = vector.load %arg1[%c0] : memref<8xi32>, vector<1xi32>
+  %6 = vector.load %arg1[%c1] : memref<8xi32>, vector<1xi32>
+
+  %8 = arith.addi %0, %4 : vector<1xi32>
+  %10 = arith.addi %2, %6 : vector<1xi32>
+
+  return %8, %10 : vector<1xi32>, vector<1xi32>
+}
+
+
+// CHECK-LABEL: func @read_read_add_add_vec0d
+//  CHECK-SAME: (%[[ARG0:.*]]: memref<8xi32>, %[[ARG1:.*]]: memref<8xi32>)
+func.func @read_read_add_add_vec0d(%arg0: memref<8xi32>, %arg1: memref<8xi32>) ->
+                                   (vector<i32>, vector<i32>){
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK: %[[V0:.*]] = vector.load %[[ARG0]][%[[C0]]] : memref<8xi32>, vector<2xi32>
+  // CHECK: %[[V1:.*]] = vector.load %[[ARG1]][%[[C0]]] : memref<8xi32>, vector<2xi32>
+  // CHECK: %[[V2:.*]] = arith.addi %[[V0]], %[[V1]] : vector<2xi32>
+  // CHECK: %[[V3:.*]] = vector.extract %[[V2]][0] : i32 from vector<2xi32>
+  // CHECK: %[[V4:.*]] = vector.splat %[[V3]] : vector<i32>
+  // CHECK: %[[V5:.*]] = vector.extract %[[V2]][1] : i32 from vector<2xi32>
+  // CHECK: %[[V6:.*]] = vector.splat %[[V5]] : vector<i32>
+  // CHECK: return %[[V4]], %[[V6]] : vector<i32>, vector<i32>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+
+  %0 = vector.load %arg0[%c0] : memref<8xi32>, vector<i32>
+  %2 = vector.load %arg0[%c1] : memref<8xi32>, vector<i32>
+
+  %4 = vector.load %arg1[%c0] : memref<8xi32>, vector<i32>
+  %6 = vector.load %arg1[%c1] : memref<8xi32>, vector<i32>
+
+  %8 = arith.addi %0, %4 : vector<i32>
+  %10 = arith.addi %2, %6 : vector<i32>
+
+  return %8, %10 : vector<i32>, vector<i32>
+}
+
+
 func.func private @use(i32)
 
 // CHECK-LABEL: func @read_read_add_write_interleaved_use