[mlir][LLVM] Add `SplitGEP` type-consistency pattern

[zero9178]

Summary

The goal of this pattern is to eliminate all GEPs that have more than two indices by splitting it into multiple GEPs.
The advantage of this change is that the resulting GEPs only ever index into one aggregate at the time. This enables handling sub-aggregates in other patterns and also creates IR with easier to deduce pointer element types.

As a minor note, getResultPtrElementType for GEPOp was rewritten since it did not properly handle dynamic indices. The way GEPOp is specified, the resulting pointer element type can *always* be deduced from its base type and indices.

[zero9178]

rebase

[gysit]

Can you add a test for the pattern or is the result always consumed by a follow up pattern?

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

787–788

IRC we introduced the loop above since using drop_begin on getIndices resulted in strange crashes depending on the machine etc. Maybe copying the indices on the line above actually fixes the problem. Was this the reason for the additional copy?

789

nit: which can be structs or arrays.

800

Is it possibly that we can continue to use the destructible type interface here? Also are we sure we want to simplify anything if the indices are dynamic. SROA and Mem2Reg will anyways fail in these cases?

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

561

nit: getOp?

591–593

I would probably fill a vector directly?

SmallVector<GEPArg> newIndices = {0};
newIndices.append(splitIter, indices.end());

or similar?

[gysit]

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

561

gepOp of course...

[zero9178]

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

787–788

Yes. The iterators of GEPIndicesAdaptor keep a pointer to the instance. When writing llvm::drop_begin(getIndices()) one creates a temporary instance that gets deleted after the call to llvm::drop_begin (the value returned by llvm::drop_begin stays alive and llvm::drop_begin only copies the iterators of the range passed to it, not the range).

800

I think at the very least we should have getResultPtrElementType always return a correct result. If we want type consistency purely for SROA and Mem2Reg then dynamic indices would just cause useless GEP splits. For type consistency for the purpose of lifting, this would create nice and easy to work with GEPs.

That said, if for now we want to do it purely for the purpose of SROA and Mem2Reg then I think it'd make more sense to disable the dynamic indices case in the pattern rather than here.

[zero9178]

Address style comments

[zero9178]

add pure static GEP splitting test

[gysit]

LGTM!

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

800

True I was worried there may be a follow up pattern applying that should not apply. But you are right it is better to prevent this in these patterns!

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

 public:
   SplitStores(MLIRContext *context, unsigned maxVectorSplitSize)
       : OpRewritePattern(context), maxVectorSplitSize(maxVectorSplitSize) {}
 
   LogicalResult matchAndRewrite(StoreOp store,
                                 PatternRewriter &rewrite) const override;
 };
 
+/// Splits GEPs with more than two indices into multiple GEPs with exactly
+/// two indices. The created GEPs are then guaranteed to index into only
+/// one aggregate at a time.
+class SplitGEP : 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

if (std::optional<Type> elemType = getElemType())
     return *elemType;
 
   return llvm::cast<LLVMPointerType>(
              extractVectorElementType(getBase().getType()))
       .getElementType();
 }
 
 Type GEPOp::getResultPtrElementType() {
-  // 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;
+  auto indices = getIndices();
+  for (GEPIndicesAdaptor<ValueRange>::value_type index :
+       llvm::drop_begin(indices)) {
+    // GEPs can only index into aggregates which are only arrays and structs.
+
+    // The resulting type if indexing into an array type is always the element
+    // type, regardless of index.
+    if (auto arrayType = dyn_cast<LLVMArrayType>(selectedType)) {
+      selectedType = arrayType.getElementType();
+      continue;
+    }
 
-    // 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;
+    // The GEP verifier ensures that any index into structs are static and
+    // that they refer to a field within the struct.
+    int64_t fieldIndex = cast<IntegerAttr>(index).getValue().getZExtValue();
+    selectedType = cast<LLVMStructType>(selectedType).getBody()[fieldIndex];
   }
 
   // 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;
 }
 
 //===----------------------------------------------------------------------===//

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

bool LLVM::GEPOp::canRewire(const DestructurableMemorySlot &slot,
 
   if (getBase() != slot.ptr || slot.elemType != getElemType())
     return false;
   if (!isFirstIndexZero(*this))
     return false;
   Type reachedType = getResultPtrElementType();
   if (!reachedType || getIndices().size() < 2)
     return false;
-  auto firstLevelIndex = cast<IntegerAttr>(getIndices()[1]);
+  auto firstLevelIndex = dyn_cast<IntegerAttr>(getIndices()[1]);
+  if (!firstLevelIndex)
+    return false;
   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);
   return true;
 }
 

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

LogicalResult SplitStores::matchAndRewrite(StoreOp store,
   // 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();
 }
 
+LogicalResult SplitGEP::matchAndRewrite(GEPOp gep,
+                                        PatternRewriter &rewriter) const {
+  FailureOr<Type> typeHint = getRequiredConsistentGEPType(gep);
+  if (succeeded(typeHint) || gep.getIndices().size() <= 2) {
+    // GEP is not canonical or a single aggregate deep, nothing to do here.
+    return failure();
+  }
+
+  auto indexToGEPArg =
+      [](GEPIndicesAdaptor<ValueRange>::value_type index) -> GEPArg {
+    if (auto integerAttr = dyn_cast<IntegerAttr>(index))
+      return integerAttr.getValue().getSExtValue();
+    return cast<Value>(index);
+  };
+
+  GEPIndicesAdaptor<ValueRange> indices = gep.getIndices();
+
+  auto splitIter = std::next(indices.begin(), 2);
+
+  // Split of the first GEP using the first two indices.
+  auto subGepOp = rewriter.create<GEPOp>(
+      gep.getLoc(), gep.getType(), gep.getSourceElementType(), gep.getBase(),
+      llvm::map_to_vector(llvm::make_range(indices.begin(), splitIter),
+                          indexToGEPArg),
+      gep.getInbounds());
+
+  // The second GEP indexes on the result pointer element type of the previous
+  // with all the remaining indices and a zero upfront. If this GEP has more
+  // than two indices remaining it'll be further split in subsequent pattern
+  // applications.
+  auto gepArgs = llvm::map_to_vector(llvm::make_range(splitIter, indices.end()),
+                                     indexToGEPArg);
+  gepArgs.insert(gepArgs.begin(), 0);
+  rewriter.replaceOpWithNewOp<GEPOp>(gep, gep.getType(),
+                                     subGepOp.getResultPtrElementType(),
+                                     subGepOp, gepArgs, 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());
     rewritePatterns.add<SplitStores>(&getContext(), maxVectorSplitSize);
+    rewritePatterns.add<SplitGEP>(&getContext());
     FrozenRewritePatternSet frozen(std::move(rewritePatterns));
 
     if (failed(applyPatternsAndFoldGreedily(getOperation(), frozen)))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> LLVM::createTypeConsistencyPass() {
   return std::make_unique<LLVMTypeConsistencyPass>();
 }