first commit (needs refinement)

Author: newling

mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h

@@ -407,8 +407,13 @@ void populateVectorTransposeNarrowTypeRewritePatterns(
     RewritePatternSet &patterns, PatternBenefit benefit = 1);
 
 /// Initialize `typeConverter` and `conversionTarget` for vector linearization.
-/// This registers (1) which operations are legal and hence should not be
-/// linearized, (2) what converted types are (rank-1 vectors) and how to
+///
+/// Definition: here 'linearization' means converting a single operation with
+/// 1+ vector operands and results of rank>1, into a single operation whose
+/// vector operands are all of rank<=1.
+///
+/// This function registers (1) which operations are legal, and hence should not
+/// be linearized, (2) what the converted types are (rank-1 vectors) and how to
 /// materialze the conversion (with shape_cast)
 ///
 /// Note: the set of legal operations can be extended by a user if for example

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

@@ -99,7 +99,7 @@ class ConvertForOpTypes
     // PR47938 tracks this issue, but it seems hard to fix. Instead, we need
     // to clone the op.
     //
-    // 2. We need to resue the original region instead of cloning it, otherwise
+    // 2. We need to reuse the original region instead of cloning it, otherwise
     // the dialect conversion framework thinks that we just inserted all the
     // cloned child ops. But what we want is to "take" the child regions and let
     // the dialect conversion framework continue recursively into ops inside

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

@@ -134,9 +134,6 @@ struct LinearizeVectorExtractStridedSlice final
     VectorType dstType =
         getTypeConverter()->convertType<VectorType>(extractOp.getType());
     assert(dstType && "vector type destination expected.");
-    if (extractOp.getVector().getType().isScalable() || dstType.isScalable())
-      return rewriter.notifyMatchFailure(extractOp,
-                                         "scalable vectors are not supported.");
 
     ArrayAttr offsets = extractOp.getOffsets();
     ArrayAttr sizes = extractOp.getSizes();
@@ -447,18 +444,21 @@ struct LinearizeVectorSplat final
 
 } // namespace
 
-/// Return true if the operation `op` does not support scalable vectors and
-/// has at least 1 scalable vector result. These ops should all eventually
-/// support scalable vectors, and this function should be removed.
+/// Some operations currently cannot be linearized if they have scalable vector
+/// results. 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.
+  // 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) {
@@ -469,10 +469,17 @@ static bool isNotLinearizableBecauseScalable(Operation *op) {
   return containsScalableResult;
 }
 
+/// This method defines a set of operations that are not linearizable,
+/// and hence considered legal for the conversion target. These ops are
+/// currently
+///
+/// 1) Ops that are not in the vector dialect, are not ConstantLike, and are not
+///    Vectorizable.
+///
+/// 2) Certain ops with scalable vector results, for which support has not yet
+///    been added.
 static bool isNotLinearizable(Operation *op) {
 
-  // Only ops that are in the vector dialect, are ConstantLike, or
-  // are Vectorizable might be linearized currently.
   StringLiteral vectorDialect = vector::VectorDialect::getDialectNamespace();
   StringRef opDialect = op->getDialect()->getNamespace();
   bool unsupported = (opDialect != vectorDialect) &&

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

@@ -17,6 +17,7 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Transforms/Patterns.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
@@ -836,8 +837,7 @@ struct TestVectorEmulateMaskedLoadStore final
   }
 };
 
-// TODO: move this code into the user project.
-namespace vendor {
+namespace bit_width_constrained_vector_linearize {
 
 /// Get the set of operand/result types to check for sufficiently
 /// small inner-most dimension size.
@@ -960,7 +960,7 @@ struct TestVectorBitWidthLinearize final
   }
 };
 
-} // namespace vendor
+} // namespace bit_width_constrained_vector_linearize
 
 struct TestVectorLinearize final
     : public PassWrapper<TestVectorLinearize, OperationPass<>> {
@@ -982,12 +982,17 @@ struct TestVectorLinearize final
     RewritePatternSet patterns(&context);
     ConversionTarget target(context);
 
+    SmallVector<Operation *> ops;
+
     vector::populateForVectorLinearize(converter, target);
 
     vector::populateVectorLinearizeBasePatterns(converter, target, patterns);
     vector::populateVectorLinearizeShuffleLikeOpsPatterns(converter, target,
                                                           patterns);
 
+    mlir::scf::populateSCFStructuralTypeConversionsAndLegality(
+        converter, patterns, target);
+
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       return signalPassFailure();
@@ -1067,7 +1072,8 @@ void registerTestVectorLowerings() {
 
   PassRegistration<TestVectorLinearize>();
 
-  PassRegistration<vendor::TestVectorBitWidthLinearize>();
+  PassRegistration<
+      bit_width_constrained_vector_linearize::TestVectorBitWidthLinearize>();
 
   PassRegistration<TestEliminateVectorMasks>();
 }