vector outputs handling

Author: Hardcode84

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

@@ -1078,7 +1078,8 @@ SLPGraph::vectorize(IRRewriter &rewriter,
       }
     };
 
-    auto handleNonVectorOutputs = [&](Value newResult) {
+    auto handleNonVectorOutputs = [&](Value newResult,
+                                      Type originalResultType) {
       // Handle the case when op results are not vectorized or have smaller
       // vector size, extract the elements from the vector.
       for (auto [i, result] : llvm::enumerate(node->ops)) {
@@ -1087,7 +1088,16 @@ SLPGraph::vectorize(IRRewriter &rewriter,
           if (getNodeForOp(useOwner))
             continue;
 
-          Value elem = rewriter.create<vector::ExtractOp>(loc, newResult, i);
+          int64_t offset = i * node->elementsCount;
+          Value elem;
+
+          if (auto vecType = dyn_cast<VectorType>(originalResultType)) {
+            elem = rewriter.create<vector::ExtractStridedSliceOp>(
+                loc, newResult, offset, vecType.getNumElements(), 1);
+          } else {
+            elem = rewriter.create<vector::ExtractOp>(loc, newResult, offset);
+          }
+
           use.set(elem);
         }
       }
@@ -1109,8 +1119,9 @@ SLPGraph::vectorize(IRRewriter &rewriter,
           VectorType::get(numElements, getElementTypeAndCount(op)->first);
       Value result = rewriter.create<vector::LoadOp>(loc, vecType, getBase(op),
                                                      getIndices(op));
-      mapping.map(op->getResult(0), result);
-      handleNonVectorOutputs(result);
+      Value originalResult = op->getResult(0);
+      mapping.map(originalResult, result);
+      handleNonVectorOutputs(result, originalResult.getType());
     } else if (maybeWriteOp(op)) {
       handleNonVectorInputs(getValueToStore(op));
       Value val = mapping.lookupOrDefault(getValueToStore(op));
@@ -1133,7 +1144,7 @@ SLPGraph::vectorize(IRRewriter &rewriter,
       newOp->getResult(0).setType(resVectorType);
 
       mapping.map(op->getResults(), newOp->getResults());
-      handleNonVectorOutputs(newOp->getResult(0));
+      handleNonVectorOutputs(newOp->getResult(0), op->getResultTypes().front());
     } else if (auto extract = dyn_cast<vector::ExtractOp>(op)) {
       // We alredy verified index is valid during graph construction, so
       // do need to check `getExtractIndex` result.

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

@@ -672,6 +672,32 @@ func.func @read_read_add_add(%arg0: memref<8xi32>, %arg1: memref<8xi32>,
 }
 
 
+// CHECK-LABEL: func @read_read_add_add_vec
+//  CHECK-SAME: (%[[ARG0:.*]]: memref<8xi32>, %[[ARG1:.*]]: memref<8xi32>)
+func.func @read_read_add_add_vec(%arg0: memref<8xi32>, %arg1: memref<8xi32>) ->
+                                   (vector<2xi32>, vector<2xi32>){
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK: %[[V0:.*]] = vector.load %[[ARG0]][%[[C0]]] : memref<8xi32>, vector<4xi32>
+  // CHECK: %[[V1:.*]] = vector.load %[[ARG1]][%[[C0]]] : memref<8xi32>, vector<4xi32>
+  // CHECK: %[[V2:.*]] = arith.addi %[[V0]], %[[V1]] : vector<4xi32>
+  // CHECK: %[[V3:.*]] = vector.extract_strided_slice %[[V2]] {offsets = [0], sizes = [2], strides = [1]} : vector<4xi32> to vector<2xi32>
+  // CHECK: %[[V4:.*]] = vector.extract_strided_slice %[[V2]] {offsets = [2], sizes = [2], strides = [1]} : vector<4xi32> to vector<2xi32>
+  // CHECK: return %[[V3]], %[[V4]] : vector<2xi32>, vector<2xi32>
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+
+  %0 = vector.load %arg0[%c0] : memref<8xi32>, vector<2xi32>
+  %2 = vector.load %arg0[%c2] : memref<8xi32>, vector<2xi32>
+
+  %4 = vector.load %arg1[%c0] : memref<8xi32>, vector<2xi32>
+  %6 = vector.load %arg1[%c2] : memref<8xi32>, vector<2xi32>
+
+  %8 = arith.addi %0, %4 : vector<2xi32>
+  %10 = arith.addi %2, %6 : vector<2xi32>
+
+  return %8, %10 : vector<2xi32>, vector<2xi32>
+}
+
 
 func.func private @use(i32)