[MLIR] Vectorize tensor.extract on 1-d tensor

[pzread]

Summary

This patch implements the vectorization of tensor.extract for the basic 1-d lookup case.

It only vectorizes the tensor.extract to a vector.gather when the op extracts value from an 1-d tensor.

Related discussion: https://github.com/iree-org/iree/issues/9198

[pzread]

Refactor.

[dcaballe]

Thanks, @pzread! LGTM! A few comments.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

328

nit: /*name=*/

563

Have you tried to move this into vectorizeOneOp instead of using a custom hook? I might be wrong here but, IIRC, custom vectorization hooks are supposed to be exceptional cases with higher priority that may fall back to other vectorization approach if the hook pattern fails. Most of the regular operation should be vectorized directly in vectorizeOneOp.

@nicolasvasilache can probably correct me/clarify.

608–610

I don't quite get what we are trying to do here. Operands coming from outside would be broadcasted or kept scalar, as needed. What is the problem? Perhaps, elaborating a bit more the comment would help.

mlir/test/Dialect/Linalg/vectorization.mlir

1181–1183

Do we have plans to support this case soon? :)

1191

nit: remove last // -----. It will run an empty test.

[pzread]

Address comments

[pzread]

Add comments.

[pzread]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

563

The vectorizeOneOp can only handle ElementwiseMappable ops, which can take either scalar or vector as their operands. tensor::extract needs to be transformed into vector.gather, so I use a custom hook to handle it.

The comment on vectorizeOneOp also mentioned custom hook is designed to handle all non-ElementwiseMappable ops.

608–610

So the problem here is that these checks in vectorizeStaticLinalgOpPrecondition is intented to check ElementwiseMappable ops (assume all operands should be converted from scalar to vector). For tensor.extract (which isn't ElementwiseMappable), its first operand is a tensor, but we have a custom hook to transform it into vector.gather.

I changed this part and added a CustomVectorizationPrecondition list to put those custom precondition checks for custom-handled ops. Let me know WDYT.

mlir/test/Dialect/Linalg/vectorization.mlir

1181–1183

I guess we need more discussions on these cases : )

[dcaballe]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

608–610

It looks like we are implementing our ad-hoc pattern rewriter here but I guess it's ok for now.
My main concern is how this custom preconditions should compose with the existing pre-conditions. I think it makes sense for them to be auto-contained pre-conditions as they have all the context about the operation being evaluated.

mlir/test/Dialect/Linalg/vectorization.mlir

1181–1183

I think @awarzynski might be looking into this :)

[pzread]

Thanks for the review @dcaballe : ) If you are okay with this revision, could you help me commit the change? thx

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

static VectorizationResult vectorizeLinalgIndex(OpBuilder &b, Operation *op,
   SmallVector<int64_t> transposition =
       llvm::to_vector<16>(llvm::seq<int64_t>(0, linalgOp.getNumLoops()));
   std::swap(transposition.back(), transposition[indexOp.getDim()]);
   auto transposeOp =
       b.create<vector::TransposeOp>(loc, broadCastOp, transposition);
   return VectorizationResult{VectorizationStatus::NewOp, transposeOp};
 }
 
+/// Helper function to vectorize the tensor.extract operations. Returns
+/// VectorizationStatus::NewOp to signal the vectorization algorithm that it
+/// should map the produced operations. This function is meant to be used as a
+/// CustomVectorizationHook.
+static VectorizationResult
+vectorizeTensorExtract(OpBuilder &b, Operation *op, LinalgOp linalgOp,
+                       const BlockAndValueMapping &bvm) {
+  tensor::ExtractOp extractOp = dyn_cast<tensor::ExtractOp>(op);
+  if (!extractOp)
+    return VectorizationResult{VectorizationStatus::Failure, nullptr};
+  auto loc = extractOp.getLoc();
+
+  // Currently only supports extraction with an 1-D index.
+  if (extractOp.getIndices().size() > 1)
+    return VectorizationResult{VectorizationStatus::Failure, nullptr};
+
+  auto indexVec = bvm.lookup(extractOp.getIndices()[0]);
+  // Compute the static loop sizes of the extract op.
+  auto targetShape = linalgOp.computeStaticLoopSizes();
+
+  SmallVector<Value> gatherIndices;
+  gatherIndices.push_back(b.create<arith::ConstantIndexOp>(loc, 0));
+
+  auto maskConstantOp = b.create<arith::ConstantOp>(
+      loc, DenseIntElementsAttr::get(
+               VectorType::get(targetShape, b.getI1Type()), true));
+
+  auto resultType =
+      VectorType::get(targetShape, extractOp.getResult().getType());
+  auto passThruConstantOp =
+      b.create<arith::ConstantOp>(loc, b.getZeroAttr(resultType));
+
+  auto gatherOp = b.create<vector::GatherOp>(
+      loc, resultType, extractOp.getTensor(), gatherIndices, indexVec,
+      maskConstantOp, passThruConstantOp);
+
+  return VectorizationResult{VectorizationStatus::NewOp, gatherOp};
+}
+
 /// Emit reduction operations if the shapes of the value to reduce is different
 /// that the result shape.
 static Operation *reduceIfNeeded(OpBuilder &b, LinalgOp linalgOp, Operation *op,
                                  Value reduceValue, Value initialValue,
                                  const BlockAndValueMapping &bvm) {
   Value reduceVec = bvm.lookup(reduceValue);
   Value outputVec = bvm.lookup(initialValue);
   auto reduceType = reduceVec.getType().dyn_cast<VectorType>();
vectorizeAsLinalgGeneric(OpBuilder &b, LinalgOp linalgOp,
   // 4b. Register CustomVectorizationHook for indexOp.
   CustomVectorizationHook vectorizeIndex =
       [&](Operation *op,
           const BlockAndValueMapping &bvm) -> VectorizationResult {
     return vectorizeLinalgIndex(b, op, linalgOp);
   };
   hooks.push_back(vectorizeIndex);
 
+  // 4c. Register CustomVectorizationHook for extractOp.
+  CustomVectorizationHook vectorizeExtract =
+      [&](Operation *op,
+          const BlockAndValueMapping &bvm) -> VectorizationResult {
+    return vectorizeTensorExtract(b, op, linalgOp, bvm);
+  };
+  hooks.push_back(vectorizeExtract);
+
   // 5. Iteratively call `vectorizeOneOp` to each op in the slice.
   for (Operation &op : block->getOperations()) {
     VectorizationResult result = vectorizeOneOp(b, linalgOp, &op, bvm, hooks);
     if (result.status == VectorizationStatus::Failure) {
       LDBG("failed to vectorize: " << op);
       return failure();
     }
     if (result.status == VectorizationStatus::NewOp) {
if (!reduceOp || !getCombinerOpKind(reduceOp)) {
       LDBG("reduction precondition failed: reduction detection failed");
       return failure();
     }
   }
   return success();
 }
 
 static LogicalResult vectorizeStaticLinalgOpPrecondition(linalg::LinalgOp op) {
+  assert(op->getNumRegions() == 1);
+  auto &opRegion = op->getRegion(0);
   // All types in the body should be a supported element type for VectorType.
-  for (Operation &innerOp : op->getRegion(0).front()) {
-    if (llvm::any_of(innerOp.getOperandTypes(), [](Type type) {
-          return !VectorType::isValidElementType(type);
+  for (Operation &innerOp : opRegion.front()) {
+    if (llvm::any_of(innerOp.getOperands(), [&](Value operand) {
+          // Ignore the operands that come from outside, since they don't need
+          // to be vectorized.
+          return &opRegion == operand.getParentRegion() &&
+                 !VectorType::isValidElementType(operand.getType());
         })) {
       return failure();
     }
     if (llvm::any_of(innerOp.getResultTypes(), [](Type type) {
           return !VectorType::isValidElementType(type);
         })) {
       return failure();
     }

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

return llvm::all_of(op.getIndexingMapsArray(), [](AffineMap m) {
     return m.isProjectedPermutation(/*allowZeroInResults=*/true);
   });
 }
 
 bool hasOnlyScalarElementwiseOp(Region &r) {
   if (!llvm::hasSingleElement(r))
     return false;
   for (Operation &op : r.front()) {
-    if (!(isa<arith::ConstantOp, func::ConstantOp, linalg::YieldOp,
-              linalg::IndexOp>(op) ||
+    if (!(isa<arith::ConstantOp, func::ConstantOp, tensor::ExtractOp,
+              linalg::YieldOp, linalg::IndexOp>(op) ||
           OpTrait::hasElementwiseMappableTraits(&op)) ||
         llvm::any_of(op.getResultTypes(),
                      [](Type type) { return !type.isIntOrIndexOrFloat(); }))
       return false;
   }
   return true;
 }
 

mlir/test/Dialect/Linalg/vectorization.mlir

 //  CHECK-SAME:     %[[ARG3:[a-zA-Z0-9]+]]: tensor<2x4xf32>
 //   CHECK-DAG:   %[[V0:.+]] = vector.transfer_read %[[ARG0]]
 //   CHECK-DAG:   %[[V1:.+]] = vector.transfer_read %[[ARG1]]
 //   CHECK-DAG:   %[[V2:.+]] = vector.transfer_read %[[ARG3]]
 //   CHECK-DAG:   %[[MUL:.+]] = arith.mulf %[[V0]], %[[V1]]
 //   CHECK-DAG:   %[[ADD:.+]] = vector.multi_reduction <add>, %[[MUL]], %[[V2]]
 //   CHECK-DAG:   vector.transfer_write %[[MUL]], %[[ARG2]]
 //   CHECK-DAG:   vector.transfer_write %[[ADD]], %[[ARG3]]
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func.func @vectorize_1d_tensor_extract(%arg0: tensor<3xf32>, %arg1: tensor<4x3xi32>) -> tensor<4x7x3x2xf32> {
+  %0 = linalg.init_tensor [4, 7, 2] : tensor<4x7x2xf32>
+  %1 = linalg.init_tensor [4, 7, 3, 2] : tensor<4x7x3x2xf32>
+  %2 = linalg.generic {
+    indexing_maps = [#map0, #map1, #map2],
+    iterator_types = ["parallel", "parallel", "parallel", "parallel"]
+  } ins(%arg1, %0 : tensor<4x3xi32>, tensor<4x7x2xf32>) outs(%1 : tensor<4x7x3x2xf32>) {
+  ^bb0(%arg2: i32, %arg3: f32, %arg4: f32):
+    %3 = arith.index_cast %arg2 : i32 to index
+    %7 = tensor.extract %arg0[%3] : tensor<3xf32>
+    linalg.yield %7 : f32
+  } -> tensor<4x7x3x2xf32>
+  return %2 : tensor<4x7x3x2xf32>
+}
+// CHECK-LABEL: func.func @vectorize_1d_tensor_extract
+// CHECK-SAME:    %[[ARG0:.*]]: tensor<3xf32>
+// CHECK-SAME:    %[[ARG1:.*]]: tensor<4x3xi32>
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<4x7x3x2xi1>
+// CHECK: %[[PASSTHRU:.*]] = arith.constant dense<0.000000e+00> : vector<4x7x3x2xf32>
+// CHECK: %[[V0:.*]] = vector.transfer_read %[[ARG1]]
+// CHECK: %[[INDICES:.*]] = arith.index_cast %[[V0]] : vector<4x7x3x2xi32> to vector<4x7x3x2xindex>
+// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]]] [%[[INDICES]]], %[[MASK]], %[[PASSTHRU]]
+// CHECK: vector.transfer_write %[[GATHER]]
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func.func @not_vectorize_nd_tensor_extract(%arg0: tensor<3x3xf32>, %arg1: tensor<4x3xi32>, %arg2: tensor<4x3xi32>) -> tensor<4x7x3x2xf32> {
+  %0 = linalg.init_tensor [4, 7, 2] : tensor<4x7x2xf32>
+  %1 = linalg.init_tensor [4, 7, 3, 2] : tensor<4x7x3x2xf32>
+  %2 = linalg.generic {
+    indexing_maps = [#map0, #map0, #map1, #map2],
+    iterator_types = ["parallel", "parallel", "parallel", "parallel"]
+  } ins(%arg1, %arg2, %0 : tensor<4x3xi32>, tensor<4x3xi32>, tensor<4x7x2xf32>) outs(%1 : tensor<4x7x3x2xf32>) {
+  ^bb0(%arg3: i32, %arg4: i32, %arg5: f32, %arg6: f32):
+    %3 = arith.index_cast %arg3 : i32 to index
+    %4 = arith.index_cast %arg4 : i32 to index
+    %7 = tensor.extract %arg0[%3, %4] : tensor<3x3xf32>
+    linalg.yield %7 : f32
+  } -> tensor<4x7x3x2xf32>
+  return %2 : tensor<4x7x3x2xf32>
+}
+// CHECK-LABEL: func.func @not_vectorize_nd_tensor_extract
+// CHECK: tensor.extract
+
+// -----