[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

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

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

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

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/Analysis/LoopAnalysis.cpp

isVectorizableLoopBodyWithOpCond(AffineForOp loop,
   // No vectorization across conditionals for now.
   auto conditionals = matcher::If();
   SmallVector<NestedMatch, 8> conditionalsMatched;
   conditionals.match(forOp, &conditionalsMatched);
   if (!conditionalsMatched.empty()) {
     return false;
   }
 
+  // No vectorization for ops with operand or result types that are not
+  // vectorizable.
+  auto types = matcher::Op([](Operation &op) -> bool {
+    if (llvm::any_of(op.getOperandTypes(), [](Type type) {
+          if (MemRefType t = dyn_cast<MemRefType>(type))
+            return !VectorType::isValidElementType(t.getElementType());
+          return !VectorType::isValidElementType(type);
+        }))
+      return true;
+    return llvm::any_of(op.getResultTypes(), [](Type type) {
+      return !VectorType::isValidElementType(type);
+    });
+  });
+  SmallVector<NestedMatch, 8> opsMatched;
+  types.match(forOp, &opsMatched);
+  if (!opsMatched.empty()) {
+    return false;
+  }
+
   // No vectorization across unknown regions.
   auto regions = matcher::Op([](Operation &op) -> bool {
     return op.getNumRegions() != 0 && !isa<AffineIfOp, AffineForOp>(op);
   });
   SmallVector<NestedMatch, 8> regionsMatched;
   regions.match(forOp, &regionsMatched);
   if (!regionsMatched.empty()) {
     return false;

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 @invalid_operand
+func.func @invalid_operand(%a : vector<4xf32>, %b : vector<4xf32>) {
+// CHECK: affine.for %{{.*}} = 0 to 10
+// CHECK:   %{{.*}} = vector.reduction <add>, %{{.*}} : vector<4xf32> into f32
+// CHECK: }
+// CHECK: return
+  affine.for %j = 0 to 10 {
+    %1 = vector.reduction <add>, %a : vector<4xf32> into f32
+  }
+  return
+}
+
+// CHECK-LABEL: func @invalid_result
+func.func @invalid_result(%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
+}