ignore non-byte-aligned types

Author: Hardcode84

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

@@ -1047,6 +1047,14 @@ class ExtractOpFromElementwise final
   }
 };
 
+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.
 ///
 /// Example:
@@ -1080,9 +1088,8 @@ class ExtractOpFromLoad final : public OpRewritePattern<vector::ExtractOp> {
                                          "scalable vectors are not supported");
 
     MemRefType memType = loadOp.getMemRefType();
-    if (isa<VectorType>(memType.getElementType()))
-      return rewriter.notifyMatchFailure(
-          op, "memrefs of vectors are not supported");
+    if (!isSupportedMemSinkElementType(memType.getElementType()))
+      return rewriter.notifyMatchFailure(op, "unsupported memref element type");
 
     int64_t rankOffset = memType.getRank() - loadVecType.getRank();
     if (rankOffset < 0)

mlir/test/Dialect/Vector/vector-sink.mlir

@@ -528,6 +528,16 @@ func.func @extract_load_scalar(%arg0: memref<?xf32>, %arg1: index) -> f32 {
   return %1 : f32
 }
 
+// CHECK-LABEL: @extract_load_index
+//  CHECK-SAME:   (%[[ARG0:.*]]: memref<?xindex>, %[[ARG1:.*]]: index)
+func.func @extract_load_index(%arg0: memref<?xindex>, %arg1: index) -> index {
+// CHECK:   %[[RES:.*]] = memref.load %[[ARG0]][%[[ARG1]]] : memref<?xindex>
+// CHECK:   return %[[RES]] : index
+  %0 = vector.load %arg0[%arg1] : memref<?xindex>, vector<4xindex>
+  %1 = vector.extract %0[0] : index from vector<4xindex>
+  return %1 : index
+}
+
 // CHECK-LABEL: @extract_load_scalar_non_zero_off
 //  CHECK-SAME:   (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index)
 func.func @extract_load_scalar_non_zero_off(%arg0: memref<?xf32>, %arg1: index) -> f32 {
@@ -598,6 +608,18 @@ func.func @negative_extract_load_scalar_from_memref_of_vec(%arg0: memref<?xvecto
   return %1 : f32
 }
 
+// CHECK-LABEL: @negative_extract_load_scalar_from_memref_of_i1
+//  CHECK-SAME:   (%[[ARG0:.*]]: memref<?xi1>, %[[ARG1:.*]]: index)
+func.func @negative_extract_load_scalar_from_memref_of_i1(%arg0: memref<?xi1>, %arg1: index) -> i1 {
+// Subbyte types are tricky, ignore them for now.
+// CHECK:   %[[RES:.*]] = vector.load %[[ARG0]][%[[ARG1]]] : memref<?xi1>, vector<8xi1>
+// CHECK:   %[[EXT:.*]] = vector.extract %[[RES]][0] : i1 from vector<8xi1>
+// CHECK:   return %[[EXT]] : i1
+  %0 = vector.load %arg0[%arg1] : memref<?xi1>, vector<8xi1>
+  %1 = vector.extract %0[0] : i1 from vector<8xi1>
+  return %1 : i1
+}
+
 // CHECK-LABEL: @negative_extract_load_no_single_use
 //  CHECK-SAME:   (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index)
 func.func @negative_extract_load_no_single_use(%arg0: memref<?xf32>, %arg1: index) -> (f32, vector<4xf32>) {