llvm
diff --git a/‎mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
Copy file name to clipboardExpand all lines: mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+31-50Lines changed: 31 additions & 50 deletions b/‎mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
Copy file name to clipboardExpand all lines: mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+31-50Lines changed: 31 additions & 50 deletions
diff --git a/‎mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
Copy file name to clipboardExpand all lines: mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+1-6Lines changed: 1 addition & 6 deletions b/‎mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
Copy file name to clipboardExpand all lines: mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+1-6Lines changed: 1 addition & 6 deletions
@@ -444,59 +444,40 @@ struct LinearizeVectorSplat final
 
 } // namespace
 
-/// Some operations currently will not be linearized if they have scalable
-/// vector results, although support should be added in the future. This
-/// function returns true if `op` is such an operation.
-static bool isNotLinearizableBecauseScalable(Operation *op) {
-
-  bool unsupported =
-      isa<vector::ExtractStridedSliceOp, vector::ExtractOp, vector::InsertOp>(
-          op);
-
-  // Case where linearization is possible even when there are scalable vector
-  // results.
-  if (!unsupported)
-    return false;
-
-  // Check if any of the results is a scalable vector type, and if there are
-  // return true (not linearizable).
-  auto types = op->getResultTypes();
-  bool containsScalableResult =
-      std::any_of(types.begin(), types.end(), [](Type type) {
-        auto vecType = dyn_cast<VectorType>(type);
-        return vecType && vecType.isScalable();
-      });
-
-  return containsScalableResult;
-}
-
-/// This method defines a set of operations that are not linearizable, and hence
-/// they are considered legal for the conversion target. These ops are
-/// currently,
-///
-/// 1) ones that are not in the vector dialect, are not ConstantLike, and are
-///    not Vectorizable, or
-///
-/// 2) have scalable vector results, for which support has not yet been added.
-static bool isNotLinearizable(Operation *op) {
+/// This method defines the set of operations that are linearizable, and hence
+/// that are considered illegal for the conversion target.
+static bool isLinearizable(Operation *op) {
 
+  // Operations such as builtin.module are not linearizable.
   StringLiteral vectorDialect = vector::VectorDialect::getDialectNamespace();
   StringRef opDialect = op->getDialect()->getNamespace();
-  bool unsupported = (opDialect != vectorDialect) &&
-                     !op->hasTrait<OpTrait::ConstantLike>() &&
-                     !op->hasTrait<OpTrait::Vectorizable>();
-  if (unsupported)
-    return true;
-
-  // vector.shape_cast cannot be linearized.
-  if (isa<vector::ShapeCastOp>(op))
-    return true;
-
-  // Some ops currently don't support scalable vectors.
-  if (isNotLinearizableBecauseScalable(op))
-    return true;
+  bool supported = (opDialect == vectorDialect) ||
+                   op->hasTrait<OpTrait::ConstantLike>() ||
+                   op->hasTrait<OpTrait::Vectorizable>();
+  if (!supported)
+    return false;
 
-  return false;
+  return TypeSwitch<Operation *, bool>(op)
+      // As type legalization is done with vector.shape_cast, shape_cast
+      // itself cannot be linearized (will create new shape_casts to linearize
+      // ad infinitum).
+      .Case<vector::ShapeCastOp>([&](auto) { return false; })
+      // vector.extract_strided_slice, vector.extract, and vector.insert
+      // operations are linearized to a rank-1 vector.shuffle by the current
+      // patterns. vector.shuffle only supports fixed size vectors, so it is
+      // impossible to use this approach to linearize these ops if they operate
+      // on scalable vectors.
+      .Case<vector::ExtractStridedSliceOp>(
+          [&](vector::ExtractStridedSliceOp extractOp) {
+            return !extractOp.getType().isScalable();
+          })
+      .Case<vector::InsertOp>([&](vector::InsertOp insertOp) {
+        return !insertOp.getType().isScalable();
+      })
+      .Case<vector::ExtractOp>([&](vector::ExtractOp extractOp) {
+        return !extractOp.getSourceVectorType().isScalable();
+      })
+      .Default([&](auto) { return true; });
 }
 
 void mlir::vector::populateForVectorLinearize(TypeConverter &typeConverter,
@@ -530,7 +511,7 @@ void mlir::vector::populateForVectorLinearize(TypeConverter &typeConverter,
 
   target.markUnknownOpDynamicallyLegal(
       [=](Operation *op) -> std::optional<bool> {
-        if (isNotLinearizable(op))
+        if (!isLinearizable(op))
           return true;
         // This will return true if, for all operand and result types `t`,
         // convertType(t) = t. This is true if there are no rank>=2 vectors.
 
@@ -837,8 +837,6 @@ struct TestVectorEmulateMaskedLoadStore final
   }
 };
 
-namespace bit_width_constrained_linearization {
-
 /// Get the set of operand/result types to check for sufficiently
 /// small inner-most dimension size.
 static SmallVector<std::pair<Type, unsigned>>
@@ -960,8 +958,6 @@ struct TestVectorBitWidthLinearize final
   }
 };
 
-} // namespace bit_width_constrained_linearization
-
 struct TestVectorLinearize final
     : public PassWrapper<TestVectorLinearize, OperationPass<>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorLinearize)
@@ -1069,8 +1065,7 @@ void registerTestVectorLowerings() {
 
   PassRegistration<TestVectorLinearize>();
 
-  PassRegistration<
-      bit_width_constrained_linearization::TestVectorBitWidthLinearize>();
+  PassRegistration<TestVectorBitWidthLinearize>();
 
   PassRegistration<TestEliminateVectorMasks>();
 }