further enhancements

Author: newling

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

@@ -409,16 +409,20 @@ void populateVectorTransposeNarrowTypeRewritePatterns(
 /// Initialize `typeConverter` and `conversionTarget` for vector linearization.
 ///
 /// 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.
+/// 1+ vector operand/result of rank>1, into a new single operation whose
+/// vector operands and results 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
-/// certain rank>1 vectors are considered valid, but adding additional
+/// certain rank>1 vectors are considered valid, by adding additional
 /// dynamically legal ops to `conversionTarget`.
+///
+/// Further note: the choice to use a dialect conversion design for
+/// linearization is to make it easy to reuse generic structural type
+/// conversions for linearizing scf/cf/func operations
 void populateForVectorLinearize(TypeConverter &typeConverter,
                                 ConversionTarget &conversionTarget);
 

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

@@ -444,8 +444,9 @@ struct LinearizeVectorSplat final
 
 } // namespace
 
-/// Some operations currently cannot be linearized if they have scalable vector
-/// results. This function returns true if `op` is such an operation.
+/// 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 =
@@ -469,15 +470,14 @@ 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
+/// 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) Ops that are not in the vector dialect, are not ConstantLike, and are not
-///    Vectorizable.
+/// 1) ones that are not in the vector dialect, are not ConstantLike, and are
+///    not Vectorizable, or
 ///
-/// 2) Certain ops with scalable vector results, for which support has not yet
-///    been added.
+/// 2) have scalable vector results, for which support has not yet been added.
 static bool isNotLinearizable(Operation *op) {
 
   StringLiteral vectorDialect = vector::VectorDialect::getDialectNamespace();
@@ -488,6 +488,10 @@ static bool isNotLinearizable(Operation *op) {
   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;

mlir/test/Dialect/Vector/linearize.mlir

@@ -320,6 +320,28 @@ func.func @test_vector_bitcast(%arg0: vector<[4]x2xf32>) -> vector<[4]x4xf16> {
   return %1 : vector<[4]x4xf16>
 }
 
+// ----- 
+
+// CHECK-LABEL: test_linearize_across_for
+func.func @test_linearize_across_for(%arg0 : vector<4xi8>) -> vector<4xi8> {
+  %0 = vector.shape_cast %arg0 : vector<4xi8> to vector<2x2xi8>
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+
+  // CHECK:  scf.for {{.*}} -> (vector<4xi8>)
+  %1 = scf.for %i = %c0 to %c4 step %c1 iter_args(%arg1 = %0) -> (vector<2x2xi8>) {
+
+    // CHECK:  arith.addi {{.*}} : vector<4xi8>
+    %2 = arith.addi %arg1, %0 : vector<2x2xi8>
+
+    // CHECK:  scf.yield {{.*}} : vector<4xi8>
+    scf.yield %2 : vector<2x2xi8>
+  }
+  %3 = vector.shape_cast %1 : vector<2x2xi8> to vector<4xi8>
+  return %3 : vector<4xi8>
+}
+
 // -----
 
 // CHECK-LABEL: linearize_vector_splat
@@ -344,4 +366,5 @@ func.func @linearize_scalable_vector_splat(%arg0: i32) -> vector<4x[2]xi32> {
   // CHECK: return %[[CAST]] : vector<4x[2]xi32>
   %0 = vector.splat %arg0 : vector<4x[2]xi32>
   return %0 : vector<4x[2]xi32>
+
 }

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

@@ -837,7 +837,7 @@ struct TestVectorEmulateMaskedLoadStore final
   }
 };
 
-namespace bit_width_constrained_vector_linearize {
+namespace bit_width_constrained_linearization {
 
 /// Get the set of operand/result types to check for sufficiently
 /// small inner-most dimension size.
@@ -960,7 +960,7 @@ struct TestVectorBitWidthLinearize final
   }
 };
 
-} // namespace bit_width_constrained_vector_linearize
+} // namespace bit_width_constrained_linearization
 
 struct TestVectorLinearize final
     : public PassWrapper<TestVectorLinearize, OperationPass<>> {
@@ -989,7 +989,6 @@ struct TestVectorLinearize final
     vector::populateVectorLinearizeBasePatterns(converter, target, patterns);
     vector::populateVectorLinearizeShuffleLikeOpsPatterns(converter, target,
                                                           patterns);
-
     mlir::scf::populateSCFStructuralTypeConversionsAndLegality(
         converter, patterns, target);
 
@@ -1073,7 +1072,7 @@ void registerTestVectorLowerings() {
   PassRegistration<TestVectorLinearize>();
 
   PassRegistration<
-      bit_width_constrained_vector_linearize::TestVectorBitWidthLinearize>();
+      bit_width_constrained_linearization::TestVectorBitWidthLinearize>();
 
   PassRegistration<TestEliminateVectorMasks>();
 }