Add more constraints for handing multi-dimensional geps for global arrays

Author: usha1830

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -2975,6 +2975,43 @@ Value *InstCombiner::getFreelyInvertedImpl(Value *V, bool WillInvertAllUses,
   return nullptr;
 }
 
+/// Return true if we should lower multi-dimensional geps
+static bool ismultiDimGep(GetElementPtrInst &GEP) {
+  // Limit handling to only 3D and 4D arrays with integer types.
+  // getelementptr [9 x [9 x [9 x i32]]], ptr @arr, i64 0, i64 %i, i64 2, i64 3
+  unsigned NumOps = GEP.getNumOperands();
+
+  // First index must be constant zero (array base)
+  if (!isa<ConstantInt>(GEP.getOperand(1)) ||
+      !cast<ConstantInt>(GEP.getOperand(1))->isZero())
+    return false;
+
+  // Limit lowering for arrays with 3 or more dimensions
+  if (NumOps < 5)
+    return false;
+
+  // Check that it's arrays all the way
+  Type *CurTy = GEP.getSourceElementType();
+  unsigned NumVar = 0;
+  for (unsigned I = 2; I < NumOps; ++I) {
+    auto *ArrTy = dyn_cast<ArrayType>(CurTy);
+    if (!ArrTy)
+      return false;
+    if (!isa<ConstantInt>(GEP.getOperand(I)))
+      ++NumVar;
+    CurTy = ArrTy->getElementType();
+  }
+
+  // Limit lowering only for one variable index
+  if (NumVar != 1)
+    return false;
+
+  if (!CurTy->isIntegerTy() || CurTy->getIntegerBitWidth() > 128)
+    return false;
+
+  return true;
+}
+
 /// Return true if we should canonicalize the gep to an i8 ptradd.
 static bool shouldCanonicalizeGEPToPtrAdd(GetElementPtrInst &GEP) {
   Value *PtrOp = GEP.getOperand(0);
@@ -2995,13 +3032,7 @@ static bool shouldCanonicalizeGEPToPtrAdd(GetElementPtrInst &GEP) {
                                  m_Shl(m_Value(), m_ConstantInt())))))
     return true;
 
-  // Flatten multidimensional GEPs with one variable index.
-  unsigned NumVarIndices = 0;
-  for (unsigned i = 1; i < GEP.getNumOperands(); ++i) {
-    if (!isa<ConstantInt>(GEP.getOperand(i)))
-      ++NumVarIndices;
-  }
-  if (NumVarIndices == 1)
+  if (ismultiDimGep(GEP))
     return true;
 
   // gep (gep %p, C1), %x, C2 is expanded so the two constants can

llvm/test/Transforms/InstCombine/canonicalize-gep-constglob.ll

@@ -2,6 +2,9 @@
 ; RUN: opt < %s -passes=instcombine -S | FileCheck %s
 
 @glob = internal global [10 x [10 x [10 x i32]]] zeroinitializer
+@glob_i8 = internal global [10 x [10 x [10 x i8]]] zeroinitializer
+@glob_i16 = internal global [10 x [10 x [10 x i16]]] zeroinitializer
+@glob_i64 = internal global [10 x [10 x [10 x i64]]] zeroinitializer
 
 define ptr @x12(i64 %x) {
 ; CHECK-LABEL: define ptr @x12(
@@ -76,3 +79,61 @@ entry:
   %c = add i32 %a, %b
   ret i32 %c
 }
+
+define i8* @flat_gep8(i64 %x) {
+; CHECK-LABEL: define ptr @flat_gep8(
+; CHECK-SAME: i64 [[X:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[GEP_IDX:%.*]] = mul i64 [[X]], 100
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr @glob_i8, i64 [[GEP_IDX]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[TMP0]], i64 35
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+entry:
+  %gep = getelementptr [10 x [10 x [10 x i8]]], ptr @glob_i8, i64 0, i64 %x, i64 3, i64 5
+  ret ptr %gep
+}
+
+define i16* @flat_gep16(i64 %x) {
+; CHECK-LABEL: define ptr @flat_gep16(
+; CHECK-SAME: i64 [[X:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[GEP_IDX:%.*]] = mul i64 [[X]], 200
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr @glob_i16, i64 [[GEP_IDX]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[TMP0]], i64 46
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+entry:
+  %gep = getelementptr [10 x [10 x [10 x i16]]], ptr @glob_i16, i64 0, i64 %x, i64 2, i64 3
+  ret ptr %gep
+}
+
+define i32* @flat_gep(i64 %x) {
+; CHECK-LABEL: define ptr @flat_gep(
+; CHECK-SAME: i64 [[X:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[GEP_IDX:%.*]] = mul i64 [[X]], 400
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr @glob, i64 [[GEP_IDX]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[TMP0]], i64 100
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+entry:
+  %gep = getelementptr [10 x [10 x [10 x i32]]], ptr @glob, i64 0, i64 %x, i64 2, i64 5
+  ret ptr %gep
+}
+
+define i64* @flat_gep64(i64 %x) {
+; CHECK-LABEL: define ptr @flat_gep64(
+; CHECK-SAME: i64 [[X:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[GEP_IDX:%.*]] = mul i64 [[X]], 800
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr @glob_i64, i64 [[GEP_IDX]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[TMP0]], i64 288
+; CHECK-NEXT:    ret ptr [[GEP]]
+;
+entry:
+  %gep = getelementptr [10 x [10 x [10 x i64]]], ptr @glob_i64, i64 0, i64 %x, i64 3, i64 6
+  ret ptr %gep
+}
+
+