[mlir][affine] SuperVectorizer only widen ops with valid result types

[caojoshua]

Summary

fixes https://github.com/llvm/llvm-project/issues/61309

[dcaballe]

Thanks for the fix! This sounds like we need a vectorization pre-condition to check that the input loop only has the supported types (e.g., no vector types). I think we should introduce that check before we go into the vectorizeLoopNest/vectorizeLoop methods, because if we bail out and we have already modified part of the input code, we may leave the IR inconsistent. Would you mind moving that check to the runOnOperation method? Thanks!

[caojoshua]

Move checks to isVectorizableLoopBodyWithOpCond()

[caojoshua]

In D147679#4249439, @dcaballe wrote:

Thanks for the fix! This sounds like we need a vectorization pre-condition to check that the input loop only has the supported types (e.g., no vector types). I think we should introduce that check before we go into the vectorizeLoopNest/vectorizeLoop methods, because if we bail out and we have already modified part of the input code, we may leave the IR inconsistent. Would you mind moving that check to the runOnOperation method? Thanks!

Yep, get your point. I instead moved the check to isVectorizableLoopBodyWithOpCond(), which seems most appropriate.

[ayzhuang]

mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp

282 ↗

(On Diff #511861)

Please also check operand types. Could be a vector_store storing some value defined outside of the loop.

[caojoshua]

mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp

282 ↗

(On Diff #511861)

Does not look like we can do this, fails a bunch of SuperVectorizer tests. I think what we have here is correct, because we only want to result type to be a valid element type.

For example, in vectorize_1d.mlir, we want to vectorize this load into a vector.transfer_read. Its result type is f32, which is valid vector element type. But its first operand, a memref, is not a valid vector type.

[ayzhuang]

mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp

282 ↗

(On Diff #511861)

You can do something similar to line 263 for memref type. But if you think the current check is enough, that's fine.

[caojoshua]

Verify ops have valid vector types

[caojoshua]

mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp

282 ↗

(On Diff #511861)

Thanks, yeah. This caught an issue a new case, where the operands have invalid type but the result is valid. For example, I added a test involving a vector.reduce, where the operand is an invalid type of vector, and the result is a valid type of float.

I think it would be good to have a case where we have:

• operand of memref with invalid vector type
• valid result type

I can't think of a scenario. I can add any recommended tests if there are any.

[ayzhuang]

LGTM, thanks!

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

static Operation *vectorizeAffineForOp(AffineForOp forOp,
   return vecForOp;
 }
 
 /// Vectorizes arbitrary operation by plain widening. We apply generic type
 /// widening of all its results and retrieve the vector counterparts for all its
 /// operands.
 static Operation *widenOp(Operation *op, VectorizationState &state) {
   SmallVector<Type, 8> vectorTypes;
-  for (Value result : op->getResults())
+  for (Value result : op->getResults()) {
+    if (!VectorType::isValidElementType(result.getType())) {
+      LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ an operation failed vectorize "
+                           "because it has an invalid result type\n");
+      return nullptr;
+    }
     vectorTypes.push_back(
         VectorType::get(state.strategy->vectorSizes, result.getType()));
+  }
 
   SmallVector<Value, 8> vectorOperands;
   for (Value operand : op->getOperands()) {
     Value vecOperand = vectorizeOperand(operand, state);
     if (!vecOperand) {
       LLVM_DEBUG(dbgs() << "\n[early-vect]+++++ an operand failed vectorize\n");
       return nullptr;
     }

mlir/test/Dialect/Affine/SuperVectorize/invalid_type.mlir

+// RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=128" -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @vector_loop
+func.func @vector_loop(%a : memref<10x20xf32>, %b : memref<10x20xf32>) {
+// CHECK: affine.for %{{.*}} = 0 to 10
+// CHECK:   affine.for %{{.*}} = 0 to 5
+// CHECK:     %{{.*}} = affine.vector_load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x20xf32>, vector<4xf32>
+// CHECK:     affine.vector_store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x20xf32>, vector<4xf32>
+// CHECK:   }
+// CHECK: }
+// CHECK: return
+  affine.for %j = 0 to 10 {
+    affine.for %i = 0 to 5 {
+      %ld0 = affine.vector_load %a[%j, %i] : memref<10x20xf32>, vector<4xf32>
+      affine.vector_store %ld0, %b[%j, %i] : memref<10x20xf32>, vector<4xf32>
+    }
+  }
+  return
+}