[mlir][llvm] add basic type consistency pattern destructuring integer stores

[zero9178]

Summary

This is a common pattern produced by clang and similar.
Essentially, it coalesces stores into adjacent integer fields into a single integer store.
This violates our definition of type-consistency that the pass is supposed to enforce and also prevents SROA and mem2reg from eliminating allocas.

This patch fixes that by splitting these stores into multiple stores.
It does so by simply using logical shift rights and truncating the produced value to the size of the field, optionally bitcasting before storing into the field.

The implementation is currently very simple, only working on struct types of a single depth and adjacent fields in that struct, with no padding inbetween. Future work could improve on these once required.

[gysit]

Thanks!

The logic looks good. I have added a few nit comments.

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

314–316

Would it make sense to sort of invert the logic of this function and return std::optional<Type>?

In any case it would be great to add a doc string for all static functions and I would also try to start the function names with a verb. E.g.,

static std::optional<Type> getConsistentType(GEPOp gep);

359

Any chance we could work with ArrayRef<Type> and use drop front / tail to remove the parts that do not map to the store?

393

ultra nit: braces would be needed due to the comment. Same right below.

459

nit: t -> type or at least ty. I would try to avoid single letter variable names.

468

nit: let's use shrOp and gepOp (below) for operations.

470

nit: fieldIntType?

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

158

nit: Let's use CST0 and CST32?

229

Would it make sense to add a test with a packed struct of the form {i16, i32, i16}. In this case, we should be able to split the store?

[zero9178]

Address review comments:

• Adjust naming nits in C++ and FileCheck
• Add test for packed struct
• Rewrite code getting written to fields to work nicely with ArrayRef
• Other nits

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

229

Correct, since no padding would be written to. I added a testcase showcasing it.

[gysit]

Nice improvement!

LGTM modulo some ultra nits.

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

314–315

358

ultra nit: Can you use three slashes as well and ideally also back ticks ` instead of ' for the variable names.

368

ultra nit: I would drop the newlines at the beginning of the function

403

ultra nit: unsigned int -> unsigned

[zero9178]

address review comments

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> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  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

 }
 
 //===----------------------------------------------------------------------===//
 // CanonicalizeAlignedGep
 //===----------------------------------------------------------------------===//
 
 /// Returns the amount of bytes the provided GEP elements will offset the
 /// pointer by. Returns nullopt if the offset could not be computed.
-static std::optional<uint64_t> gepToByteOffset(DataLayout &layout, Type base,
-                                               ArrayRef<uint32_t> indices) {
-  uint64_t offset = indices[0] * layout.getTypeSize(base);
+static std::optional<uint64_t> gepToByteOffset(DataLayout &layout, GEPOp gep) {
 
-  Type currentType = base;
+  SmallVector<uint32_t> indices;
+  // Ensures all indices are static and fetches them.
+  for (auto index : gep.getIndices()) {
+    IntegerAttr indexInt = llvm::dyn_cast_if_present<IntegerAttr>(index);
+    if (!indexInt)
+      return std::nullopt;
+    indices.push_back(indexInt.getInt());
+  }
+
+  uint64_t offset = indices[0] * layout.getTypeSize(gep.getSourceElementType());
+
+  Type currentType = gep.getSourceElementType();
   for (uint32_t index : llvm::drop_begin(indices)) {
     bool shouldCancel =
         TypeSwitch<Type, bool>(currentType)
             .Case([&](LLVMArrayType arrayType) {
               if (arrayType.getNumElements() <= index)
                 return true;
               offset += index * layout.getTypeSize(arrayType.getElementType());
               currentType = arrayType.getElementType();
while (distanceToStart < offset) {
 
     if (shouldCancel)
       return failure();
   }
 
   return success();
 }
 
-LogicalResult
-CanonicalizeAlignedGep::matchAndRewrite(GEPOp gep,
-                                        PatternRewriter &rewriter) const {
+// Returns the consistent type for the GEP if the GEP is not type-consistent.
+// Returns failure if the GEP is already consistent.
+static FailureOr<Type> getRequiredConsistentGEPType(GEPOp gep) {
   // GEP of typed pointers are not supported.
   if (!gep.getElemType())
     return failure();
 
   std::optional<Type> maybeBaseType = gep.getElemType();
   if (!maybeBaseType)
     return failure();
   Type baseType = *maybeBaseType;
 
   Type typeHint = isElementTypeInconsistent(gep.getBase(), baseType);
   if (!typeHint)
     return failure();
+  return typeHint;
+}
 
-  SmallVector<uint32_t> indices;
-  // Ensures all indices are static and fetches them.
-  for (auto index : gep.getIndices()) {
-    IntegerAttr indexInt = llvm::dyn_cast_if_present<IntegerAttr>(index);
-    if (!indexInt)
-      return failure();
-    indices.push_back(indexInt.getInt());
+LogicalResult
+CanonicalizeAlignedGep::matchAndRewrite(GEPOp gep,
+                                        PatternRewriter &rewriter) const {
+  FailureOr<Type> typeHint = getRequiredConsistentGEPType(gep);
+  if (failed(typeHint)) {
+    // GEP is already canonical, nothing to do here.
   }
 
   DataLayout layout = DataLayout::closest(gep);
-  std::optional<uint64_t> desiredOffset =
-      gepToByteOffset(layout, gep.getSourceElementType(), indices);
+  std::optional<uint64_t> desiredOffset = gepToByteOffset(layout, gep);
   if (!desiredOffset)
     return failure();
 
   SmallVector<GEPArg> newIndices;
   if (failed(
-          findIndicesForOffset(layout, typeHint, *desiredOffset, newIndices)))
+          findIndicesForOffset(layout, *typeHint, *desiredOffset, newIndices)))
     return failure();
 
   rewriter.replaceOpWithNewOp<GEPOp>(
-      gep, LLVM::LLVMPointerType::get(getContext()), typeHint, gep.getBase(),
+      gep, LLVM::LLVMPointerType::get(getContext()), *typeHint, gep.getBase(),
       newIndices, gep.getInbounds());
 
   return success();
 }
 
+// Returns the list of fields of 'structType' that are written to by a store
+// 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) {
+
+  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 int fieldSize = dataLayout.getTypeSize(body[exclusiveEnd]);
+    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)
+    return failure();
+  return body.take_front(exclusiveEnd);
+}
+
+LogicalResult
+SplitIntegerStores::matchAndRewrite(StoreOp store,
+                                    PatternRewriter &rewriter) const {
+  IntegerType sourceType = dyn_cast<IntegerType>(store.getValue().getType());
+  if (!sourceType) {
+    // We currently only support integer 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 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.
+    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);
+  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);
+    }
+
+    // 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;
+  }
+
+  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());
     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 @index_to_struct(%arg: i32) {
   %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, struct<"bar", (i32, i32)>)> : (i32) -> !llvm.ptr
   // CHECK: %[[GEP0:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, struct<"bar", (i32, i32)>)>
   // CHECK: %[[GEP1:.*]] = llvm.getelementptr %[[GEP0]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"bar", (i32, i32)>
   %7 = llvm.getelementptr %1[4] : (!llvm.ptr) -> !llvm.ptr, i8
   // CHECK: llvm.store %[[ARG]], %[[GEP1]]
   llvm.store %arg, %7 : i32, !llvm.ptr
   llvm.return
 }
+
+// -----
+
+// CHECK-LABEL: llvm.func @coalesced_store_ints
+// CHECK-SAME: %[[ARG:.*]]: i64
+llvm.func @coalesced_store_ints(%arg: i64) {
+  // CHECK: %[[CST0:.*]] = llvm.mlir.constant(0 : i64) : i64
+  // CHECK: %[[CST32:.*]] = llvm.mlir.constant(32 : i64) : i64
+
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32)> : (i32) -> !llvm.ptr
+
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32)>
+  // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST0]]
+  // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i32
+  // CHECK: llvm.store %[[TRUNC]], %[[GEP]]
+  // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST32]] : i64
+  // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i32
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<"foo", (i32, i32)>
+  // CHECK: llvm.store %[[TRUNC]], %[[GEP]]
+  llvm.store %arg, %1 : i64, !llvm.ptr
+  // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @coalesced_store_ints_offset
+// CHECK-SAME: %[[ARG:.*]]: i64
+llvm.func @coalesced_store_ints_offset(%arg: i64) {
+  // CHECK: %[[CST0:.*]] = llvm.mlir.constant(0 : i64) : i64
+  // CHECK: %[[CST32:.*]] = llvm.mlir.constant(32 : i64) : i64
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i64, i32, i32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i64, i32, i32)> : (i32) -> !llvm.ptr
+  %3 = llvm.getelementptr %1[0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i64, i32, i32)>
+
+  // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST0]]
+  // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i32
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i64, i32, i32)>
+  // CHECK: llvm.store %[[TRUNC]], %[[GEP]]
+  // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST32]] : i64
+  // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i32
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 2] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<"foo", (i64, i32, i32)>
+  // CHECK: llvm.store %[[TRUNC]], %[[GEP]]
+  llvm.store %arg, %3 : i64, !llvm.ptr
+  // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @coalesced_store_floats
+// CHECK-SAME: %[[ARG:.*]]: i64
+llvm.func @coalesced_store_floats(%arg: i64) {
+  // CHECK: %[[CST0:.*]] = llvm.mlir.constant(0 : i64) : i64
+  // CHECK: %[[CST32:.*]] = llvm.mlir.constant(32 : i64) : i64
+  %0 = llvm.mlir.constant(1 : i32) : i32
+
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (f32, f32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (f32, f32)> : (i32) -> !llvm.ptr
+
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (f32, f32)>
+  // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST0]]
+  // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i32
+  // CHECK: %[[BIT_CAST:.*]] = llvm.bitcast %[[TRUNC]] : i32 to f32
+  // CHECK: llvm.store %[[BIT_CAST]], %[[GEP]]
+  // CHECK: %[[SHR:.*]] = llvm.lshr %[[ARG]], %[[CST32]] : i64
+  // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i32
+  // CHECK: %[[BIT_CAST:.*]] = llvm.bitcast %[[TRUNC]] : i32 to f32
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1] : (!llvm.ptr)  -> !llvm.ptr, !llvm.struct<"foo", (f32, f32)>
+  // CHECK: llvm.store %[[BIT_CAST]], %[[GEP]]
+  llvm.store %arg, %1 : i64, !llvm.ptr
+  // CHECK-NOT: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.return
+}
+
+// -----
+
+// Padding test purposefully not modified.
+
+// CHECK-LABEL: llvm.func @coalesced_store_padding_inbetween
+// CHECK-SAME: %[[ARG:.*]]: i64
+llvm.func @coalesced_store_padding_inbetween(%arg: i64) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i16, i32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i16, i32)> : (i32) -> !llvm.ptr
+  // CHECK: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.store %arg, %1 : i64, !llvm.ptr
+  llvm.return
+}
+
+// -----
+
+// Padding test purposefully not modified.
+
+// CHECK-LABEL: llvm.func @coalesced_store_padding_end
+// CHECK-SAME: %[[ARG:.*]]: i64
+llvm.func @coalesced_store_padding_end(%arg: i64) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i16)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i16)> : (i32) -> !llvm.ptr
+  // CHECK: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.store %arg, %1 : i64, !llvm.ptr
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @coalesced_store_past_end
+// CHECK-SAME: %[[ARG:.*]]: i64
+llvm.func @coalesced_store_past_end(%arg: i64) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32)> : (i32) -> !llvm.ptr
+  // CHECK: llvm.store %[[ARG]], %[[ALLOCA]]
+  llvm.store %arg, %1 : i64, !llvm.ptr
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @coalesced_store_packed_struct
+// CHECK-SAME: %[[ARG:.*]]: i64
+llvm.func @coalesced_store_packed_struct(%arg: i64) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[CST0:.*]] = llvm.mlir.constant(0 : i64) : i64
+  // CHECK: %[[CST16:.*]] = llvm.mlir.constant(16 : i64) : i64
+  // CHECK: %[[CST48:.*]] = llvm.mlir.constant(48 : i64) : i64
+
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", packed (i16, i32, i16)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", packed (i16, i32, i16)> : (i32) -> !llvm.ptr
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", packed (i16, i32, 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]]
+  // CHECK: %[[TRUNC:.*]] = llvm.trunc %[[SHR]] : i64 to i32
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", packed (i16, i32, i16)>
+  // CHECK: llvm.store %[[TRUNC]], %[[GEP]]
+  // 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
+}