[Mlir-commits] [mlir] f18a861 - [mlir][Vector] Enable masked vectorization of linalg.fill

[Diego Caballero]

Author: Diego Caballero
Date: 2023-03-29T19:53:29Z
New Revision: f18a8612995e3b4b7af9d7430374915724cdde51

URL: https://github.com/llvm/llvm-project/commit/f18a8612995e3b4b7af9d7430374915724cdde51
DIFF: https://github.com/llvm/llvm-project/commit/f18a8612995e3b4b7af9d7430374915724cdde51.diff

LOG: [mlir][Vector] Enable masked vectorization of linalg.fill

linalg.fill was already vectorizable with masks but not supported in the
dynamic pre-checks.

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D146856

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
    mlir/test/Dialect/Linalg/vectorization.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 6b27b412f183e..98ee5e2889758 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -1291,19 +1291,13 @@ static LogicalResult reductionPreconditions(LinalgOp op) {
 
 static LogicalResult vectorizeDynamicLinalgOpPrecondition(linalg::LinalgOp op) {
   // TODO: Masking only supports dynamic generic ops for now.
-  if (!isa<linalg::GenericOp>(op))
+  if (!isa<linalg::GenericOp, linalg::FillOp>(op))
     return failure();
 
   // TODO: Index vectorization assumes static shape.
   if (op.hasIndexSemantics())
     return failure();
 
-  // TODO: 0-d vectors are not supported yet.
-  if (llvm::any_of(op.getIndexingMapsArray(), [](AffineMap map) {
-        return map.isEmpty() || map.getResults().empty();
-      }))
-    return failure();
-
   LDBG("Dynamically-shaped op meets vectorization pre-conditions\n");
   return success();
 }

diff  --git a/mlir/test/Dialect/Linalg/vectorization.mlir b/mlir/test/Dialect/Linalg/vectorization.mlir
index 26e27c108ce81..105d95225fa6c 100644
--- a/mlir/test/Dialect/Linalg/vectorization.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization.mlir
@@ -2535,3 +2535,24 @@ transform.sequence failures(propagate) {
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   transform.structured.masked_vectorize %0 vector_sizes [8, 32]
 }
+
+// -----
+
+func.func @vectorize_dynamic_fill(%A : tensor<?x?xf32>, %arg0 : f32) -> tensor<?x?xf32> {
+  %0 = linalg.fill ins(%arg0 : f32) outs(%A : tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL: func.func @vectorize_dynamic_fill
+//   CHECK: %[[DIM0:.*]] = tensor.dim
+//   CHECK: %[[DIM1:.*]] = tensor.dim
+//   CHECK: %[[MASK:.*]] = vector.create_mask %[[DIM0]], %[[DIM1]] : vector<8x16xi1>
+//   CHECK: %[[BCAST:.*]] = vector.broadcast %{{.*}} : f32 to vector<8x16xf32>
+//   CHECK: vector.mask %[[MASK]] { vector.transfer_write %[[BCAST]], {{.*}} {in_bounds = [true, true]} : vector<8x16xf32>, tensor<?x?xf32> } : vector<8x16xi1>
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  transform.structured.masked_vectorize %0 vector_sizes [8, 16]
+}
+