From cd52cd7b55fba84c7c1fb24d03308d816d709b5e Mon Sep 17 00:00:00 2001
From: William Moses <gh@wsmoses.com>
Date: Tue, 7 May 2024 18:44:29 -0400
Subject: [PATCH] Allow gepoperator in EnzymeComputeByteOffsetOfGEP (#1865)

---
 enzyme/Enzyme/CApi.cpp | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/enzyme/Enzyme/CApi.cpp b/enzyme/Enzyme/CApi.cpp
index dc58942da147..929729cb6c90 100644
--- a/enzyme/Enzyme/CApi.cpp
+++ b/enzyme/Enzyme/CApi.cpp
@@ -1263,13 +1263,15 @@ LLVMValueRef EnzymeComputeByteOffsetOfGEP(LLVMBuilderRef B_r, LLVMValueRef V_r,
   IRBuilder<> &B = *unwrap(B_r);
   auto T = cast<IntegerType>(unwrap(T_r));
   auto width = T->getBitWidth();
-  auto gep = cast<GetElementPtrInst>(unwrap(V_r));
+  auto uw = unwrap(V_r);
+  GEPOperator *gep = isa<GetElementPtrInst>(uw)
+                         ? cast<GEPOperator>(cast<GetElementPtrInst>(uw))
+                         : cast<GEPOperator>(cast<ConstantExpr>(uw));
   auto &DL = B.GetInsertBlock()->getParent()->getParent()->getDataLayout();
 
   MapVector<Value *, APInt> VariableOffsets;
   APInt Offset(width, 0);
-  bool success =
-      collectOffset(cast<GEPOperator>(gep), DL, width, VariableOffsets, Offset);
+  bool success = collectOffset(gep, DL, width, VariableOffsets, Offset);
   (void)success;
   assert(success);
   Value *start = ConstantInt::get(T, Offset);