address reviwer comments

Author: nirvedhmeshram

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

@@ -810,10 +810,12 @@ static Value calculateGatherOffset(RewriterBase &rewriter,
 
 enum VectorMemoryAccessKind { ScalarBroadcast, Contiguous, Gather };
 
-
-/// Find the non constant dim in a linalgOp. This is used for finding contiguous
-/// loads and it is expected that only one dim will be non constant, if thats
-/// not the case this function will assert.
+/// Find the non-unit dim in a linalgOp.
+/// When executing this hook, it is expected that only one dim will be non-unit.
+/// Other cases (i.e. reading n-D vectors) should've been labelled as gather
+/// loads before calling this method. This is used for finding contiguous loads
+/// (represented as `tensor.extract`) within `linalg.generic` Ops. Note that
+/// this condition is expected to hold for statically shaped Linalg Ops only.
 static uint64_t getNonUnitLoopDim(LinalgOp linalgOp) {
   uint64_t nonUnitDim = 0;
   uint64_t countNonUnitDim = 0;
@@ -823,8 +825,10 @@ static uint64_t getNonUnitLoopDim(LinalgOp linalgOp) {
       countNonUnitDim++;
     }
   }
-  assert(countNonUnitDim == 1 &&
-         "Expected only one non unit loop dim in this linalg op");
+
+  assert(linalgOp.hasDynamicShape() ||
+         countNonUnitDim == 1 && "For statically shaped Linalg Ops, only one "
+                                 "non-unit loop dim is expected");
   return nonUnitDim;
 }
 
@@ -908,6 +912,8 @@ static bool isContiguousLoadIdx(LinalgOp &linalgOp, Value &val,
   Operation *defOp = val.getDefiningOp();
   assert(defOp && "This is neither a block argument nor an operation result");
 
+  // Given the assumption on the loop ranges above, we expect only 1 non-unit
+  // loop dim.
   auto nonUnitLoopDim = getNonUnitLoopDim(linalgOp);
 
   if (auto indexOp = dyn_cast<linalg::IndexOp>(defOp)) {

mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir

@@ -258,7 +258,7 @@ module attributes {transform.with_named_sequence} {
 
 #map = affine_map<(d0, d1) -> (d0, d1)>
 #map1 = affine_map<(d0, d1, d2) -> (d0 + d1 + d2)>
-func.func @vectorize_nd_tensor_extract_without_outer_unit_dim(%arg0: tensor<8x128x768xf32>, %arg1 : index) -> tensor<8x1xf32> {
+func.func @vectorize_nd_tensor_extract_load_1d_column_vector_using_gather_load(%arg0: tensor<8x128x768xf32>, %arg1 : index) -> tensor<8x1xf32> {
   %c0 = arith.constant 0 : index
   %0 = tensor.empty() : tensor<8x1xf32>
   %1 = linalg.generic {
@@ -276,15 +276,15 @@ func.func @vectorize_nd_tensor_extract_without_outer_unit_dim(%arg0: tensor<8x12
 }
 
 module attributes {transform.with_named_sequence} {
-  transform.named_sequence @__transform_main(%arg2: !transform.any_op {transform.readonly}) {
-    %0 = transform.structured.match ops{["linalg.generic"]} in %arg2 : (!transform.any_op) -> !transform.any_op
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op
     %1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
     %2 = transform.structured.vectorize_children_and_apply_patterns %1 {vectorize_nd_extract} : (!transform.any_op) -> !transform.any_op
     transform.yield
   }
 }
 
-// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_without_outer_unit_dim
+// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_load_1d_column_vector_using_gather_load
 // CHECK-SAME: %[[ARG0:.*]]: tensor<8x128x768xf32>
 // CHECK-SAME: %[[ARG1:.*]]: index
 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index