EnzymeAD / Enzyme

High-performance automatic differentiation of LLVM and MLIR.
https://enzyme.mit.edu
Other
1.21k stars 100 forks source link

No reverse pass found for Petsc Functions #1375

Open akdemironur opened 11 months ago

akdemironur commented 11 months ago

Here is my failing code. The issue is that if I move the lines with PetscCallVoid (which are using a user struct) inside the for loop, or if I comment them out, there are no errors. However, in order to execute the for loop, I need that coneSize. I tried preprocessing it and storing it in the user struct, but when I attempted to retrieve coneSize from an array inside the struct, it failed as well. How can I solve this? (I'm using the latest release of Enzyme, v0.0.79)

Edit: I realized that if I hardcode coneSize, for example set it to 4, the code fails too. So struct might not be the thing which is triggering the failure. Petsc function are triggering the error. I tried writing a wrapper function in my source code for Petsc functions. But it didn't worked too. I added sample error. Actually these functions are not affecting my derivatives. If there is a way to tag them for that I could use it too.

No reverse pass found for DMPlexGetConeSize
declare !dbg !1321 i32 @DMPlexGetConeSize(%struct._p_DM* noundef, i32 noundef, i32* noundef) local_unnamed_addr #4
void formCellResidual(PetscReal *cellResidual, const PetscReal *x, PetscInt cellO, AppCtx *user)
{
    *cellResidual = 0;
    PetscReal TO = x[cellO], cellOCenter[3];
    const PetscInt *cone;
    PetscInt coneSize;
    DMLabel faceType;
    PetscCallVoid(DMPlexComputeCellGeometryFVM(user->dm, cellO, NULL, cellOCenter, NULL));
    PetscCallVoid(DMPlexGetConeSize(user->dm, cellO, &coneSize));
    PetscCallVoid(DMGetLabel(user->dm, "faceType", &faceType));
    PetscCallVoid(DMPlexGetCone(user->dm, cellO, &cone));
    for (PetscInt i = 0; i < coneSize; i++)
    {
        // some uber super calculations
    }
}
        __enzyme_autodiff(formCellResidual, enzyme_dupnoneed, &f, &grad_f, enzyme_dup, x, grad_x, enzyme_const, cellO,
                          enzyme_const, user);

Additionally, I have a side question. I'm currently using Enzyme, which I really appreciate. However, I sometimes come across errors, and unfortunately, the error messages or LLVM IR codes are not very meaningful to me. If you could provide a starting point for me to contribute to Enzyme, I'd be eager to give it a try.

akdemironur commented 11 months ago

Also I tried using Enzyme with LLDEnzyme-14.so as https://github.com/EnzymeAD/Enzyme/issues/648 but I got the same error.

wsmoses commented 11 months ago

Mind pasting the whole log and also perhaps the command line (and maybe even tarball) of the project you're compiling?

wsmoses commented 11 months ago

Regardless, very briefly skimming the docs, I assume you should mark the function as inactive using the new syntax in https://github.com/EnzymeAD/Enzyme/pull/1363

e.g.

__attribute__((enzyme_inactive))
PetscErrorCode DMPlexGetConeSize(DM dm, PetscInt p, PetscInt *size)

cc @jedbrown do you think petsc upstream would accept an attribute like the above being added (and relevant elsewhere). It will have no impact on the program if not using Enzyme (besides a possible compiler warning about an unused attribute), and if using Enzyme will inform the AD that the function should not be differentiated.

wsmoses commented 11 months ago

Regarding the last part of your question, I opened an issue recently that we should transition more of our error messages away from llvm assertions into compile time errors to make it easier for folks (including possibly for debugging this, depending on what error log you were seeing).

If you have cycles (or other things that you'd like to help contribute and improve), you're more than welcome to and I and others would be happy to help!

akdemironur commented 11 months ago

I created this minimal case. Also I noticed that if I keep parameters of PETSc function outside of my function it works.

I have attached the compilation error output and the corresponding log files. PETSc is a large library, so installing it will take your time. However, if it is necessary, I can provide a Dockerfile to attempt to compile this code.

void __enzyme_autodiff(void *, ...);
int enzyme_const, enzyme_dup, enzyme_out, enzyme_dupnoneed;

#include <petsc.h>

static char help[] = "";

typedef struct
{
    DM dm;
    PetscInt coneSize;
} AppCtx;

void formCellResidual(PetscReal *cellResidual, const PetscReal *x, PetscInt cellO, AppCtx *user)
{
    *cellResidual = 0;
    PetscInt *coneSize;
    DMPlexGetConeSize(user->dm, cellO, &coneSize); // this line fails
    // DMPlexGetConeSize(user->dm, cellO, &user->coneSize); // this line works
    for (PetscInt i = 0; i < user->coneSize; i++)
    {
        *cellResidual += x[i];
    }
}

void formCellJacobian(PetscReal *grad_x, const PetscReal *x, PetscInt cellO, AppCtx *user)
{
    PetscReal f, grad_f;
    grad_f = 1.0;
    __enzyme_autodiff(formCellResidual, enzyme_dupnoneed, &f, &grad_f, enzyme_dup, x, grad_x, enzyme_const, cellO,
                      enzyme_const, user);
}

int main(int argc, char **argv)
{
    PetscCall(PetscInitialize(&argc, &argv, NULL, help));
    PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Hello, Enzyme"));
    PetscCall(PetscFinalize());
}

enzyme-case-079668.c.log enzyme-case-079668.sh.log error.log

I am using CMake & Ninja so here is my compile_commands.json:

[
  {
    "arguments": [
      "/usr/bin/clang-14",
      "-fPIC",
      "-Qunused-arguments",
      "-g",
      "-O3",
      "-Wno-implicit-function-declaration",
      "-I/enzyme-case/include",
      "-isystem",
      "/reqs/petsc/include",
      "-Wall",
      "-g",
      "-Wextra",
      "-Xclang",
      "-load",
      "-Xclang",
      "/reqs/Enzyme/enzyme/build/Enzyme/ClangEnzyme-14.so",
      "-flegacy-pass-manager",
      "-fno-vectorize",
      "-fno-slp-vectorize",
      "-fno-unroll-loops",
      "-fuse-ld=lld",
      "-flto",
      "-c",
      "-I/reqs/petsc/petsc_clang14_release/include",
      "-o",
      "CMakeFiles/enzyme-case.dir/src/enzyme-case.c.o",
      "/enzyme-case/src/enzyme-case.c"
    ],
    "directory": "/enzyme-case/build",
    "file": "/enzyme-case/src/enzyme-case.c",
    "output": "/enzyme-case/build/CMakeFiles/enzyme-case.dir/src/enzyme-case.c.o"
  }
]
jedbrown commented 11 months ago

@wsmoses We could use a macro that expands to this attribute, but the question is where to put it. If the answer is "everywhere except certain numerical routines", we'd probably just put the attribute into PETSC_EXTERN (assuming the attribute can be used at the declaration site and not repeated at the definition (e.g., like visibility attributes), making a new PETSC_EXTERN_ACTIVE (or something more precise) for those functions that are active.

akdemironur commented 11 months ago

Regardless, very briefly skimming the docs, I assume you should mark the function as inactive using the new syntax in #1363

e.g.

__attribute__((enzyme_inactive))
PetscErrorCode DMPlexGetConeSize(DM dm, PetscInt p, PetscInt *size)

cc @jedbrown do you think petsc upstream would accept an attribute like the above being added (and relevant elsewhere). It will have no impact on the program if not using Enzyme (besides a possible compiler warning about an unused attribute), and if using Enzyme will inform the AD that the function should not be differentiated.

Btw, I've tried writing a wrapper function with enzyme_inactive attribute, but that didn't worked for my case.

Regarding the last part of your question, I opened an issue recently that we should transition more of our error messages away from llvm assertions into compile time errors to make it easier for folks (including possibly for debugging this, depending on what error log you were seeing).

If you have cycles (or other things that you'd like to help contribute and improve), you're more than welcome to and I and others would be happy to help!

I'm starting by reading your papers and code. If you have any suggestion for a newcomer that would be great! Thank you!

wsmoses commented 11 months ago

what was the error/source code with the enzyme_inactive.

Relatedly @tgymnich any thoughts on us adding PETSc to enzyme explorer?

akdemironur commented 11 months ago

I think that was related to my build system, after cleaning build directory, exact same code worked. Sorry for the inconvenience. for anyone interested here is the working version. This is more than enough for now :)

__attribute__((enzyme_inactive)) PetscErrorCode
EnzymeDMPlexGetConeSize(DM dm, PetscInt cellO, PetscInt *coneSize) {
  PetscCall(DMPlexGetConeSize(dm, cellO, coneSize));
  return 0;
}