[mlir][LLVM] Make `SplitStores` pattern capable of writing to sub-aggregates

[zero9178]

Summary

The pattern was previously only capable of storing into struct fields which are primitive types. If the struct contained a nested struct it immediately aborted the pattern rewrite.

This patch introduces the capability of recursively splitting stores into sub-aggregates as well. This is achieved by splitting an aggregate sized integer from the original store argument and letting repeated pattern applications further split it into field stores.

Additionally, the pattern is also capable of handling partial writes into aggregates, which is a pattern clang generates as well. Special care had to be taken to make sure no stores are created that weren't in the original code.

[gysit]

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

395

I guess there is a slight risk here that we may not be able to split the store into the nested array/struct?

508

[zero9178]

address review comments

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

395

Yes. This transformation is generally optimistic in that it always assumes that integer stores into structs it generates can be further split up. Checking otherwise would require doing a recursive check of the fields/elements of the aggregate.
That said, I think its a rather unlikely case that an integer store would create consistent stores into some fields followed by an inconsistent store into a subaggregate. And even if, the generated code stays correct.

[zero9178]

add explicit check making sure that when splitting an integer only partially storing into an aggregate that it can also be split afterwards

[gysit]

Thanks LGTM modulo ultra nit.

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

552

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

 /// operation writing `storeSize` bytes at `storeOffset` within the struct.
 /// `storeOffset` is required to cleanly point to an immediate field within
 /// the struct.
 /// If the write operation were to write to any padding, write beyond the
 /// struct, partially write to a field, or contains currently unsupported
 /// types, failure is returned.
 static FailureOr<ArrayRef<Type>>
 getWrittenToFields(const DataLayout &dataLayout, LLVMStructType structType,
-                   int storeSize, unsigned storeOffset) {
+                   unsigned storeSize, unsigned storeOffset) {
   ArrayRef<Type> body = structType.getBody();
   unsigned currentOffset = 0;
   body = body.drop_until([&](Type type) {
     if (!structType.isPacked()) {
       unsigned alignment = dataLayout.getTypeABIAlignment(type);
       currentOffset = llvm::alignTo(currentOffset, alignment);
     }
 
     // currentOffset is guaranteed to be equal to offset since offset is either
     // 0 or stems from a type-consistent GEP indexing into just a single
     // aggregate.
     if (currentOffset == storeOffset)
       return true;
 
     assert(currentOffset < storeOffset &&
            "storeOffset should cleanly point into an immediate field");
 
     currentOffset += dataLayout.getTypeSize(type);
     return false;
   });
 
   size_t exclusiveEnd = 0;
   for (; exclusiveEnd < body.size() && storeSize > 0; exclusiveEnd++) {
-    // Not yet recursively handling aggregates, only primitives.
-    if (!isa<IntegerType, FloatType>(body[exclusiveEnd]))
-      return failure();
-
     if (!structType.isPacked()) {
       unsigned alignment = dataLayout.getTypeABIAlignment(body[exclusiveEnd]);
       // No padding allowed inbetween fields at this point in time.
       if (!llvm::isAligned(llvm::Align(alignment), currentOffset))
         return failure();
     }
 
     unsigned fieldSize = dataLayout.getTypeSize(body[exclusiveEnd]);
+    if (fieldSize > storeSize) {
+      // Partial writes into an aggregate are okay since subsequent pattern
+      // applications can further split these up into writes into the
+      // sub-elements.
+      if (isa<LLVMStructType, LLVMArrayType>(body[exclusiveEnd]))
+        return body.take_front(exclusiveEnd + 1);
+
+      return failure();
+    }
     currentOffset += fieldSize;
     storeSize -= fieldSize;
   }
 
-  // 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)
+  // If the storeSize is not 0 at this point we are  writing past the aggregate
+  // as a whole. Abort.
+  if (storeSize > 0)
     return failure();
   return body.take_front(exclusiveEnd);
 }
 
 /// 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,
static void splitVectorStore(const DataLayout &dataLayout, Location loc,
   }
 }
 
 /// 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,
+                              Value value, unsigned storeSize,
+                              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);
+    // If we are doing a partial write into a direct field the remaining
+    // `storeSize` will be less than the size of the field. We have to truncate
+    // to the `storeSize` to avoid creating a store that wasn't in the original
+    // code.
+    IntegerType fieldIntType =
+        rewriter.getIntegerType(std::min(fieldSize, storeSize) * 8);
     Value valueToStore = rewriter.create<TruncOp>(loc, fieldIntType, shrOp);
 
     // 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;
+    storeSize -= fieldSize;
   }
 }
 
 LogicalResult SplitStores::matchAndRewrite(StoreOp store,
                                            PatternRewriter &rewriter) const {
   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();
   }
 
   auto dataLayout = DataLayout::closest(store);
 
+  unsigned storeSize = dataLayout.getTypeSize(sourceType);
   unsigned offset = 0;
   Value address = store.getAddr();
   if (auto gepOp = address.getDefiningOp<GEPOp>()) {
     // Currently only handle canonical GEPs with exactly two indices,
     // indexing a single aggregate deep.
-    // Recursing into sub-structs is left as a future exercise.
     // If the GEP is not canonical we have to fail, otherwise we would not
     // create type-consistent IR.
     if (gepOp.getIndices().size() != 2 ||
         succeeded(getRequiredConsistentGEPType(gepOp)))
       return failure();
 
-    // A GEP might point somewhere into the middle of an aggregate with the
-    // store storing into multiple adjacent elements. Destructure into
-    // the base address with an offset.
+    // If the element's size the GEP is pointing at is smaller than the store
+    // size, it is pointing into the middle of an aggregate with the store
+    // storing into multiple adjacent elements. Destructure into the base
+    // address of the aggregate with a store offset.
+    if (storeSize > dataLayout.getTypeSize(gepOp.getResultPtrElementType())) {
       std::optional<uint64_t> byteOffset = gepToByteOffset(dataLayout, gepOp);
       if (!byteOffset)
         return failure();
 
       offset = *byteOffset;
       typeHint = gepOp.getSourceElementType();
       address = gepOp.getBase();
     }
+  }
 
   auto structType = typeHint.dyn_cast<LLVMStructType>();
   if (!structType) {
     // TODO: Handle array types in the future.
     return failure();
   }
 
   FailureOr<ArrayRef<Type>> writtenToFields =
-      getWrittenToFields(dataLayout, structType,
-                         /*storeSize=*/dataLayout.getTypeSize(sourceType),
-                         /*storeOffset=*/offset);
+      getWrittenToFields(dataLayout, structType, storeSize, offset);
   if (failed(writtenToFields))
     return failure();
 
   if (writtenToFields->size() <= 1) {
     // Other patterns should take care of this case, we are only interested in
     // splitting field stores.
     return failure();
   }
 
   if (isa<IntegerType>(sourceType)) {
     splitIntegerStore(dataLayout, store.getLoc(), rewriter, address,
-                      store.getValue(), offset, *writtenToFields);
+                      store.getValue(), storeSize, offset, *writtenToFields);
     rewriter.eraseOp(store);
     return success();
   }
 
   // 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();

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

llvm.func @bitcast_insertion(%arg: i32) {
   // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x f32
   %1 = llvm.alloca %0 x f32 : (i32) -> !llvm.ptr
   // CHECK: %[[BIT_CAST:.*]] = llvm.bitcast %[[ARG]] : i32 to f32
   // CHECK: llvm.store %[[BIT_CAST]], %[[ALLOCA]]
   llvm.store %arg, %1 : i32, !llvm.ptr
   // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
   llvm.return
 }
+
+// -----
+
+// CHECK-LABEL: llvm.func @overlapping_int_aggregate_store
+// CHECK-SAME: %[[ARG:.*]]: i64
+llvm.func @overlapping_int_aggregate_store(%arg: i64) {
+  // CHECK: %[[CST0:.*]] = llvm.mlir.constant(0 : i64) : i64
+  // CHECK: %[[CST16:.*]] = llvm.mlir.constant(16 : i64) : i64
+
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i16, struct<(i16, i16, i16)>)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i16, struct<(i16, i16, i16)>)> : (i32) -> !llvm.ptr
+
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i16, struct<(i16, i16, i16)>)>
+  // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST0]]
+  // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i16
+  // CHECK: llvm.store %[[TRUNC]], %[[GEP]]
+
+  // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST16]] : i64
+  // CHECK: [[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i48
+  // CHECK: %[[TOP_GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<"foo", (i16, struct<(i16, i16, i16)>)>
+
+  // Normal integer splitting of [[TRUNC]] follows:
+
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[TOP_GEP]][0, 0] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<(i16, i16, i16)>
+  // CHECK: llvm.store %{{.*}}, %[[GEP]]
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[TOP_GEP]][0, 1] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<(i16, i16, i16)>
+  // CHECK: llvm.store %{{.*}}, %[[GEP]]
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[TOP_GEP]][0, 2] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<(i16, i16, i16)>
+  // CHECK: llvm.store %{{.*}}, %[[GEP]]
+
+  llvm.store %arg, %1 : i64, !llvm.ptr
+  // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @overlapping_vector_aggregate_store
+// CHECK-SAME: %[[ARG:.*]]: vector<4xi16>
+llvm.func @overlapping_vector_aggregate_store(%arg: vector<4 x i16>) {
+  // 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", (i16, struct<(i16, i16, i16)>)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i16, struct<(i16, i16, i16)>)> : (i32) -> !llvm.ptr
+
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<"foo", (i16, struct<(i16, i16, i16)>)>
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST0]] : i32]
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP]]
+
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST1]] : i32]
+  // CHECK: %[[GEP0:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<"foo", (i16, struct<(i16, i16, i16)>)>
+  // CHECK: %[[GEP1:.*]] = llvm.getelementptr %[[GEP0]][0, 0] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<(i16, i16, i16)>
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP1]]
+
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST2]] : i32]
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1, 1] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<"foo", (i16, struct<(i16, i16, i16)>)>
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP]]
+
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractelement %[[ARG]][%[[CST3]] : i32]
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1, 2] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<"foo", (i16, struct<(i16, i16, i16)>)>
+  // CHECK: llvm.store %[[EXTRACT]], %[[GEP]]
+
+  llvm.store %arg, %1 : vector<4 x i16>, !llvm.ptr
+  // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @partially_overlapping_aggregate_store
+// CHECK-SAME: %[[ARG:.*]]: i64
+llvm.func @partially_overlapping_aggregate_store(%arg: i64) {
+  // CHECK: %[[CST0:.*]] = llvm.mlir.constant(0 : i64) : i64
+  // CHECK: %[[CST16:.*]] = llvm.mlir.constant(16 : i64) : i64
+
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i16, struct<(i16, i16, i16, i16)>)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i16, struct<(i16, i16, i16, i16)>)> : (i32) -> !llvm.ptr
+
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i16, struct<(i16, i16, i16, i16)>)>
+  // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST0]]
+  // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i16
+  // CHECK: llvm.store %[[TRUNC]], %[[GEP]]
+
+  // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST16]] : i64
+  // CHECK: [[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i48
+  // CHECK: %[[TOP_GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i16, struct<(i16, i16, i16, i16)>)>
+
+  // Normal integer splitting of [[TRUNC]] follows:
+
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[TOP_GEP]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(i16, i16, i16, i16)>
+  // CHECK: llvm.store %{{.*}}, %[[GEP]]
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[TOP_GEP]][0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(i16, i16, i16, i16)>
+  // CHECK: llvm.store %{{.*}}, %[[GEP]]
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[TOP_GEP]][0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(i16, i16, i16, i16)>
+  // CHECK: llvm.store %{{.*}}, %[[GEP]]
+
+  // It is important that there are no more stores at this point.
+  // Specifically a store into the fourth field of %[[TOP_GEP]] would
+  // incorrectly change the semantics of the code.
+  // CHECK-NOT: llvm.store %{{.*}}, %{{.*}}
+
+  llvm.store %arg, %1 : i64, !llvm.ptr
+
+  llvm.return
+}