Skip to content
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

Jump to bottom

Add cblas level 1 routines and corresponding regression tests #364

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
+638 −1
Open

Add cblas level 1 routines and corresponding regression tests #364

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
7756342
Add cblas sscal
tansongchen Oct 25, 2021
5bd85f7
Implementation for BLAS level 1 and 2 routines
tansongchen Jan 4, 2022
bc0f779
Add many tests
tansongchen Jan 19, 2022
59df75a
Make tests more concise and fix several gradient computation bugs
tansongchen Jan 31, 2022
83c19d8
Improve and add test cases for BLAS level 1 routines
tansongchen Feb 2, 2022
b9c7431
Fix operand bundles and simplify test cases for BLAS
tansongchen Feb 9, 2022
17b7f8d
Temporarily remove generalization
tansongchen Feb 14, 2022
b1818f4
Improve adjoint of swap
tansongchen May 26, 2022
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
157 changes: 156 additions & 1 deletion enzyme/Enzyme/AdjointGenerator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4900,7 +4900,8 @@ class AdjointGenerator

std::string extractBLAS(StringRef in, std::string &prefix,
std::string &suffix) {
std::string extractable[] = {"ddot", "sdot", "dnrm2", "snrm2"};
std::string extractable[] = {"ddot", "sdot", "dnrm2", "snrm2",
"dswap", "sswap", "dcopy", "scopy"};
std::string prefixes[] = {"", "cblas_", "cublas_"};
std::string suffixes[] = {"", "_", "_64_"};
for (auto ex : extractable) {
Expand All @@ -4926,6 +4927,8 @@ class AdjointGenerator
IRBuilder<> allocationBuilder(gutils->inversionAllocs);
allocationBuilder.setFastMathFlags(getFast());

Module *currentModule = gutils->oldFunc->getParent();

if (funcName == "dnrm2" || funcName == "snrm2") {
if (!gutils->isConstantInstruction(&call)) {

Expand Down Expand Up @@ -5076,6 +5079,7 @@ class AdjointGenerator
}
return true;
}

if (funcName == "ddot" || funcName == "sdot") {
if (!gutils->isConstantInstruction(&call)) {
Type *innerType;
Expand Down Expand Up @@ -5659,6 +5663,157 @@ class AdjointGenerator
}
return true;
}

if ((funcName == "sswap" || funcName == "dswap") &&
called->isDeclaration()) {
// swap(n, x, incx, y, incy)
Type *innerType;
if (funcName == "dswap") {
innerType = Type::getDoubleTy(call.getContext());
} else if (funcName == "sswap") {
innerType = Type::getFloatTy(call.getContext());
} else {
assert(false && "Unreachable");
}
assert(Mode != DerivativeMode::ForwardMode &&
Mode != DerivativeMode::ForwardModeSplit &&
"forward mode not handled");
auto arg_n = call.getArgOperand(0), arg_x = call.getArgOperand(1),
arg_incx = call.getArgOperand(2), arg_y = call.getArgOperand(3),
arg_incy = call.getArgOperand(4);
auto type_n = arg_n->getType(), type_x = arg_x->getType(),
type_incx = arg_incx->getType(), type_y = arg_y->getType(),
type_incy = arg_incy->getType();
auto has_diff_x = !gutils->isConstantValue(arg_x),
has_diff_y = !gutils->isConstantValue(arg_y);
if (Mode == DerivativeMode::ReverseModeCombined ||
Mode == DerivativeMode::ReverseModeGradient) {
IRBuilder<> Builder2(call.getParent());
getReverseBuilder(Builder2);
auto new_n = lookup(gutils->getNewFromOriginal(arg_n), Builder2),
new_incx = lookup(gutils->getNewFromOriginal(arg_incx), Builder2),
new_incy = lookup(gutils->getNewFromOriginal(arg_incy), Builder2);
if (has_diff_x && has_diff_y) {
auto diff_x =
lookup(gutils->invertPointerM(arg_x, Builder2), Builder2),
diff_y =
lookup(gutils->invertPointerM(arg_y, Builder2), Builder2);
#if LLVM_VERSION_MAJOR >= 11
auto swapCall =
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add to differentialuseanalysis.h a corresponding note that neither primal x, nor primal y are needed for the reverse pass here.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Which function should I add to? and is there some example to see how to add them

FunctionCallee(call.getFunctionType(), call.getCalledOperand());
#else
auto swapCall = call.getCalledFunction();
#endif
SmallVector<Value *, 5> args = {new_n, diff_x, new_incx, diff_y,
new_incy};
SmallVector<ValueType, 5> vtypes = {
ValueType::None, ValueType::Shadow, ValueType::None,
ValueType::Shadow, ValueType::None};
auto defs = gutils->getInvertedBundles(&call, vtypes, Builder2, true);
Builder2.CreateCall(swapCall, args, defs);
} else if (has_diff_x || has_diff_y) {
auto toset =
has_diff_x
? lookup(gutils->invertPointerM(arg_x, Builder2), Builder2)
: lookup(gutils->invertPointerM(arg_y, Builder2), Builder2);
auto tosettype = has_diff_x ? type_x : type_y;
auto tosetinc = has_diff_x ? new_incx : new_incy;
auto memset = getOrInsertMemsetStrided(
*currentModule, cast<PointerType>(tosettype), type_n, 0);
SmallVector<Value *, 4> args = {
toset, ConstantFP::get(innerType, 0.0), new_n, tosetinc};
SmallVector<ValueType, 4> vtypes = {ValueType::Shadow,
ValueType::None, ValueType::None,
ValueType::None};
auto defs = gutils->getInvertedBundles(&call, vtypes, BuilderZ, true);
Builder2.CreateCall(memset, args, defs);
}
}
if (gutils->knownRecomputeHeuristic.find(&call) !=
gutils->knownRecomputeHeuristic.end()) {
if (!gutils->knownRecomputeHeuristic[&call]) {
gutils->cacheForReverse(BuilderZ, newCall,
getIndex(&call, CacheType::Self));
}
}

if (Mode == DerivativeMode::ReverseModeGradient) {
eraseIfUnused(call, /*erase*/ true, /*check*/ false);
} else {
eraseIfUnused(call);
}
return true;
}

if ((funcName == "scopy" || funcName == "dcopy") &&
called->isDeclaration()) {
// scopy(n, x, incx, y, incy)
Type *innerType;
std::string axpyName;
if (funcName == "dcopy") {
innerType = Type::getDoubleTy(call.getContext());
axpyName = (prefix + "daxpy" + suffix).str();
} else if (funcName == "scopy") {
innerType = Type::getFloatTy(call.getContext());
axpyName = (prefix + "saxpy" + suffix).str();
} else {
assert(false && "Unreachable");
}
auto arg_n = call.getArgOperand(0), arg_x = call.getArgOperand(1),
arg_incx = call.getArgOperand(2), arg_y = call.getArgOperand(3),
arg_incy = call.getArgOperand(4);
auto type_n = arg_n->getType(), type_x = arg_x->getType(),
type_incx = arg_incx->getType(), type_y = arg_y->getType(),
type_incy = arg_incy->getType();
bool has_diff_x = !gutils->isConstantValue(arg_x),
has_diff_y = !gutils->isConstantValue(arg_y);
if (Mode == DerivativeMode::ReverseModeCombined ||
Mode == DerivativeMode::ReverseModeGradient) {
IRBuilder<> Builder2(call.getParent());
getReverseBuilder(Builder2);
auto new_n = lookup(gutils->getNewFromOriginal(arg_n), Builder2),
new_incx = lookup(gutils->getNewFromOriginal(arg_incx), Builder2),
new_incy = lookup(gutils->getNewFromOriginal(arg_incy), Builder2);
if (has_diff_x) {
if (has_diff_y) {
auto diff_x =
lookup(gutils->invertPointerM(arg_x, Builder2), Builder2),
diff_y =
lookup(gutils->invertPointerM(arg_y, Builder2), Builder2);
auto axpyCall = currentModule->getOrInsertFunction(
axpyName, Builder2.getVoidTy(), type_n, innerType, type_y,
type_incy, type_x, type_incx);
SmallVector<Value *, 6> args = {
new_n, ConstantFP::get(innerType, 1.0),
diff_y, new_incy,
diff_x, new_incx};
Builder2.CreateCall(
axpyCall, args,
gutils->getInvertedBundles(&call,
{ValueType::None, ValueType::Shadow,
ValueType::None, ValueType::Shadow,
ValueType::None},
Builder2, true));
tansongchen marked this conversation as resolved.
Show resolved Hide resolved
}
}
}

if (gutils->knownRecomputeHeuristic.find(&call) !=
gutils->knownRecomputeHeuristic.end()) {
if (!gutils->knownRecomputeHeuristic[&call]) {
gutils->cacheForReverse(BuilderZ, newCall,
getIndex(&call, CacheType::Self));
}
}

if (Mode == DerivativeMode::ReverseModeGradient) {
eraseIfUnused(call, /*erase*/ true, /*check*/ false);
} else {
eraseIfUnused(call);
}
return true;
}

llvm::errs() << " fallback?\n";
return false;
}
Expand Down
82 changes: 82 additions & 0 deletions enzyme/Enzyme/Utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -293,6 +293,88 @@ Function *getOrInsertMemcpyStrided(Module &M, PointerType *T, Type *IT,
return F;
}

Function *getOrInsertMemsetStrided(Module &M, PointerType *T, Type *IT,
unsigned align) {
Type *elementType = T->getPointerElementType();
assert(elementType->isFloatingPointTy());
std::string name = "__enzyme_memset_" + tofltstr(elementType) + "_" +
std::to_string(cast<IntegerType>(IT)->getBitWidth()) +
"_align" + std::to_string(align) + "stride";
FunctionType *FT = FunctionType::get(Type::getVoidTy(M.getContext()),
{T, elementType, IT, IT}, false);

#if LLVM_VERSION_MAJOR >= 9
Function *F = cast<Function>(M.getOrInsertFunction(name, FT).getCallee());
#else
Function *F = cast<Function>(M.getOrInsertFunction(name, FT));
#endif

if (!F->empty())
return F;

F->setLinkage(Function::LinkageTypes::InternalLinkage);
F->addFnAttr(Attribute::ArgMemOnly);
F->addFnAttr(Attribute::NoUnwind);
F->addFnAttr(Attribute::AlwaysInline);
F->addParamAttr(0, Attribute::NoCapture);
F->addParamAttr(0, Attribute::WriteOnly);

BasicBlock *entry = BasicBlock::Create(M.getContext(), "entry", F);
BasicBlock *body = BasicBlock::Create(M.getContext(), "for.body", F);
BasicBlock *end = BasicBlock::Create(M.getContext(), "for.end", F);

auto dst = F->arg_begin();
dst->setName("dst");
auto val = dst + 1;
val->setName("val");
auto num = val + 1;
num->setName("num");
auto stride = num + 1;
stride->setName("stride");

{
IRBuilder<> B(entry);
B.CreateCondBr(B.CreateICmpEQ(num, ConstantInt::get(num->getType(), 0)),
end, body);
}

{
IRBuilder<> B(body);
B.setFastMathFlags(getFast());
PHINode *idx = B.CreatePHI(num->getType(), 2, "idx");
idx->addIncoming(ConstantInt::get(num->getType(), 0), entry);

#if LLVM_VERSION_MAJOR > 7
Value *dsti = B.CreateInBoundsGEP(dst->getType()->getPointerElementType(),
dst, idx, "dst.i");
#else
Value *dsti = B.CreateInBoundsGEP(dst, idx, "dst.i");
#endif

StoreInst *dsts = B.CreateStore(val, dsti);

if (align) {
#if LLVM_VERSION_MAJOR >= 10
dsts->setAlignment(Align(align));
#else
dsts->setAlignment(align);
#endif
}

Value *next =
B.CreateNUWAdd(idx, ConstantInt::get(num->getType(), 1), "idx.next");
idx->addIncoming(next, body);
B.CreateCondBr(B.CreateICmpEQ(num, next), end, body);
}

{
IRBuilder<> B(end);
B.CreateRetVoid();
}

return F;
}

// TODO implement differential memmove
Function *getOrInsertDifferentialFloatMemmove(Module &M, Type *T,
unsigned dstalign,
Expand Down
4 changes: 4 additions & 0 deletions enzyme/Enzyme/Utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -601,6 +601,10 @@ llvm::Function *getOrInsertMemcpyStrided(llvm::Module &M, llvm::PointerType *T,
llvm::Type *IT, unsigned dstalign,
unsigned srcalign);

/// Create function for type that performs memset with a stride
llvm::Function *getOrInsertMemsetStrided(llvm::Module &M, llvm::PointerType *T,
llvm::Type *IT, unsigned align);

/// Create function for type that performs the derivative memmove on floating
/// point memory
llvm::Function *
Expand Down
93 changes: 93 additions & 0 deletions enzyme/test/Enzyme/ReverseMode/blas/cblas_dcopy.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
;RUN: %opt < %s %loadEnzyme -enzyme -mem2reg -instsimplify -simplifycfg -S | FileCheck %s

@enzyme_const = common global i32 0, align 4

define void @wrapper(i32 %n, double* %x, i32 %incx, double* %y, i32 %incy) {
entry:
tail call void @cblas_dcopy(i32 %n, double* %x, i32 %incx, double* %y, i32 %incy)
ret void
}

declare void @cblas_dcopy(i32, double*, i32, double*, i32)

define void @wrapperMod(i32 %n, double* %x, i32 %incx, double* %y, i32 %incy) {
entry:
tail call void @cblas_dcopy(i32 %n, double* %x, i32 %incx, double* %y, i32 %incy)
store double 0.000000e+00, double* %x, align 8
store double 1.000000e+00, double* %y, align 8
ret void
}

define void @active(i32 %n, double* %x, double* %_x, i32 %incx, double* %y, double* %_y, i32 %incy) {
entry:
tail call void (i8*, ...) @__enzyme_autodiff(i8* bitcast (void (i32, double*, i32, double*, i32)* @wrapper to i8*), i32 %n, double* %x, double* %_x, i32 %incx, double* %y, double* %_y, i32 %incy)
ret void
}

declare void @__enzyme_autodiff(i8*, ...)

define void @inactiveX(i32 %n, double* %x, double* nocapture readnone %_x, i32 %incx, double* %y, double* %_y, i32 %incy) {
entry:
%0 = load i32, i32* @enzyme_const, align 4
tail call void (i8*, ...) @__enzyme_autodiff(i8* bitcast (void (i32, double*, i32, double*, i32)* @wrapper to i8*), i32 %n, i32 %0, double* %x, i32 %incx, double* %y, double* %_y, i32 %incy)
ret void
}

define void @inactiveY(i32 %n, double* %x, double* %_x, i32 %incx, double* %y, double* nocapture readnone %_y, i32 %incy) {
entry:
%0 = load i32, i32* @enzyme_const, align 4
tail call void (i8*, ...) @__enzyme_autodiff(i8* bitcast (void (i32, double*, i32, double*, i32)* @wrapper to i8*), i32 %n, double* %x, double* %_x, i32 %incx, i32 %0, double* %y, i32 %incy)
ret void
}

define void @activeMod(i32 %n, double* %x, double* %_x, i32 %incx, double* %y, double* %_y, i32 %incy) {
entry:
tail call void (i8*, ...) @__enzyme_autodiff(i8* bitcast (void (i32, double*, i32, double*, i32)* @wrapperMod to i8*), i32 %n, double* %x, double* %_x, i32 %incx, double* %y, double* %_y, i32 %incy)
ret void
}

;CHECK: define void @active
;CHECK-NEXT: entry
;CHECK-NEXT: call void @[[active:.+]](

;CHECK: define void @inactiveX
;CHECK-NEXT: entry
;CHECK-NEXT: call void @[[inactiveX:.+]](

;CHECK: define void @inactiveY
;CHECK-NEXT: entry
;CHECK-NEXT: call void @[[inactiveY:.+]](

;CHECK: define void @activeMod
;CHECK-NEXT: entry
;CHECK-NEXT: call void @[[activeMod:.+]](

;CHECK:define internal void @[[active]](i32 %n, double* %x, double* %"x'", i32 %incx, double* %y, double* %"y'", i32 %incy)
;CHECK-NEXT:entry:
;CHECK-NEXT: tail call void @cblas_dcopy(i32 %n, double* %x, i32 %incx, double* %y, i32 %incy)
;CHECK-NEXT: call void @cblas_daxpy(i32 %n, double 1.000000e+00, double* %"y'", i32 %incy, double* %"x'", i32 %incx)
;CHECK-NEXT: ret void
;CHECK-NEXT:}

;CHECK:define internal void @[[inactiveX]](i32 %n, double* %x, i32 %incx, double* %y, double* %"y'", i32 %incy)
;CHECK-NEXT:entry:
;CHECK-NEXT: tail call void @cblas_dcopy(i32 %n, double* %x, i32 %incx, double* %y, i32 %incy)
;CHECK-NEXT: ret void
;CHECK-NEXT:}

;CHECK:define internal void @[[inactiveY]](i32 %n, double* %x, double* %"x'", i32 %incx, double* %y, i32 %incy)
;CHECK-NEXT:entry:
;CHECK-NEXT: tail call void @cblas_dcopy(i32 %n, double* %x, i32 %incx, double* %y, i32 %incy)
;CHECK-NEXT: ret void
;CHECK-NEXT:}

;CHECK:define internal void @[[activeMod]](i32 %n, double* %x, double* %"x'", i32 %incx, double* %y, double* %"y'", i32 %incy)
;CHECK-NEXT:entry:
;CHECK-NEXT: tail call void @cblas_dcopy(i32 %n, double* %x, i32 %incx, double* %y, i32 %incy)
;CHECK-NEXT: store double 0.000000e+00, double* %x, align 8
;CHECK-NEXT: store double 1.000000e+00, double* %y, align 8
;CHECK-NEXT: store double 0.000000e+00, double* %"y'", align 8
;CHECK-NEXT: store double 0.000000e+00, double* %"x'", align 8
;CHECK-NEXT: call void @cblas_daxpy(i32 %n, double 1.000000e+00, double* %"y'", i32 %incy, double* %"x'", i32 %incx)
;CHECK-NEXT: ret void
;CHECK-NEXT:}
Loading