Skip to content

complex blas reverse mode gemv/gemm #3190

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ymardoukhi wants to merge 10 commits into
base: main
Choose a base branch
from
Open

complex blas reverse mode gemv/gemm #3190

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
f6adcb7
remove guard for complex-valued reverse mode; fixed `dot/ger` naming.
ymardoukhi Aug 17, 2026
dd6fa42
alpha/beta ABI fixes
ymardoukhi Aug 31, 2026
839f500
CBLAS transpose flags
ymardoukhi Aug 31, 2026
f8e84c1
active alpha/beta for gemm (FrobInnerProd)
ymardoukhi Aug 24, 2026
73a54e6
fix Constant<> splat for complex fpType
ymardoukhi Aug 25, 2026
944d543
complex CBLAS gemv integration test
ymardoukhi Aug 25, 2026
5d29f83
complex CBLAS gemm integration test
ymardoukhi Aug 26, 2026
380ac64
add zgemv/zgemm "overwrite" caching tests
ymardoukhi Aug 26, 2026
69f9f7d
Merge branch 'main' into feat/complex-blas-reverse-mode
ymardoukhi Sep 7, 2026
8ea7052
Merge branch 'main' into feat/complex-blas-reverse-mode
ymardoukhi Sep 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 43 additions & 9 deletions enzyme/Enzyme/Utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1655,13 +1655,24 @@ void callSPMVDiagUpdate(IRBuilder<> &B, Module &M, BlasInfo blas,
return;
}

llvm::Constant *getRealValuedConstant(llvm::Type *fpType, double val) {
if (auto VT = dyn_cast<VectorType>(fpType)) {
auto elTy = VT->getElementType();
return ConstantVector::get(
{ConstantFP::get(elTy, val), ConstantFP::get(elTy, 0.0)});
}
return ConstantFP::get(fpType, val);
}

llvm::CallInst *
getorInsertInnerProd(llvm::IRBuilder<> &B, llvm::Module &M, BlasInfo blas,
IntegerType *IT, Type *BlasPT, Type *BlasIT, Type *fpTy,
llvm::ArrayRef<llvm::Value *> args,
const llvm::ArrayRef<llvm::OperandBundleDef> bundles,
bool byRef, bool cublas, bool julia_decl) {
assert(fpTy->isFloatingPointTy());
// relax to receive complex values.
bool isComplex = (blas.floatType == "c" || blas.floatType == "z");
assert(fpTy->isFloatingPointTy() || isComplex);

// add inner_prod call if not already present
std::string prod_name = "__enzyme_inner_prod" + blas.floatType + blas.suffix;
Expand All @@ -1673,8 +1684,9 @@ getorInsertInnerProd(llvm::IRBuilder<> &B, llvm::Module &M, BlasInfo blas,
if (!F->empty())
return B.CreateCall(F, args, bundles);

// add dot call if not already present
std::string dot_name = blas.prefix + blas.floatType + "dot" + blas.suffix;
// for the complex varient use `dotc`.
std::string dot_name = blas.prefix + blas.floatType +
(isComplex ? "dotc" : "dot") + blas.suffix;
auto FDotT =
FunctionType::get(fpTy, {BlasIT, BlasPT, BlasIT, BlasPT, BlasIT}, false);
auto FDot = M.getOrInsertFunction(dot_name, FDotT);
Expand Down Expand Up @@ -4140,11 +4152,21 @@ llvm::Value *transpose(std::string floatType, IRBuilder<> &B, llvm::Value *V,
}

} else if (T->isIntegerTy(32)) {
auto is111 = B.CreateICmpEQ(V, ConstantInt::get(V->getType(), 111));
auto sel1 = B.CreateSelect(
B.CreateICmpEQ(V, ConstantInt::get(V->getType(), 112)),
ConstantInt::get(V->getType(), 111), ConstantInt::get(V->getType(), 0));
return B.CreateSelect(is111, ConstantInt::get(V->getType(), 112), sel1);
// CBLAS_TRANSPOSE: CblasNoTrans=111, CblasTrans=112, CblasConjTrans=113.
// See the Fortran codeblock above: 111<->113/'N'<->'C'. 112 case not handled
if (floatType == "z" || floatType == "c") {
auto is111 = B.CreateICmpEQ(V, ConstantInt::get(V->getType(), 111));
auto sel1 = B.CreateSelect(
B.CreateICmpEQ(V, ConstantInt::get(V->getType(), 113)),
ConstantInt::get(V->getType(), 111), ConstantInt::get(V->getType(), 0));
return B.CreateSelect(is111, ConstantInt::get(V->getType(), 113), sel1);
} else {
auto is111 = B.CreateICmpEQ(V, ConstantInt::get(V->getType(), 111));
auto sel1 = B.CreateSelect(
B.CreateICmpEQ(V, ConstantInt::get(V->getType(), 112)),
ConstantInt::get(V->getType(), 111), ConstantInt::get(V->getType(), 0));
return B.CreateSelect(is111, ConstantInt::get(V->getType(), 112), sel1);
}
} else {
std::string s;
llvm::raw_string_ostream ss(s);
Expand Down Expand Up @@ -4205,11 +4227,23 @@ llvm::Value *transpose(std::string floatType, llvm::IRBuilder<> &B,
}

// cblas
if (!cublas)
if (!cublas) {
// CBLAS_TRANSPOSE swap 111<->113/'N'<->'C'

if (floatType == "c" || floatType == "z") {
auto is111 = B.CreateICmpEQ(V, ConstantInt::get(V->getType(), 111));
auto sel1 = B.CreateSelect(
B.CreateICmpEQ(V, ConstantInt::get(V->getType(), 113)),
ConstantInt::get(V->getType(), 111),
ConstantInt::get(V->getType(), 0));
return B.CreateSelect(is111, ConstantInt::get(V->getType(), 113),
sel1);
}
return B.CreateSelect(
B.CreateICmpEQ(V, ConstantInt::get(V->getType(), 111)),
ConstantInt::get(V->getType(), 112),
ConstantInt::get(V->getType(), 111));
}
}

if (byRef) {
Expand Down
3 changes: 3 additions & 0 deletions enzyme/Enzyme/Utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -915,6 +915,9 @@ getorInsertInnerProd(llvm::IRBuilder<> &B, llvm::Module &M, BlasInfo blas,
const llvm::ArrayRef<llvm::OperandBundleDef> bundles,
bool byRef, bool cublas, bool julia_decl);

// Costruct Constant<1.0> for the complex numbers correctly i.e. 1.0 + 0.0i.
llvm::Constant *getRealValuedConstant(llvm::Type *fpType, double val);

/// Create function for type that performs memcpy with a stride
llvm::Function *getOrInsertMemcpyStrided(llvm::Module &M,
llvm::Type *elementType,
Expand Down
Loading
Loading