[mlir][LLVM] Make `SplitIntegerStores` capable of splitting vectors as well

[zero9178]

Summary

The original plan was to turn this into its own pattern, but one of the difficulties was deeming when splitting the vector is required.
SplitIntegerStores essentially already did that by checking for field overlap.
Therefore, it was renamed to SplitStores and extended to splitting stores with values of vector and integer type.

The vector splitting is done in a simple manner by simply using extractelement to get each vector element. Subsequent pattern applications are responsible for further cleaning up the output and making it type-consistent.

Worst case, if the code cannot be transformed into a type-consistent form (due to e.g. the code explicitly doing partial writes to elements or similar), we might needlessly do a vector split.

[gysit]

mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.td

44

nit: should we rename to max-vector-split-size since this may be used for loads as well at some point.

mlir/include/mlir/Dialect/LLVMIR/Transforms/TypeConsistency.h

64–65

ultra nit: unsigned int -> unsigned.

mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp

544

Would it make sense to put the handling of the vector and and the struct case below into separate function to make this function a little smaller?

[zero9178]

• Address review comments
• Add test and capability for vector store to also work on an offset

[gysit]

Nice LGTM!

[zero9178]

clang-format

mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.td

def LLVMTypeConsistency
     : Pass<"llvm-type-consistency", "::mlir::LLVM::LLVMFuncOp"> {
   let summary = "Rewrites to improve type consistency";
   let description = [{
     Set of rewrites to improve the coherency of types within an LLVM dialect
     program. This will adjust operations operating on pointers so they interpret
     their associated pointee type as consistently as possible.
   }];
   let constructor = "::mlir::LLVM::createTypeConsistencyPass()";
+
+  let options = [
+    Option<"maxVectorSplitSize", "max-vector-split-size", "unsigned",
+           /*default=*/"512",
+           "Maximum size in bits of a vector value in a load or store operation"
+           " operating on multiple elements that should still be split">,
+  ];
 }
 
 def NVVMOptimizeForTarget : Pass<"llvm-optimize-for-nvvm-target"> {
   let summary = "Optimize NVVM IR";
   let constructor = "::mlir::NVVM::createOptimizeForTargetPass()";
 }
 
 def DIScopeForLLVMFuncOp : Pass<"ensure-debug-info-scope-on-llvm-func", "LLVM::LLVMFuncOp"> {

mlir/include/mlir/Dialect/LLVMIR/Transforms/TypeConsistency.h

 class CanonicalizeAlignedGep : public OpRewritePattern<GEPOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(GEPOp gep,
                                 PatternRewriter &rewriter) const override;
 };
 
-/// Splits stores of integers which write into multiple adjacent stores
-/// of a pointer. The integer is then split and stores are generated for
-/// every field being stored in a type-consistent manner.
-/// This is currently done on a best-effort basis.
-class SplitIntegerStores : public OpRewritePattern<StoreOp> {
+/// Splits stores which write into multiple adjacent elements of an aggregate
+/// through a pointer. Currently, integers and vector are split and stores
+/// are generated for every element being stored to in a type-consistent manner.
+/// This is done on a best-effort basis.
+class SplitStores : public OpRewritePattern<StoreOp> {
+  unsigned maxVectorSplitSize;
+
 public:
-  using OpRewritePattern::OpRewritePattern;
+  SplitStores(MLIRContext *context, unsigned maxVectorSplitSize)
+      : OpRewritePattern(context), maxVectorSplitSize(maxVectorSplitSize) {}
 
   LogicalResult matchAndRewrite(StoreOp store,
                                 PatternRewriter &rewrite) const override;
 };
 
 } // namespace LLVM
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LLVMIR_TRANSFORMS_TYPECONSISTENCY_H

mlir/lib/Dialect/LLVMIR/Transforms/TypeConsistency.cpp

getWrittenToFields(const DataLayout &dataLayout, LLVMStructType structType,
 
   // If the storeSize is not 0 at this point we are either partially writing
   // into a field or writing past the aggregate as a whole. Abort.
   if (storeSize != 0)
     return failure();
   return body.take_front(exclusiveEnd);
 }
 
-LogicalResult
-SplitIntegerStores::matchAndRewrite(StoreOp store,
+/// Splits a store of the vector `value` into `address` at `storeOffset` into
+/// multiple stores of each element with the goal of each generated store
+/// becoming type-consistent through subsequent pattern applications.
+static void splitVectorStore(const DataLayout &dataLayout, Location loc,
+                             RewriterBase &rewriter, Value address,
+                             TypedValue<VectorType> value,
+                             unsigned storeOffset) {
+  VectorType vectorType = value.getType();
+  unsigned elementSize = dataLayout.getTypeSize(vectorType.getElementType());
+
+  // Extract every element in the vector and store it in the given address.
+  for (size_t index : llvm::seq<size_t>(0, vectorType.getNumElements())) {
+    auto pos =
+        rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(index));
+    auto extractOp = rewriter.create<ExtractElementOp>(loc, value, pos);
+
+    // For convenience, we do indexing by calculating the final byte offset.
+    // Other patterns will turn this into a type-consistent GEP.
+    auto gepOp = rewriter.create<GEPOp>(
+        loc, address.getType(), rewriter.getI8Type(), address,
+        ArrayRef<GEPArg>{storeOffset + index * elementSize});
+
+    rewriter.create<StoreOp>(loc, extractOp, gepOp);
+  }
+}
+
+/// Splits a store of the integer `value` into `address` at `storeOffset` into
+/// multiple stores to each 'writtenFields', making each store operation
+/// type-consistent.
+static void splitIntegerStore(const DataLayout &dataLayout, Location loc,
+                              RewriterBase &rewriter, Value address,
+                              Value value, unsigned storeOffset,
+                              ArrayRef<Type> writtenToFields) {
+  unsigned currentOffset = storeOffset;
+  for (Type type : writtenToFields) {
+    unsigned fieldSize = dataLayout.getTypeSize(type);
+
+    // Extract the data out of the integer by first shifting right and then
+    // truncating it.
+    auto pos = rewriter.create<ConstantOp>(
+        loc, rewriter.getIntegerAttr(value.getType(),
+                                     (currentOffset - storeOffset) * 8));
+
+    auto shrOp = rewriter.create<LShrOp>(loc, value, pos);
+
+    IntegerType fieldIntType = rewriter.getIntegerType(fieldSize * 8);
+    Value valueToStore = rewriter.create<TruncOp>(loc, fieldIntType, shrOp);
+    if (fieldIntType != type) {
+      // Bitcast to the right type. `fieldIntType` was explicitly created
+      // to be of the same size as `type` and must currently be a primitive as
+      // well.
+      valueToStore = rewriter.create<BitcastOp>(loc, type, valueToStore);
+    }
+
+    // We create an `i8` indexed GEP here as that is the easiest (offset is
+    // already known). Other patterns turn this into a type-consistent GEP.
+    auto gepOp =
+        rewriter.create<GEPOp>(loc, address.getType(), rewriter.getI8Type(),
+                               address, ArrayRef<GEPArg>{currentOffset});
+    rewriter.create<StoreOp>(loc, valueToStore, gepOp);
+
+    // No need to care about padding here since we already checked previously
+    // that no padding exists in this range.
+    currentOffset += fieldSize;
+  }
+}
+
+LogicalResult SplitStores::matchAndRewrite(StoreOp store,
                                            PatternRewriter &rewriter) const {
-  IntegerType sourceType = dyn_cast<IntegerType>(store.getValue().getType());
-  if (!sourceType) {
-    // We currently only support integer sources.
+  Type sourceType = store.getValue().getType();
+  if (!isa<IntegerType, VectorType>(sourceType)) {
+    // We currently only support integer and vector sources.
     return failure();
   }
 
   Type typeHint = isElementTypeInconsistent(store.getAddr(), sourceType);
   if (!typeHint) {
     // Nothing to do, since it is already consistent.
     return failure();
   }
LogicalResult SplitStores::matchAndRewrite(StoreOp store,
 
   FailureOr<ArrayRef<Type>> writtenToFields =
       getWrittenToFields(dataLayout, structType,
                          /*storeSize=*/dataLayout.getTypeSize(sourceType),
                          /*storeOffset=*/offset);
   if (failed(writtenToFields))
     return failure();
 
-  unsigned currentOffset = offset;
-  for (Type type : *writtenToFields) {
-    unsigned fieldSize = dataLayout.getTypeSize(type);
-
-    // Extract the data out of the integer by first shifting right and then
-    // truncating it.
-    auto pos = rewriter.create<ConstantOp>(
-        store.getLoc(),
-        rewriter.getIntegerAttr(sourceType, (currentOffset - offset) * 8));
-
-    auto shrOp = rewriter.create<LShrOp>(store.getLoc(), store.getValue(), pos);
-
-    IntegerType fieldIntType = rewriter.getIntegerType(fieldSize * 8);
-    Value valueToStore =
-        rewriter.create<TruncOp>(store.getLoc(), fieldIntType, shrOp);
-    if (fieldIntType != type) {
-      // Bitcast to the right type. `fieldIntType` was explicitly created
-      // to be of the same size as `type` and must currently be a primitive as
-      // well.
-      valueToStore =
-          rewriter.create<BitcastOp>(store.getLoc(), type, valueToStore);
-    }
+  if (writtenToFields->size() <= 1) {
+    // Other patterns should take care of this case, we are only interested in
+    // splitting field stores.
+    return failure();
 
-    // We create an `i8` indexed GEP here as that is the easiest (offset is
-    // already known). Other patterns turn this into a type-consistent GEP.
-    auto gepOp = rewriter.create<GEPOp>(store.getLoc(), address.getType(),
-                                        rewriter.getI8Type(), address,
-                                        ArrayRef<GEPArg>{currentOffset});
-    rewriter.create<StoreOp>(store.getLoc(), valueToStore, gepOp);
-
-    // No need to care about padding here since we already checked previously
-    // that no padding exists in this range.
-    currentOffset += fieldSize;
+  if (isa<IntegerType>(sourceType)) {
+    splitIntegerStore(dataLayout, store.getLoc(), rewriter, address,
+                      store.getValue(), offset, *writtenToFields);
+    rewriter.eraseOp(store);
+    return success();
   }
 
-  rewriter.eraseOp(store);
+  // Add a reasonable bound to not split very large vectors that would end up
+  // generating lots of code.
+  if (dataLayout.getTypeSizeInBits(sourceType) > maxVectorSplitSize)
+    return failure();
 
+  // Vector types are simply split into its elements and new stores generated
+  // with those. Subsequent pattern applications will split these stores further
+  // if required.
+  splitVectorStore(dataLayout, store.getLoc(), rewriter, address,
+                   cast<TypedValue<VectorType>>(store.getValue()), offset);
+  rewriter.eraseOp(store);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Type consistency pass
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct LLVMTypeConsistencyPass
     : public LLVM::impl::LLVMTypeConsistencyBase<LLVMTypeConsistencyPass> {
   void runOnOperation() override {
     RewritePatternSet rewritePatterns(&getContext());
     rewritePatterns.add<AddFieldGetterToStructDirectUse<LoadOp>>(&getContext());
     rewritePatterns.add<AddFieldGetterToStructDirectUse<StoreOp>>(
         &getContext());
     rewritePatterns.add<CanonicalizeAlignedGep>(&getContext());
-    rewritePatterns.add<SplitIntegerStores>(&getContext());
+    rewritePatterns.add<SplitStores>(&getContext(), maxVectorSplitSize);
     FrozenRewritePatternSet frozen(std::move(rewritePatterns));
 
     if (failed(applyPatternsAndFoldGreedily(getOperation(), frozen)))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> LLVM::createTypeConsistencyPass() {
   return std::make_unique<LLVMTypeConsistencyPass>();
 }

mlir/test/Dialect/LLVMIR/type-consistency.mlir

llvm.func @coalesced_store_packed_struct(%arg: i64) {
   // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST48]]
   // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i16
   // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", packed (i16, i32, i16)>
   // CHECK: llvm.store %[[TRUNC]], %[[GEP]]
   llvm.store %arg, %1 : i64, !llvm.ptr
   // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
   llvm.return
 }
+
+// -----
+
+// CHECK-LABEL: llvm.func @vector_write_split
+// CHECK-SAME: %[[ARG:.*]]: vector<4xi32>
+llvm.func @vector_write_split(%arg: vector<4xi32>) {
+  // CHECK: %[[CST0:.*]] = llvm.mlir.constant(0 : i32) : i32
+  // CHECK: %[[CST1:.*]] = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[CST2:.*]] = llvm.mlir.constant(2 : i32) : i32
+  // CHECK: %[[CST3:.*]] = llvm.mlir.constant(3 : i32) : i32
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i32, i32, i32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32, i32, i32)> : (i32) -> !llvm.ptr
+
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32, i32, i32)>
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST0]] : i32] : vector<4xi32>
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP]] : i32, !llvm.ptr
+
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST1]] : i32] : vector<4xi32>
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32, i32, i32)>
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP]] : i32, !llvm.ptr
+
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST2]] : i32] : vector<4xi32>
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32, i32, i32)>
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP]] : i32, !llvm.ptr
+
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST3]] : i32] : vector<4xi32>
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 3] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32, i32, i32)>
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP]] : i32, !llvm.ptr
+
+  llvm.store %arg, %1 : vector<4xi32>, !llvm.ptr
+  // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @vector_write_split_offset
+// CHECK-SAME: %[[ARG:.*]]: vector<4xi32>
+llvm.func @vector_write_split_offset(%arg: vector<4xi32>) {
+  // CHECK: %[[CST0:.*]] = llvm.mlir.constant(0 : i32) : i32
+  // CHECK: %[[CST1:.*]] = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[CST2:.*]] = llvm.mlir.constant(2 : i32) : i32
+  // CHECK: %[[CST3:.*]] = llvm.mlir.constant(3 : i32) : i32
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i64, i32, i32, i32, i32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i64, i32, i32, i32, i32)> : (i32) -> !llvm.ptr
+  %2 = llvm.getelementptr %1[0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i64, i32, i32, i32, i32)>
+
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST0]] : i32] : vector<4xi32>
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i64, i32, i32, i32, i32)>
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP]] : i32, !llvm.ptr
+
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST1]] : i32] : vector<4xi32>
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i64, i32, i32, i32, i32)>
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP]] : i32, !llvm.ptr
+
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST2]] : i32] : vector<4xi32>
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 3] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i64, i32, i32, i32, i32)>
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP]] : i32, !llvm.ptr
+
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST3]] : i32] : vector<4xi32>
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 4] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i64, i32, i32, i32, i32)>
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP]] : i32, !llvm.ptr
+
+  llvm.store %arg, %2 : vector<4xi32>, !llvm.ptr
+  // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.return
+}
+
+// -----
+
+// Small test that a split vector store will be further optimized (to than e.g.
+// split integer loads to structs as shown here)
+
+// CHECK-LABEL: llvm.func @vector_write_split_struct
+// CHECK-SAME: %[[ARG:.*]]: vector<2xi64>
+llvm.func @vector_write_split_struct(%arg: vector<2xi64>) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i32, i32, i32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32, i32, i32)> : (i32) -> !llvm.ptr
+
+  // CHECK-COUNT-4: llvm.store %{{.*}}, %{{.*}} : i32, !llvm.ptr
+
+  llvm.store %arg, %1 : vector<2xi64>, !llvm.ptr
+  // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @type_consistent_vector_store
+// CHECK-SAME: %[[ARG:.*]]: vector<4xi32>
+llvm.func @type_consistent_vector_store(%arg: vector<4xi32>) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (vector<4xi32>)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (vector<4xi32>)> : (i32) -> !llvm.ptr
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (vector<4xi32>)>
+  // CHECK: llvm.store %[[ARG]], %[[GEP]]
+  llvm.store %arg, %1 : vector<4xi32>, !llvm.ptr
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @type_consistent_vector_store_other_type
+// CHECK-SAME: %[[ARG:.*]]: vector<4xi32>
+llvm.func @type_consistent_vector_store_other_type(%arg: vector<4xi32>) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[BIT_CAST:.*]] = llvm.bitcast %[[ARG]] : vector<4xi32> to vector<4xf32>
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (vector<4xf32>)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (vector<4xf32>)> : (i32) -> !llvm.ptr
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (vector<4xf32>)>
+  // CHECK: llvm.store %[[BIT_CAST]], %[[GEP]]
+  llvm.store %arg, %1 : vector<4xi32>, !llvm.ptr
+  // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.return
+}