review comments

Author: Hardcode84

mlir/include/mlir/Dialect/Arith/Utils/Utils.h

@@ -121,8 +121,8 @@ struct ArithBuilder {
 };
 
 /// ArithBuilder specialized specifically for tensor/memref indexing
-/// calculations. Those calculations generally should never signed overflow, so
-/// we can set oveflow flags accordingly.
+/// calculations. Those calculations generally should never signed overflow and
+/// always use signed integers, so we can set oveflow flags accordingly.
 struct ArithIndexingBuilder : public ArithBuilder {
   ArithIndexingBuilder(OpBuilder &b, Location loc)
       : ArithBuilder(b, loc, arith::IntegerOverflowFlags::nsw) {}

mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td

@@ -458,7 +458,9 @@ def ApplySinkVectorPatternsOp : Op<Transform_Dialect,
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
     Patterns that remove redundant Vector Ops by re-ordering them with
-    e.g. elementwise Ops:
+    e.g. elementwise Ops.
+
+    Example:
     ```
     %at = vector.transpose %a, [1, 0]: vector<4x2xf32> to vector<2x4xf32>
     %bt = vector.transpose %b, [1, 0]: vector<4x2xf32> to vector<2x4xf32>
@@ -480,8 +482,11 @@ def ApplySinkVectorMemPatternsOp : Op<Transform_Dialect,
     "apply_patterns.vector.sink_mem_ops",
     [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
   let description = [{
-    Patterns that remove redundant Vector Ops by merging them with load/store
-    ops
+    Patterns that replace redundant Vector Ops (followed by
+    `vector.load`/`vector.store`) with either vector.load/vector.store or
+    `memref.load`/`memref.store`. Currently limited to 1-element vectors.
+
+    Example:
     ```
     vector.load %arg0[%arg1] : memref<?xf32>, vector<4xf32>
     vector.extract %0[1] : f32 from vector<4xf32>

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

@@ -1047,15 +1047,18 @@ class ExtractOpFromElementwise final
   }
 };
 
+/// Check if the element type is suitable for vector.load/store sinking.
+/// Element type must be index or byte-aligned integer or floating-point type.
 static bool isSupportedMemSinkElementType(Type type) {
   if (isa<IndexType>(type))
     return true;
 
-  // Non-byte-aligned types are tricky, skip them.
   return type.isIntOrFloat() && type.getIntOrFloatBitWidth() % 8 == 0;
 }
 
-/// Pattern to rewrite vector.extract(vector.load) -> vector/memref.load.
+/// Pattern to rewrite `vector.extract(vector.load) -> vector/memref.load.
+/// Only index and byte-aligned integer and floating-point element types are
+/// supported for now.
 ///
 /// Example:
 /// ```
@@ -1088,8 +1091,11 @@ class ExtractOpFromLoad final : public OpRewritePattern<vector::ExtractOp> {
                                          "scalable vectors are not supported");
 
     MemRefType memType = loadOp.getMemRefType();
+
+    // Non-byte-aligned types are tricky and may require special handling,
+    // ignore them for now.
     if (!isSupportedMemSinkElementType(memType.getElementType()))
-      return rewriter.notifyMatchFailure(op, "unsupported memref element type");
+      return rewriter.notifyMatchFailure(op, "unsupported element type");
 
     int64_t rankOffset = memType.getRank() - loadVecType.getRank();
     if (rankOffset < 0)
@@ -1161,7 +1167,7 @@ class StoreOpFromSplatOrBroadcast final
 
     if (vecType.getNumElements() != 1)
       return rewriter.notifyMatchFailure(
-          op, "only 1-element, vectors are supported");
+          op, "only 1-element vectors are supported");
 
     Operation *splat = op.getValueToStore().getDefiningOp();
     if (!isa_and_present<vector::BroadcastOp, vector::SplatOp>(splat))
@@ -2253,6 +2259,7 @@ void mlir::vector::populateSinkVectorOpsPatterns(RewritePatternSet &patterns,
 
 void mlir::vector::populateSinkVectorMemOpsPatterns(RewritePatternSet &patterns,
                                                     PatternBenefit benefit) {
+  // TODO: Consider converting these patterns to canonicalizations.
   patterns.add<ExtractOpFromLoad, StoreOpFromSplatOrBroadcast>(
       patterns.getContext(), benefit);
 }