Skip to content

Commit

Permalink
Add Support for Functions with Constant Array Type Parameters in Jaco…
Browse files Browse the repository at this point in the history 
…bian Mode

Now One must be able to find the Jacobian of functions with Constant Arrays in the parameter list.
For example, a function of the form:
```cpp
void func(double arr[3], double x, double y, double* output){
	output[0]=arr[2]*x*y;
	.
	.
	.
	output[n-1]=arr[0]*arr[1]*arr[2];
}
```

will generate a Jacobian of size n x 5.

Corresponding tests for the same have been written.

Previously this feature was not implememnted because it was necessary to know the size
of an array to correctly determine the size of the output jacobian matrix. This has now
been achieved for constant arrays, as we know the array sizes for them and hence we can
precisely locate where each jacobian entry for a array parameter is located. This is
achieved with the help of m_IndependentVarsSize, which stores the number of actual parameters
that each function parameter corresponds to.

Prior to this commit a std::map<ValueDecl, Expr> was used to map variable declarations to
their corresponding jacobian expression, so that we can lookup the latter during the reversemode
computations quickly. While this is fine for primitive variables(because each variable will only
correspond to one jacobian expression), an array declaration will correspond to multiple jacobian
expressions, depending on which index of the array is being referred to. Since the ValueDecl can
only refer to the name of the array and not name+index, we must update the map to use the name of
the array + index as the key. This can be done easily by using a string representation of the
ValueDecl name with the index as a suffix as the key of the map. Hence variables like
m_ExprVariables, m_VectorOutputString have been introduced to map array declarations, indexed by
position in array to their corresponding jacobian expressions.

Closes #472
  • Loading branch information
@Nirhar
Nirhar committed Aug 12, 2022
1 parent 5505764 commit fb567d0
Show file tree
Hide file tree
Showing 3 changed files with 252 additions and 29 deletions.
8 changes: 8 additions & 0 deletions include/clad/Differentiator/ReverseModeVisitor.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,9 @@ namespace clad {
class ExternalRMVSource;
class MultiplexExternalRMVSource;

using VectorOutputString =
std::vector<std::unordered_map<std::string, clang::Expr*>>;

/// A visitor for processing the function code in reverse mode.
/// Used to compute derivatives by clad::gradient.
class ReverseModeVisitor
Expand All @@ -38,6 +41,10 @@ namespace clad {
// several private/protected members of the visitor classes.
friend class ErrorEstimationHandler;
llvm::SmallVector<const clang::ValueDecl*, 16> m_IndependentVars;
llvm::SmallVector<int, 16> m_IndependentVarsSize;
std::unordered_map<std::string, clang::Expr*> m_ExprVariables;
VectorOutputString m_VectorOutputString;

/// In addition to a sequence of forward-accumulated Stmts (m_Blocks), in
/// the reverse mode we also accumulate Stmts for the reverse pass which
/// will be executed on return.
Expand All @@ -62,6 +69,7 @@ namespace clad {
std::vector<Stmts> m_LoopBlock;
unsigned outputArrayCursor = 0;
unsigned numParams = 0;
unsigned numActualParams = 0;
bool isVectorValued = false;
bool use_enzyme = false;
// FIXME: Should we make this an object instead of a pointer?
Expand Down
144 changes: 115 additions & 29 deletions lib/Differentiator/ReverseModeVisitor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -373,28 +373,54 @@ namespace clad {

// Creates the ArraySubscriptExprs for the independent variables
size_t idx = 0;
auto size_type = m_Context.getSizeType();
unsigned size_type_bits = m_Context.getIntWidth(size_type);
for (auto arg : args) {
// FIXME: fix when adding array inputs, now we are just skipping all
// array/pointer inputs (not treating them as independent variables).
if (utils::isArrayOrPointerType(arg->getType())) {
if (utils::isArrayOrPointerType(arg->getType())) { //is a array or pointer type parameter
if (arg->getName() == "p")
m_Variables[arg] = m_Result;

ParmVarDecl* parg=dyn_cast<ParmVarDecl>(const_cast<ValueDecl*>(arg));
QualType qt = parg->getOriginalType();
assert(qt->isConstantArrayType() && "Only Constant type arrays are allowed to be parameters of function whose jacobian is to be found. Non Constant types and Pointer type are not supported");
ConstantArrayType* t = dyn_cast<ConstantArrayType>(const_cast<Type*>(qt.getTypePtr()));
int sizeOfArray = (int)(t->getSize().roundToDouble(false));
m_IndependentVars.push_back(arg);
m_IndependentVarsSize.push_back(sizeOfArray);

for(int j=0;j<sizeOfArray;j++){
std::string name=arg->getNameAsString()+"["+std::to_string(j)+"]";
// Create the idx literal.
auto i =
IntegerLiteral::Create(m_Context, llvm::APInt(size_type_bits, idx),
size_type, noLoc);
// Create the jacobianMatrix[idx] expression.
auto result_at_i =
m_Sema.CreateBuiltinArraySubscriptExpr(m_Result, noLoc, i, noLoc)
.get();
m_ExprVariables[name]=result_at_i;
idx+=1;
numActualParams++;
}
}else{ //is normal variable parameter
// Create the idx literal.
auto i =
IntegerLiteral::Create(m_Context, llvm::APInt(size_type_bits, idx),
size_type, noLoc);
// Create the jacobianMatrix[idx] expression.
auto result_at_i =
m_Sema.CreateBuiltinArraySubscriptExpr(m_Result, noLoc, i, noLoc)
.get();
m_Variables[arg] = result_at_i;
m_ExprVariables[arg->getNameAsString()]=result_at_i;
idx += 1;
continue;
numActualParams++;
m_IndependentVars.push_back(arg);
m_IndependentVarsSize.push_back(1);
}
auto size_type = m_Context.getSizeType();
unsigned size_type_bits = m_Context.getIntWidth(size_type);
// Create the idx literal.
auto i =
IntegerLiteral::Create(m_Context, llvm::APInt(size_type_bits, idx),
size_type, noLoc);
// Create the jacobianMatrix[idx] expression.
auto result_at_i =
m_Sema.CreateBuiltinArraySubscriptExpr(m_Result, noLoc, i, noLoc)
.get();
m_Variables[arg] = result_at_i;
idx += 1;
m_IndependentVars.push_back(arg);

}
}

Expand Down Expand Up @@ -1067,7 +1093,44 @@ namespace clad {
auto ASI = SplitArraySubscript(ASE);
const Expr* Base = ASI.first;
const auto& Indices = ASI.second;
StmtDiff BaseDiff = Visit(Base);
StmtDiff BaseDiff;

//Check is we are visiting an Independent variable expressed as an array subscript expression in Jacobian Mode

if(isVectorValued){
if(auto dyn = dyn_cast<VarDecl>(dyn_cast<DeclRefExpr>(Base)->getDecl())){
//Check if this an independent variable
for(auto i:m_IndependentVars){
if(dyn->getNameAsString()==i->getNameAsString()){
llvm::APSInt intIdx;
auto isIdxValid =
clad_compat::Expr_EvaluateAsInt(ASE->getIdx(), intIdx, m_Context);

// FIXME: We assume that inside the index is just an Integer Expression
// and not a Complex expression
assert(isIdxValid && "Only Integer Literals allowed as array indices of Independent Variables");

int index = intIdx.getExtValue();
std::string indVarName = dyn->getNameAsString()+"["+std::to_string(index)+"]";
auto it = m_VectorOutputString[outputArrayCursor].find(indVarName);

// Create the (jacobianMatrix[idx] += dfdx) statement.
if (dfdx()) {
auto add_assign = BuildOp(BO_AddAssign, it->second, dfdx());
// Add it to the body statements.
addToCurrentBlock(add_assign, direction::reverse);
}
break;
}
}
BaseDiff = StmtDiff(dyn_cast<DeclRefExpr>(Clone(Base)));
}else{
BaseDiff = Visit(Base);
}
}else{
BaseDiff = Visit(Base);
}

llvm::SmallVector<Expr*, 4> clonedIndices(Indices.size());
llvm::SmallVector<Expr*, 4> reverseIndices(Indices.size());
llvm::SmallVector<Expr*, 4> forwSweepDerivativeIndices(Indices.size());
Expand Down Expand Up @@ -1918,21 +1981,44 @@ namespace clad {

std::unordered_map<const clang::ValueDecl*, clang::Expr*>
temp_m_Variables;
for (unsigned i = 0; i < numParams; i++) {
auto size_type = m_Context.getSizeType();
unsigned size_type_bits = m_Context.getIntWidth(size_type);
llvm::APInt idxValue(size_type_bits,
i + (outputArrayCursor * numParams));
auto idx = IntegerLiteral::Create(m_Context, idxValue,
size_type, noLoc);
// Create the jacobianMatrix[idx] expression.
auto result_at_i = m_Sema
.CreateBuiltinArraySubscriptExpr(
m_Result, noLoc, idx, noLoc)
.get();
temp_m_Variables[m_IndependentVars[i]] = result_at_i;
std::unordered_map<std::string,clang::Expr*> temp_m_VariablesStr;
auto size_type = m_Context.getSizeType();
unsigned size_type_bits = m_Context.getIntWidth(size_type);
for (unsigned i = 0, j=0; i < numParams; i++) {
auto arg = m_IndependentVars[i];
ParmVarDecl* parg=dyn_cast<ParmVarDecl>(const_cast<ValueDecl*>(arg));
if(parg->getOriginalType()->isConstantArrayType()){
int arrSize = m_IndependentVarsSize[i];
for(int k=0;k<arrSize;k++,j++){
llvm::APInt idxValue(size_type_bits,
j + (outputArrayCursor * numActualParams));
auto idx = IntegerLiteral::Create(m_Context, idxValue,
size_type, noLoc);
auto result_at_i = m_Sema
.CreateBuiltinArraySubscriptExpr(
m_Result, noLoc, idx, noLoc)
.get();
std::string sName = arg->getNameAsString()+"["+std::to_string(k)+"]";
temp_m_VariablesStr[sName]=result_at_i;
}
}else{
llvm::APInt idxValue(size_type_bits,
j + (outputArrayCursor * numActualParams));
auto idx = IntegerLiteral::Create(m_Context, idxValue,
size_type, noLoc);
// Create the jacobianMatrix[idx] expression.
auto result_at_i = m_Sema
.CreateBuiltinArraySubscriptExpr(
m_Result, noLoc, idx, noLoc)
.get();
temp_m_Variables[m_IndependentVars[i]] = result_at_i;
temp_m_VariablesStr[arg->getNameAsString()]=result_at_i;
j++;
}

}
m_VectorOutput.push_back(temp_m_Variables);
m_VectorOutputString.push_back(temp_m_VariablesStr);
}

auto dfdf = ConstantFolder::synthesizeLiteral(m_Context.IntTy,
Expand Down
129 changes: 129 additions & 0 deletions test/Jacobian/Jacobian.C
View file
Original file line number Diff line number Diff line change
Expand Up @@ -304,6 +304,118 @@ void f_1_jac_0(double a, double b, double c, double output[], double *jacobianMa
// CHECK-NEXT: }
// CHECK-NEXT:}

void f_5(float a[3], float output[]){
output[0]=a[0]*a[1];
output[1]=a[1]*a[2];
output[2]=a[0]*a[2];
}
void f_5_jac(float a[3], float output[], float *jacobianMatrix);
// CHECK: void f_5_jac(float a[3], float output[], float *jacobianMatrix) {
// CHECK-NEXT: float _t0;
// CHECK-NEXT: float _t1;
// CHECK-NEXT: float _t2;
// CHECK-NEXT: float _t3;
// CHECK-NEXT: float _t4;
// CHECK-NEXT: float _t5;
// CHECK-NEXT: _t1 = a[0];
// CHECK-NEXT: _t0 = a[1];
// CHECK-NEXT: output[0] = a[0] * a[1];
// CHECK-NEXT: _t3 = a[1];
// CHECK-NEXT: _t2 = a[2];
// CHECK-NEXT: output[1] = a[1] * a[2];
// CHECK-NEXT: _t5 = a[0];
// CHECK-NEXT: _t4 = a[2];
// CHECK-NEXT: output[2] = a[0] * a[2];
// CHECK-NEXT: {
// CHECK-NEXT: float _r4 = 1 * _t4;
// CHECK-NEXT: jacobianMatrix[6UL] += _r4;
// CHECK-NEXT: float _r5 = _t5 * 1;
// CHECK-NEXT: jacobianMatrix[8UL] += _r5;
// CHECK-NEXT: }
// CHECK-NEXT: {
// CHECK-NEXT: float _r2 = 1 * _t2;
// CHECK-NEXT: jacobianMatrix[4UL] += _r2;
// CHECK-NEXT: float _r3 = _t3 * 1;
// CHECK-NEXT: jacobianMatrix[5UL] += _r3;
// CHECK-NEXT: }
// CHECK-NEXT: {
// CHECK-NEXT: float _r0 = 1 * _t0;
// CHECK-NEXT: jacobianMatrix[0UL] += _r0;
// CHECK-NEXT: float _r1 = _t1 * 1;
// CHECK-NEXT: jacobianMatrix[1UL] += _r1;
// CHECK-NEXT: }
// CHECK-NEXT: }

void f_6(float a[1], float b, float output[]) {
output[0] = a[0] * a[0] * a[0];
output[1] = a[0] * a[0] * a[0] + b * b * b;
output[2] = 2 * (a[0] + b);
}
void f_6_jac(float a[1], float b, float output[], float *jacobianMatrix);
// CHECK: void f_6_jac(float a[1], float b, float output[], float *jacobianMatrix) {
// CHECK-NEXT: float _t0;
// CHECK-NEXT: float _t1;
// CHECK-NEXT: float _t2;
// CHECK-NEXT: float _t3;
// CHECK-NEXT: float _t4;
// CHECK-NEXT: float _t5;
// CHECK-NEXT: float _t6;
// CHECK-NEXT: float _t7;
// CHECK-NEXT: float _t8;
// CHECK-NEXT: float _t9;
// CHECK-NEXT: float _t10;
// CHECK-NEXT: float _t11;
// CHECK-NEXT: float _t12;
// CHECK-NEXT: _t2 = a[0];
// CHECK-NEXT: _t1 = a[0];
// CHECK-NEXT: _t3 = _t2 * _t1;
// CHECK-NEXT: _t0 = a[0];
// CHECK-NEXT: output[0] = a[0] * a[0] * a[0];
// CHECK-NEXT: _t6 = a[0];
// CHECK-NEXT: _t5 = a[0];
// CHECK-NEXT: _t7 = _t6 * _t5;
// CHECK-NEXT: _t4 = a[0];
// CHECK-NEXT: _t10 = b;
// CHECK-NEXT: _t9 = b;
// CHECK-NEXT: _t11 = _t10 * _t9;
// CHECK-NEXT: _t8 = b;
// CHECK-NEXT: output[1] = a[0] * a[0] * a[0] + b * b * b;
// CHECK-NEXT: _t12 = (a[0] + b);
// CHECK-NEXT: output[2] = 2 * (a[0] + b);
// CHECK-NEXT: {
// CHECK-NEXT: float _r12 = 1 * _t12;
// CHECK-NEXT: float _r13 = 2 * 1;
// CHECK-NEXT: jacobianMatrix[4UL] += _r13;
// CHECK-NEXT: jacobianMatrix[5UL] += _r13;
// CHECK-NEXT: }
// CHECK-NEXT: {
// CHECK-NEXT: float _r4 = 1 * _t4;
// CHECK-NEXT: float _r5 = _r4 * _t5;
// CHECK-NEXT: jacobianMatrix[2UL] += _r5;
// CHECK-NEXT: float _r6 = _t6 * _r4;
// CHECK-NEXT: jacobianMatrix[2UL] += _r6;
// CHECK-NEXT: float _r7 = _t7 * 1;
// CHECK-NEXT: jacobianMatrix[2UL] += _r7;
// CHECK-NEXT: float _r8 = 1 * _t8;
// CHECK-NEXT: float _r9 = _r8 * _t9;
// CHECK-NEXT: jacobianMatrix[3UL] += _r9;
// CHECK-NEXT: float _r10 = _t10 * _r8;
// CHECK-NEXT: jacobianMatrix[3UL] += _r10;
// CHECK-NEXT: float _r11 = _t11 * 1;
// CHECK-NEXT: jacobianMatrix[3UL] += _r11;
// CHECK-NEXT: }
// CHECK-NEXT: {
// CHECK-NEXT: float _r0 = 1 * _t0;
// CHECK-NEXT: float _r1 = _r0 * _t1;
// CHECK-NEXT: jacobianMatrix[0UL] += _r1;
// CHECK-NEXT: float _r2 = _t2 * _r0;
// CHECK-NEXT: jacobianMatrix[0UL] += _r2;
// CHECK-NEXT: float _r3 = _t3 * 1;
// CHECK-NEXT: jacobianMatrix[0UL] += _r3;
// CHECK-NEXT: }
// CHECK-NEXT: }


#define TEST(F, x, y, z) { \
result[0] = 0; result[1] = 0; result[2] = 0;\
result[3] = 0; result[4] = 0; result[5] = 0;\
Expand Down Expand Up @@ -335,4 +447,21 @@ int main() {
TEST(f_3, 1, 2, 3); // CHECK-EXEC: Result is = {22.69, 0.00, 0.00, 0.00, -17.48, 0.00, 0.00, 0.00, -41.58}
TEST(f_4, 1, 2, 3); // CHECK-EXEC: Result is = {84.00, 42.00, 0.00, 0.00, 126.00, 84.00, 126.00, 0.00, 42.00}
TEST_F_1_SINGLE_PARAM(1, 2, 3); // CHECK-EXEC: Result is = {3.00, 3.00, -2.00}


auto d_f_5 = clad::jacobian(f_5);
float a5[3]={3,4,5};
float op5[3]={0};
float jc5[9]={0};
d_f_5.execute(a5,op5,jc5);
printf("Result is = {%.2f, %.2f, %.2f, %.2f, %.2f, %.2f, %.2f, %.2f, %.2f}\n", jc5[0],jc5[1],jc5[2],jc5[3],jc5[4],jc5[5],jc5[6],jc5[7],jc5[8]);
//CHECK-EXEC: Result is = {4.00, 3.00, 0.00, 0.00, 5.00, 4.00, 5.00, 0.00, 3.00}

auto d_f_6 = clad::jacobian(f_6);
float a6[1]={3};float b6=5;
float op6[3]={0};
float jc6[6]={0};
d_f_6.execute(a6,b6,op6,jc6);
printf("Result is = {%.2f, %.2f, %.2f, %.2f, %.2f, %.2f}\n", jc6[0],jc6[1],jc6[2],jc6[3],jc6[4],jc6[5]);
//CHECK-EXEC: Result is = {27.00, 0.00, 27.00, 75.00, 2.00, 2.00}
}

0 comments on commit fb567d0

Please sign in to comment.