Disable contraction vectorization with broadcasts

Author: adam-smnk

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

@@ -2103,30 +2103,37 @@ vectorizeInsertSliceOpPrecondition(tensor::InsertSliceOp sliceOp,
 ///   vector::TransferWriteOp - Write the result vector back to the
 ///   destination
 /// The operands shapes are preserved and loaded directly into vectors.
-/// Any further permutations or numerical casting remain within contraction.
+/// Any further permutations or numerical casting remain within contraction op.
 static LogicalResult
 vectorizeAsLinalgContraction(RewriterBase &rewriter, VectorizationState &state,
                              LinalgOp linalgOp,
                              SmallVectorImpl<Value> &newResults) {
   Location loc = linalgOp.getLoc();
   MLIRContext *ctx = linalgOp.getContext();
 
+  // For simplicity, contraction vectorization is limited to linalg named ops.
+  // Generic op is ignored as not every arbitrary contraction body can be
+  // expressed by a vector.contract.
   if (!isa<ContractionOpInterface>(linalgOp.getOperation()))
     return failure();
 
   OpOperand *outOperand = linalgOp.getDpsInitOperand(0);
   Operation *reduceOp = matchLinalgReduction(outOperand);
   auto maybeKind = getCombinerOpKind(reduceOp);
-  if (!maybeKind)
+  if (!maybeKind) {
+    LDBG("Failed to determine contraction combining kind.\n");
     return failure();
+  }
 
   // Check that all dimensions are present in the input operands.
   // Arbitrary broadcasts are not supported by the vector contraction.
-  // Broadcasts are expected to be materialized before vectorization.
+  // Broadcasts are expected to be decomposed before vectorization.
   AffineMap lhsMap = linalgOp.getIndexingMapsArray()[0];
   AffineMap rhsMap = linalgOp.getIndexingMapsArray()[1];
-  if (getUnusedDimsBitVector({lhsMap, rhsMap}).any())
+  if (getUnusedDimsBitVector({lhsMap, rhsMap}).any()) {
+    LDBG("Contractions with broadcasts are not supported.\n");
     return failure();
+  }
 
   // Load operands.
   SmallVector<Value> vecOperands;
@@ -2659,20 +2666,10 @@ FailureOr<VectorizationResult> mlir::linalg::vectorize(
               return failure();
             }
 
-            // For simplicity, contraction vectorization is limited to linalg
-            // named ops. Generic op is ignored as not every arbitrary
-            // contraction body can be expressed by a vector.contract.
             if (createNamedContraction &&
-                isa<ContractionOpInterface>(linalgOp.getOperation())) {
-              // Attempt vectorizing directly into a named contraction.
-              // In case of failure, fall back to the generic path.
-              LogicalResult res = vectorizeAsLinalgContraction(
-                  rewriter, state, linalgOp, results);
-              if (succeeded(res))
-                return success();
-
-              LDBG("Failed to vectorize as a named contraction.\n");
-            }
+                isa<ContractionOpInterface>(linalgOp.getOperation()))
+              return vectorizeAsLinalgContraction(rewriter, state, linalgOp,
+                                                  results);
 
             LDBG("Vectorize generic by broadcasting to the canonical vector "
                  "shape\n");

mlir/test/Dialect/Linalg/vectorization/contraction-interface.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -transform-interpreter -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -transform-interpreter -split-input-file -verify-diagnostics | FileCheck %s
 
 ///----------------------------------------------------------------------------------------
 /// Tests for vectorizing operations implementing contraction op interface.
@@ -214,14 +214,15 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-/// Contractions' arbitrarty broadcasts are not supported in contraction interface
+/// Contractions with arbitrarty broadcasts are not supported in contraction interface
 /// vectorization.
 /// Dimension broadcasts are expected to be decomposed first which removes ambiguity
 /// caused by possible variants of dimensions materialization.
 /// For example, whether the below target LHS input layout is (m, k) or (k, m).
 
 func.func @negative_matmul_broadcast(%A: tensor<4xf32>, %B: tensor<4x16xf32>,
     %C: tensor<8x16xf32>) -> tensor<8x16xf32> {
+  // expected-error @+1 {{Attempted to vectorize, but failed}}
   %0 = linalg.matmul
     indexing_maps = [affine_map<(m, n, k) -> (k)>, // broadcast
                      affine_map<(m, n, k) -> (k, n)>,
@@ -231,10 +232,6 @@ func.func @negative_matmul_broadcast(%A: tensor<4xf32>, %B: tensor<4x16xf32>,
   return %0 : tensor<8x16xf32>
 }
 
-// CHECK-LABEL: func.func @negative_matmul_broadcast(
-// CHECK-NOT: vector.contract
-// CHECK: vector.multi_reduction
-
 module attributes {transform.with_named_sequence} {
   transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
     %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op