Generalize blas attributor (#2171)

EnzymeAD · Nov 19, 2024 · e42f5aa · e42f5aa
1 parent d9d6338
commit e42f5aa
Showing 1 changed file with 19 additions and 2 deletions.
diff --git a/enzyme/Enzyme/Utils.cpp b/enzyme/Enzyme/Utils.cpp
@@ -929,8 +929,25 @@ void callMemcpyStridedBlas(llvm::IRBuilder<> &B, llvm::Module &M, BlasInfo blas,
 
   FunctionType *FT = FunctionType::get(copy_retty, tys, false);
   auto fn = M.getOrInsertFunction(copy_name, FT);
-  Function *F = cast<Function>(fn.getCallee());
-  attributeKnownFunctions(*F);
+  Value *callVal = fn.getCallee();
+  Function *called = nullptr;
+  while (!called) {
+    if (auto castinst = dyn_cast<ConstantExpr>(callVal))
+      if (castinst->isCast()) {
+        callVal = castinst->getOperand(0);
+        continue;
+      }
+    if (auto fn = dyn_cast<Function>(callVal)) {
+      called = fn;
+      break;
+    }
+    if (auto alias = dyn_cast<GlobalAlias>(callVal)) {
+      callVal = alias->getAliasee();
+      continue;
+    }
+    break;
+  }
+  attributeKnownFunctions(*called);
 
   B.CreateCall(fn, args, bundles);
 }