[mlir][llvm] Type consistency transformations

[Moxinilian]

Summary

This revision introduces new rewrites to improve the type consistency of
a program expressed in the LLVM dialect.

Type consistency means that a given opaque pointer is consistently used
assuming the same pointee type, in a best effort basis. The introduced
rewrites modify the program to improve type consistency while preserving
the same semantics. This is useful for two main reasons:

• Transformation passes in the LLVM dialect like SROA or Mem2Reg can analyse code better if type information and structure is used in a consistent manner. Opaque pointers make this difficult to enforce, but type consistency improvements increase the amount of occurences where reasonable analysis can pick up on transformable patterns.
• While LLVM IR is not particularly picky about inconsistent type uses, it may be of interest to lift LLVM IR into higher level dialects. Having more instances of consistent type information would help lifting into dialects that do care about consistent types.

In order to detect cases of inconsistent uses, operations returning an
LLVMPointer can implement the GetResultPtrElementType interface, which
allows getting a hint of which type the provided pointer should see its
pointee as, if such hint is available. The provided rewrites will then
use this hint to attempt to modify operations using the pointers so they
use the hinted type when dealing with the pointer.

Two transformations have been implemented in this revision:

• When a load/store uses a struct ptr directly to write to the first element of the struct, introduce a GEP to the first element so the type structure is preserved.
• When a GEP statically indexes in a pointer with a base type inconsistent with the hint, try to find indices using the hint as a base type that would result in the same offset, and replace the GEP with this indexing.

More transformations are possible and I hope this is only a beginning
for this simplification effort.

[Moxinilian]

Fix formatting.

[gysit]

mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td

173–174

Soon there will hopefully only be opaque pointers!

189

As this is an LLVM interface we may not need the result number. An LLVM operation can only have zero or one result. If it has multiple results they need to be packed into a struct which we don't want to support I guess.

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

33–34

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

22

nit: Can you do a pass over the includes some like sys/types.h definitely seem unnecessary.

67

nit: since there is an indexType in mlir maybe rename to subelementType?

72–75

Is this second attempt needed? It seems both structs and arrays expect a i32 index?

98

I suspect this should get the element type in the typed pointer case.

Can we bail out earlier and not handle the typed pointer case?

102

I think you should be able to do:

getAddrMutable().assign(properPtr)

117

Should't there be a check here that the first subelement type matches the inconsistent element type? Only in that case the rewrite seems beneficial?

140–142

maybe factor out a static helper method that is called areCastCompatible?

149–152

159

nit: As above getValueMutable().assign() may work.

175

nit: let's use uint64_t or size_t consistently.

242

I probably need an online explanation of this tomorrow. I would have expected that the alignment needs to relative the base address?

294

nit: I hope typed pointers will be gone :). Also I think it is fine to simplify and only support the opaque pointer case!

[Moxinilian]

Address comments + fix bugs + add bitcast in direct struct load.

[Moxinilian]

Improve revision message.

[Moxinilian]

Mark comments as done.

[gysit]

Nice Thanks!

LGTM!

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

291

ultra nit: here we end up if a scalar element in a struct or array is read at an offset? If yes it may make sense to put a short comment that says that can happen if the access to a leaf type is not perfectly aligned. Also maybe add a TODO that we may want to support vector types?

[Moxinilian]

Address comment.

mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td

InterfaceMethod<
       /*defaultImpl=*/ [{
         ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         op.setTbaaAttr(attr);
       }]
       >
   ];
 }
 
+def GetResultPtrElementType : OpInterface<"GetResultPtrElementType"> {
+  let description = [{
+    An interface for operations that yield an LLVMPointer. Allows the
+    operation to provide the type of the element an LLVMPointer opaque
+    pointer points to, if known. This is only a hint as to how to
+    interpret a given pointer, translating how the current operation
+    understands it.
+  }];
+
+  let cppNamespace = "::mlir::LLVM";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/        [{Returns the the element type hint of the n-th
+                         LLVMPointer result, if known. Returns nullptr if the
+                         requested result is not an LLVMPointer or if the
+                         element type is unknown.}],
+      /*returnType=*/  "Type",
+      /*methodName=*/  "getResultPtrElementType",
+      /*args=*/        (ins "size_t":$nth)
+      >
+  ];
+}
+
+
 //===----------------------------------------------------------------------===//
 // LLVM dialect type interfaces.
 //===----------------------------------------------------------------------===//
 
 // An interface for LLVM pointer element types.
 def LLVM_PointerElementTypeInterface
     : TypeInterface<"PointerElementTypeInterface"> {
   let cppNamespace = "::mlir::LLVM";

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 def LLVM_FDivOp : LLVM_FloatArithmeticOp<"fdiv", "FDiv">;
 def LLVM_FRemOp : LLVM_FloatArithmeticOp<"frem", "FRem">;
 def LLVM_FNegOp : LLVM_UnaryFloatArithmeticOp<
   LLVM_ScalarOrVectorOf<LLVM_AnyFloat>, "fneg", "FNeg">;
 
 // Memory-related operations.
 def LLVM_AllocaOp : LLVM_Op<"alloca",
     [DeclareOpInterfaceMethods<PromotableAllocationOpInterface>,
-     DeclareOpInterfaceMethods<DestructurableAllocationOpInterface>]>,
+     DeclareOpInterfaceMethods<DestructurableAllocationOpInterface>,
+     DeclareOpInterfaceMethods<GetResultPtrElementType>]>,
   LLVM_MemOpPatterns {
   let arguments = (ins AnyInteger:$arraySize,
                    OptionalAttr<I64Attr>:$alignment,
                    OptionalAttr<TypeAttr>:$elem_type,
                    UnitAttr:$inalloca);
   let results = (outs Res<LLVM_AnyPointer, "",
                           [MemAlloc<AutomaticAllocationScopeResource>]>:$res);
   string llvmInstName = "Alloca";
let builders = [
     ];
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def LLVM_GEPOp : LLVM_Op<"getelementptr", [Pure,
     DeclareOpInterfaceMethods<PromotableOpInterface>,
     DeclareOpInterfaceMethods<SafeMemorySlotAccessOpInterface>,
-    DeclareOpInterfaceMethods<DestructurableAccessorOpInterface>]> {
+    DeclareOpInterfaceMethods<DestructurableAccessorOpInterface>,
+    DeclareOpInterfaceMethods<GetResultPtrElementType>]> {
   let arguments = (ins LLVM_ScalarOrVectorOf<LLVM_AnyPointer>:$base,
                    Variadic<LLVM_ScalarOrVectorOf<AnyInteger>>:$dynamicIndices,
                    DenseI32ArrayAttr:$rawConstantIndices,
                    OptionalAttr<TypeAttr>:$elem_type,
                    UnitAttr:$inbounds);
   let results = (outs LLVM_ScalarOrVectorOf<LLVM_AnyPointer>:$res);
   let skipDefaultBuilders = 1;
 

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

 //===- Passes.h - LLVM Pass Construction and Registration -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LLVMIR_TRANSFORMS_PASSES_H
 #define MLIR_DIALECT_LLVMIR_TRANSFORMS_PASSES_H
 
 #include "mlir/Dialect/LLVMIR/Transforms/LegalizeForExport.h"
 #include "mlir/Dialect/LLVMIR/Transforms/OptimizeForNVVM.h"
 #include "mlir/Dialect/LLVMIR/Transforms/RequestCWrappers.h"
+#include "mlir/Dialect/LLVMIR/Transforms/TypeConsistency.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 namespace LLVM {
 
 /// Create a pass to add DIScope to LLVMFuncOp that are missing it.
 std::unique_ptr<Pass> createDIScopeForLLVMFuncOpPass();

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

let description = [{
     attribute that instructs the conversion to LLVM to emit the C wrapper for
     the function. This pass is expected to be applied immediately before the
     conversion of builtin functions to LLVM to avoid the attribute being
     dropped by other passes.
   }];
   let constructor = "::mlir::LLVM::createRequestCWrappersPass()";
 }
 
+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()";
+}
+
 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"> {
   let summary = "Materialize LLVM debug info subprogram attribute on every LLVMFuncOp";
   let description = [{

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

+//===- TypeConsistency.h - Rewrites to improve type consistency -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Set of rewrites to improve the coherency of types within an LLVM dialect
+// program. This will adjust operations around a given pointer so they interpret
+// its pointee type as consistently as possible.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LLVMIR_TRANSFORMS_TYPECONSISTENCY_H
+#define MLIR_DIALECT_LLVMIR_TRANSFORMS_TYPECONSISTENCY_H
+
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace LLVM {
+
+#define GEN_PASS_DECL_LLVMTYPECONSISTENCY
+#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc"
+
+/// Creates a pass that adjusts operations operating on pointers so they
+/// interpret pointee types as consistently as possible.
+std::unique_ptr<Pass> createTypeConsistencyPass();
+
+/// Transforms uses of pointers to a whole struct to uses of pointers to the
+/// first element of a struct. This is achieved by inserting when possible a GEP
+/// to the first element.
+template <class User>
+class AddFieldGetterToStructDirectUse : public OpRewritePattern<User> {
+public:
+  using OpRewritePattern<User>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(User user,
+                                PatternRewriter &rewriter) const override;
+};
+
+/// Canonicalizes GEPs of which the base type and the pointer's type hint do not
+/// match. This is done by replacing the original GEP into a GEP with the type
+/// hint as a base type when an element of the hinted type aligns with the
+/// original GEP.
+class CanonicalizeAlignedGep : public OpRewritePattern<GEPOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GEPOp gep,
+                                PatternRewriter &rewriter) const override;
+};
+
+} // namespace LLVM
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LLVMIR_TRANSFORMS_TYPECONSISTENCY_H

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

LogicalResult AllocaOp::verify() {
   if (auto targetExtType = dyn_cast<LLVMTargetExtType>(elemTy);
       targetExtType && !targetExtType.supportsMemOps())
     return emitOpError()
            << "this target extension type cannot be used in alloca";
 
   return success();
 }
 
+Type AllocaOp::getResultPtrElementType(size_t nth) {
+  if (nth != 0)
+    return nullptr;
+  // This will become trivial once non-opaque pointers are gone.
+  return getElemType().has_value() ? *getElemType()
+                                   : getResult().getType().getElementType();
+}
+
 //===----------------------------------------------------------------------===//
 // LLVM::BrOp
 //===----------------------------------------------------------------------===//
 
 /// Check if the `loopAttr` references correct symbols.
 static LogicalResult verifyLoopAnnotationAttr(LoopAnnotationAttr loopAttr,
                                               Operation *op) {
   if (!loopAttr)
Type LLVM::GEPOp::getSourceElementType() {
   if (std::optional<Type> elemType = getElemType())
     return *elemType;
 
   return llvm::cast<LLVMPointerType>(
              extractVectorElementType(getBase().getType()))
       .getElementType();
 }
 
+Type GEPOp::getResultPtrElementType(size_t nth) {
+  if (nth != 0)
+    return nullptr;
+
+  // Ensures all indices are static and fetches them.
+  SmallVector<IntegerAttr> indices;
+  for (auto index : getIndices()) {
+    IntegerAttr indexInt = llvm::dyn_cast_if_present<IntegerAttr>(index);
+    if (!indexInt)
+      return nullptr;
+    indices.push_back(indexInt);
+  }
+
+  // Set the initial type currently being used for indexing. This will be
+  // updated as the indices get walked over.
+  Type selectedType = getSourceElementType();
+
+  // Follow the indexed elements in the gep.
+  for (IntegerAttr index : llvm::drop_begin(indices)) {
+    // Ensure the structure of the type being indexed can be reasoned about.
+    // This includes rejecting any potential typed pointer.
+    auto destructurable =
+        llvm::dyn_cast<DestructurableTypeInterface>(selectedType);
+    if (!destructurable)
+      return nullptr;
+
+    // Follow the type at the index the gep is accessing, making it the new type
+    // used for indexing.
+    Type field = destructurable.getTypeAtIndex(index);
+    if (!field)
+      return nullptr;
+    selectedType = field;
+  }
+
+  // When there are no more indices, the type currently being used for indexing
+  // is the type of the value pointed at by the returned indexed pointer.
+  return selectedType;
+}
+
 //===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//
 
 /// Returns true if the given type is supported by atomic operations. All
 /// integer and float types with limited bit width are supported. Additionally,
 /// depending on the operation pointers may be supported as well.
 static bool isTypeCompatibleWithAtomicOp(Type type, bool isPointerTypeAllowed) {

mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp

 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Interfaces for AllocaOp
 //===----------------------------------------------------------------------===//
 
-static Type getAllocaElementType(LLVM::AllocaOp alloca) {
-  // This will become trivial once non-opaque pointers are gone.
-  return alloca.getElemType().has_value()
-             ? *alloca.getElemType()
-             : alloca.getResult().getType().getElementType();
-}
-
 llvm::SmallVector<MemorySlot> LLVM::AllocaOp::getPromotableSlots() {
   if (!getOperation()->getBlock()->isEntryBlock())
     return {};
 
-  return {MemorySlot{getResult(), getAllocaElementType(*this)}};
+  return {MemorySlot{getResult(), getResultPtrElementType(0)}};
 }
 
 Value LLVM::AllocaOp::getDefaultValue(const MemorySlot &slot,
                                       RewriterBase &rewriter) {
   return rewriter.create<LLVM::UndefOp>(getLoc(), slot.elemType);
 }
 
 void LLVM::AllocaOp::handleBlockArgument(const MemorySlot &slot,
if (defaultValue && defaultValue.use_empty())
     rewriter.eraseOp(defaultValue.getDefiningOp());
   rewriter.eraseOp(*this);
 }
 
 SmallVector<DestructurableMemorySlot> LLVM::AllocaOp::getDestructurableSlots() {
   if (!mlir::matchPattern(getArraySize(), m_One()))
     return {};
 
-  Type elemType = getAllocaElementType(*this);
+  Type elemType = getResultPtrElementType(0);
   auto destructurable = dyn_cast<DestructurableTypeInterface>(elemType);
   if (!destructurable)
     return {};
 
   std::optional<DenseMap<Attribute, Type>> destructuredType =
       destructurable.getSubelementIndexMap();
   if (!destructuredType)
     return {};
bool LLVM::GEPOp::canUsesBeRemoved(
   return forwardToUsers(*this, newBlockingUses);
 }
 
 DeletionKind LLVM::GEPOp::removeBlockingUses(
     const SmallPtrSetImpl<OpOperand *> &blockingUses, RewriterBase &rewriter) {
   return DeletionKind::Delete;
 }
 
-/// Returns the type the resulting pointer of the GEP points to. If such a type
-/// is not clear, returns null type.
-static Type computeReachedGEPType(LLVM::GEPOp gep) {
-  if (gep.getIndices().empty())
-    return {};
-
-  // Ensures all indices are static and fetches them.
-  SmallVector<IntegerAttr> indices;
-  for (auto index : gep.getIndices()) {
-    IntegerAttr indexInt = llvm::dyn_cast_if_present<IntegerAttr>(index);
-    if (!indexInt)
-      return {};
-    indices.push_back(indexInt);
-  }
-
-  // Check the pointer indexing only targets the first element.
-  if (indices[0].getInt() != 0)
-    return {};
-
-  // Set the initial type currently being used for indexing. This will be
-  // updated as the indices get walked over.
-  std::optional<Type> maybeSelectedType = gep.getElemType();
-  if (!maybeSelectedType)
-    return {};
-  Type selectedType = *maybeSelectedType;
-
-  // Follow the indexed elements in the gep.
-  for (IntegerAttr index : llvm::drop_begin(indices)) {
-    // Ensure the structure of the type being indexed can be reasoned about.
-    // This includes rejecting any potential typed pointer.
-    auto destructurable =
-        llvm::dyn_cast<DestructurableTypeInterface>(selectedType);
-    if (!destructurable)
-      return {};
-
-    // Follow the type at the index the gep is accessing, making it the new type
-    // used for indexing.
-    Type field = destructurable.getTypeAtIndex(index);
-    if (!field)
-      return {};
-    selectedType = field;
-  }
-
-  // When there are no more indices, the type currently being used for indexing
-  // is the type of the value pointed at by the returned indexed pointer.
-  return selectedType;
+static bool isFirstIndexZero(LLVM::GEPOp gep) {
+  IntegerAttr index =
+      llvm::dyn_cast_if_present<IntegerAttr>(gep.getIndices()[0]);
+  return index && index.getInt() == 0;
 }
 
 LogicalResult LLVM::GEPOp::ensureOnlySafeAccesses(
     const MemorySlot &slot, SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
   if (getBase() != slot.ptr)
     return success();
   if (slot.elemType != getElemType())
     return failure();
-  Type reachedType = computeReachedGEPType(*this);
+  if (!isFirstIndexZero(*this))
+    return failure();
+  Type reachedType = getResultPtrElementType(0);
   if (!reachedType)
     return failure();
   mustBeSafelyUsed.emplace_back<MemorySlot>({getResult(), reachedType});
   return success();
 }
 
 bool LLVM::GEPOp::canRewire(const DestructurableMemorySlot &slot,
                             SmallPtrSetImpl<Attribute> &usedIndices,
                             SmallVectorImpl<MemorySlot> &mustBeSafelyUsed) {
   auto basePtrType = llvm::dyn_cast<LLVM::LLVMPointerType>(getBase().getType());
   if (!basePtrType)
     return false;
 
   // Typed pointers are not supported. This should be removed once typed
   // pointers are removed from the LLVM dialect.
   if (!basePtrType.isOpaque())
     return false;
 
   if (getBase() != slot.ptr || slot.elemType != getElemType())
     return false;
-  Type reachedType = computeReachedGEPType(*this);
+  if (!isFirstIndexZero(*this))
+    return false;
+  Type reachedType = getResultPtrElementType(0);
   if (!reachedType || getIndices().size() < 2)
     return false;
   auto firstLevelIndex = cast<IntegerAttr>(getIndices()[1]);
   assert(slot.elementPtrs.contains(firstLevelIndex));
   if (!llvm::isa<LLVM::LLVMPointerType>(slot.elementPtrs.at(firstLevelIndex)))
     return false;
   mustBeSafelyUsed.emplace_back<MemorySlot>({getResult(), reachedType});
   usedIndices.insert(firstLevelIndex);

mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRLLVMIRTransforms
   DIScopeForLLVMFuncOp.cpp
   LegalizeForExport.cpp
   OptimizeForNVVM.cpp
   RequestCWrappers.cpp
+  TypeConsistency.cpp
 
   DEPENDS
   MLIRLLVMPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRFuncDialect
   MLIRLLVMDialect
   MLIRPass
   MLIRTransforms
   MLIRNVVMDialect
   )

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

+//===- TypeConsistency.cpp - Rewrites to improve type consistency ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/LLVMIR/Transforms/TypeConsistency.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMInterfaces.h"
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/MathExtras.h"
+#include <cstdint>
+#include <optional>
+#include <sys/types.h>
+
+namespace mlir {
+namespace LLVM {
+#define GEN_PASS_DEF_LLVMTYPECONSISTENCY
+#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc"
+} // namespace LLVM
+} // namespace mlir
+
+using namespace mlir;
+using namespace LLVM;
+
+//===----------------------------------------------------------------------===//
+// Utils
+//===----------------------------------------------------------------------===//
+
+/// Checks that a pointer value has a pointee type hint consistent with the
+/// expected type. Returns the type it actually hints to if it differs, or
+/// nullptr if the type is consistent or impossible to analyze.
+static Type isElementTypeInconsistent(Value addr, Type expectedType) {
+  auto defOp = dyn_cast_or_null<GetResultPtrElementType>(addr.getDefiningOp());
+  if (!defOp)
+    return nullptr;
+
+  Type elemType = defOp.getResultPtrElementType(0);
+  if (!elemType)
+    return nullptr;
+
+  if (elemType == expectedType)
+    return nullptr;
+
+  return elemType;
+}
+
+//===----------------------------------------------------------------------===//
+// AddFieldGetterToStructDirectUse
+//===----------------------------------------------------------------------===//
+
+/// Gets the type of the first subelement of `type` if `type` is destructurable,
+/// nullptr otherwise.
+static Type getFirstSubelementType(Type type) {
+  auto destructurable = dyn_cast<DestructurableTypeInterface>(type);
+  if (!destructurable)
+    return nullptr;
+
+  Type indexType = destructurable.getTypeAtIndex(
+      IntegerAttr::get(IntegerType::get(type.getContext(), 32), 0));
+  if (indexType)
+    return indexType;
+
+  indexType = destructurable.getTypeAtIndex(
+      IntegerAttr::get(IntegerType::get(type.getContext(), 64), 0));
+  if (indexType)
+    return indexType;
+
+  return nullptr;
+}
+
+/// Attempts to extract a pointer to the first field of an `elemType` from the
+/// address pointer of the provided MemOp, and rewire the MemOp so it uses that
+/// pointer instead.
+template <class MemOp>
+static LogicalResult
+attemptFieldIndirection(MemOp op, PatternRewriter &rewriter, Type elemType) {
+  PatternRewriter::InsertionGuard guard(rewriter);
+
+  rewriter.setInsertionPointAfterValue(op.getAddr());
+  SmallVector<GEPArg> firstTypeIndices{0, 0};
+
+  // TODO: Simplify once typed pointers are removed from LLVM dialect.
+  Value properPtr =
+      op.getAddr().getType().isOpaque()
+          ? rewriter.create<GEPOp>(
+                op->getLoc(), LLVM::LLVMPointerType::get(op.getContext()),
+                elemType, op.getAddr(), firstTypeIndices)
+          : rewriter.create<GEPOp>(op->getLoc(),
+                                   LLVM::LLVMPointerType::get(op.getContext()),
+                                   op.getAddr(), firstTypeIndices);
+
+  rewriter.updateRootInPlace(op, [&]() {
+    MutableOperandRange addr = op.getAddrMutable();
+    addr.clear();
+    addr.append(properPtr);
+  });
+
+  return success();
+}
+
+template <>
+LogicalResult AddFieldGetterToStructDirectUse<LoadOp>::matchAndRewrite(
+    LoadOp load, PatternRewriter &rewriter) const {
+  Type inconsistentElementType =
+      isElementTypeInconsistent(load.getAddr(), load.getType());
+  if (!inconsistentElementType)
+    return failure();
+  if (!getFirstSubelementType(inconsistentElementType))
+    return failure();
+  return attemptFieldIndirection<LoadOp>(load, rewriter,
+                                         inconsistentElementType);
+}
+
+template <>
+LogicalResult AddFieldGetterToStructDirectUse<StoreOp>::matchAndRewrite(
+    StoreOp store, PatternRewriter &rewriter) const {
+  PatternRewriter::InsertionGuard guard(rewriter);
+
+  Type inconsistentElementType =
+      isElementTypeInconsistent(store.getAddr(), store.getValue().getType());
+  if (!inconsistentElementType)
+    return failure();
+  Type firstType = getFirstSubelementType(inconsistentElementType);
+  if (!firstType)
+    return failure();
+
+  DataLayout layout = DataLayout::closest(store);
+  // Check that the first field has the right type or can at least be bitcast
+  // to the right type.
+  if (firstType != store.getValue().getType() &&
+      (isa<LLVMStructType, LLVMArrayType>(firstType) ||
+       layout.getTypeSize(firstType) !=
+           layout.getTypeSize(store.getValue().getType())))
+    return failure();
+
+  if (failed(attemptFieldIndirection<StoreOp>(store, rewriter,
+                                              inconsistentElementType)))
+    return failure();
+
+  Value replaceValue;
+  if (firstType == store.getValue().getType()) {
+    replaceValue = store.getValue();
+  } else {
+    rewriter.setInsertionPointAfterValue(store.getValue());
+    replaceValue = rewriter.create<BitcastOp>(store->getLoc(), firstType,
+                                              store.getValue());
+  }
+
+  rewriter.updateRootInPlace(store, [&]() {
+    MutableOperandRange value = store.getValueMutable();
+    value.clear();
+    value.append(replaceValue);
+  });
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// 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) {
+  size_t offset = indices[0] * layout.getTypeSize(base);
+
+  Type currentType = base;
+  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();
+              return false;
+            })
+            .Case([&](LLVMStructType structType) {
+              ArrayRef<Type> body = structType.getBody();
+              if (body.size() <= index)
+                return true;
+              for (uint32_t i = 0; i < index; i++) {
+                if (!structType.isPacked())
+                  offset = llvm::alignTo(offset,
+                                         layout.getTypeABIAlignment(body[i]));
+                offset += layout.getTypeSize(body[i]);
+              }
+              currentType = body[index];
+              return false;
+            })
+            .Default([](Type) { return true; });
+
+    if (shouldCancel)
+      return std::nullopt;
+  }
+
+  return offset;
+}
+
+/// Fills in `equivalentIndicesOut` with GEP indices that would be equivalent to
+/// offsetting a pointer by `offset` bytes, assuming the GEP has `base` as base
+/// type.
+static LogicalResult
+findIndicesForOffset(DataLayout &layout, Type base, uint64_t offset,
+                     SmallVectorImpl<GEPArg> &equivalentIndicesOut) {
+  uint64_t baseSize = layout.getTypeSize(base);
+  uint64_t rootIndex = offset / baseSize;
+  offset %= baseSize;
+  if (rootIndex > std::numeric_limits<uint32_t>::max())
+    return failure();
+  equivalentIndicesOut.push_back(rootIndex);
+
+  Type currentType = base;
+  while (offset > 0) {
+    bool shouldCancel =
+        TypeSwitch<Type, bool>(currentType)
+            .Case([&](LLVMArrayType arrayType) {
+              uint64_t elemSize =
+                  layout.getTypeSize(arrayType.getElementType());
+              uint64_t index = offset / elemSize;
+              offset %= elemSize;
+              equivalentIndicesOut.push_back(index);
+              currentType = arrayType.getElementType();
+              return index > std::numeric_limits<uint32_t>::max();
+            })
+            .Case([&](LLVMStructType structType) {
+              ArrayRef<Type> body = structType.getBody();
+              uint32_t index = 0;
+              for (Type elem : body) {
+                if (!structType.isPacked())
+                  offset =
+                      llvm::alignDown(offset, layout.getTypeABIAlignment(elem));
+                if (offset < layout.getTypeSize(elem)) {
+                  equivalentIndicesOut.push_back(index);
+                  currentType = elem;
+                  return false;
+                }
+                offset -= layout.getTypeSize(elem);
+                index++;
+              }
+              return true;
+            })
+            .Default([](Type) { return true; });
+
+    if (shouldCancel)
+      return failure();
+  }
+
+  return success();
+}
+
+LogicalResult
+CanonicalizeAlignedGep::matchAndRewrite(GEPOp gep,
+                                        PatternRewriter &rewriter) const {
+  std::optional<Type> maybeBaseType = gep.getElemType();
+  if (!maybeBaseType)
+    return failure();
+  Type baseType = *maybeBaseType;
+
+  Type typeHint = isElementTypeInconsistent(gep.getBase(), baseType);
+  if (!typeHint)
+    return failure();
+
+  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());
+  }
+
+  DataLayout layout = DataLayout::closest(gep);
+  std::optional<uint64_t> desiredOffset =
+      gepToByteOffset(layout, gep.getSourceElementType(), indices);
+  if (!desiredOffset)
+    return failure();
+
+  SmallVector<GEPArg> newIndices;
+  if (failed(
+          findIndicesForOffset(layout, typeHint, *desiredOffset, newIndices)))
+    return failure();
+
+  // TODO: Simplify when opaque pointers are gone.
+  if (gep.getElemType())
+    rewriter.replaceOpWithNewOp<GEPOp>(
+        gep, LLVM::LLVMPointerType::get(getContext()), typeHint, gep.getBase(),
+        newIndices, gep.getInbounds());
+  else
+    rewriter.replaceOpWithNewOp<GEPOp>(
+        gep, LLVM::LLVMPointerType::get(getContext()), gep.getBase(),
+        newIndices, gep.getInbounds());
+
+  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());
+    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

+// RUN: mlir-opt %s --pass-pipeline="builtin.module(llvm.func(llvm-type-consistency))" --split-input-file | FileCheck %s
+
+// CHECK-LABEL: llvm.func @same_address
+llvm.func @same_address(%arg: i32) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i32, i32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32, i32)> : (i32) -> !llvm.ptr
+  // CHECK: = llvm.getelementptr %[[ALLOCA]][0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32, i32)>
+  %7 = llvm.getelementptr %1[8] : (!llvm.ptr) -> !llvm.ptr, i8
+  llvm.store %arg, %7 : i32, !llvm.ptr
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @same_address_keep_inbounds
+llvm.func @same_address_keep_inbounds(%arg: i32) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i32, i32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32, i32)> : (i32) -> !llvm.ptr
+  // CHECK: = llvm.getelementptr inbounds %[[ALLOCA]][0, 2] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32, i32)>
+  %7 = llvm.getelementptr inbounds %1[8] : (!llvm.ptr) -> !llvm.ptr, i8
+  llvm.store %arg, %7 : i32, !llvm.ptr
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @struct_use_instead_of_first_field
+llvm.func @struct_use_instead_of_first_field(%arg: i32) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i32, i32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32, i32)> : (i32) -> !llvm.ptr
+  // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32, i32)>
+  // CHECK: llvm.store %{{.*}}, %[[GEP]] : i32
+  llvm.store %arg, %1 : i32, !llvm.ptr
+  llvm.return
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @struct_use_instead_of_first_field_same_size
+// CHECK-SAME: (%[[ARG:.*]]: f32)
+llvm.func @struct_use_instead_of_first_field_same_size(%arg: f32) {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK-DAG: %[[ALLOCA:.*]] = llvm.alloca %{{.*}} x !llvm.struct<"foo", (i32, i32, i32)>
+  %1 = llvm.alloca %0 x !llvm.struct<"foo", (i32, i32, i32)> : (i32) -> !llvm.ptr
+  // CHECK-DAG: %[[GEP:.*]] = llvm.getelementptr %[[ALLOCA]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<"foo", (i32, i32, i32)>
+  // CHECK-DAG: %[[BITCAST:.*]] = llvm.bitcast %[[ARG]] : f32 to i32
+  // CHECK: llvm.store %[[BITCAST]], %[[GEP]] : i32
+  llvm.store %arg, %1 : f32, !llvm.ptr
+  llvm.return
+}