EnzymeAD / Enzyme

High-performance automatic differentiation of LLVM and MLIR.

Home Page:https://enzyme.mit.edu

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Enzyme crash

davidedelvento opened this issue · comments

I'm trying to get more and more familiar with Enzyme, and I am slowly understanding some subtleties. For example the fact that you cannot specify the argument types if you are using only the LLD plugin, you need the Clang plugin for that. Now, in this particular case I can get away with using the Clang plugin only but I am leaving the LLD also because I will soon need to go multisource.

So I modified one of the examples to get that working, and then expanded the example further to have it multi-dimensional (second call to enzyme for the sumAndMulArray function) and that is when it crashed

Source code

#include <stdio.h>

int enzyme_dup;
int enzyme_dupnoneed;
int enzyme_out;
int enzyme_const;

double __enzyme_autodiff(void*, ...);

double sumAndMul(double* array, int size, double mul) {
  double sum = 0;
  for(int i=0; i<size; i++) {
    sum += array[i];
  }
  return sum * mul;
}

double sumAndMulArray(double* array, int size, double* mul) {
  double sum = 0;
  for(int i=0; i<size; i++) {
    sum += array[i] * mul[i];
  }
  return sum;
}

int main() {
    double array[4] = {1.0, 2.0, 3.0, 4.0};
    // double d_array[4] = {1.0, 2.0, 3.0, 4.0};
    double d_array[4] = {.0, .0, .0, .0};
    int size = 4;
    double mul = 3.0;

    double tot = sumAndMul(array, size, mul);
    double d_mul = __enzyme_autodiff((void*)sumAndMul,
                     enzyme_dup,   array, d_array,
                     enzyme_const, size,
                     enzyme_out,   mul);

    printf("d_array=%f %f %f %f, mul=%f\n", d_array[0],  d_array[1],  d_array[2],  d_array[3], mul);
    printf("tot=%f, der=%f\n", tot, d_mul);

    double mul_array[4] = {3.0, 4.0, 5.0, 6.0};

    tot = sumAndMulArray(array, size, mul_array);
    d_mul = __enzyme_autodiff((void*)sumAndMulArray,
                     enzyme_dup,   array, d_array,
                     enzyme_const, size,
                     enzyme_out,   mul_array);

    printf("d_array=%f %f %f %f, mul=%f\n", d_array[0],  d_array[1],  d_array[2],  d_array[3], mul);
    printf("tot=%f, der=%f\n", tot, d_mul);
}

Makefile

CC=clang
CFLAGS = -fopenmp=libomp -flto
LIBS=-lm
ENZ     = ${HOME}/repositories/ENZYME/Enzyme-v0.0.100/enzyme/build-v0.0.100-with-llvm-v16.0.2/Enzyme/LLDEnzyme-16.so

all : test.o
        $(CC) -fuse-ld=lld $(CFLAGS) $^ -o test.exe $(LIBS) -v -Wl,--load-pass-plugin=$(ENZ)

test.o : test.c
        $(CC) $(CFLAGS) -c $^ -v -fplugin=${HOME}/repositories/ENZYME/Enzyme-v0.0.100/enzyme/build-v0.0.100-with-llvm-v16.0.2/Enzyme/ClangEnzyme-16.so

.PHONY : clean
clean :
        rm -f *.o *.exe

crash

; Function Attrs: mustprogress noinline nounwind optnone willreturn memory(readwrite) uwtable
define internal { ptr } @diffesumAndMulArray(ptr noundef
 %0, ptr
 %1, i32 noundef
 %2, ptr noundef
 %3, double
 %4) #7 {
  %6 = call i32 @llvm.smax.i32(i32 %2, i32 0)
  %7 = zext i32 %6 to i64                                                                                                                                                                                                   br label %8
                                                                                                                                                                                                                          8:                                                ; preds = %14, %5
  %9 = phi i64 [ %11, %14 ], [ 0, %5 ]                                                                                                                                                                                      %10 = phi double [ 0.000000e+00, %5 ], [ %21, %14 ]
  %11 = add nuw nsw i64 %9, 1
  %12 = trunc i64 %9 to i32
  %13 = icmp slt i32 %12, %2
  br i1 %13, label %14, label %23

14:                                               ; preds = %8
  %15 = sext i32 %12 to i64
  %16 = getelementptr inbounds double, ptr %0, i64 %15
  %17 = load double, ptr %16, align 8
  %18 = sext i32 %12 to i64
  %19 = getelementptr inbounds double, ptr %3, i64 %18
  %20 = load double, ptr %19, align 8
  %21 = call double @llvm.fmuladd.f64(double %17, double %20, double %10) #8
  %22 = add nsw i32 %12, 1                                                                                                                                                                                                  br label %8, !llvm.loop !11
                                                                                                                                                                                                                          23:                                               ; preds = %8
  br label %30                                                                                                                                                                                                            
24:                                               ; No predecessors!
  %25 = alloca i64, align 8
  %26 = alloca double, align 8
  store double 0.000000e+00, ptr %26, align 8

27:                                               ; No predecessors!

28:                                               ; No predecessors!

29:                                               ; No predecessors!

30:                                               ; preds = %23
  store double %4, ptr %26, align 8
}

ptr %2
clang: ../Enzyme/DiffeGradientUtils.cpp:215: Value *DiffeGradientUtils::diffe(Value *, IRBuilder<> &): Assertion `!val->getType()->isPointerTy()' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace, preprocessed source, and associated run script.

So assuming the first invocation to Enzyme (for sumAndMul) and its two-step compilation line is correct (I mean, it works, but it may be by accident), am I doing something wrong wrt how I extended it to the second invocation (for sumAndMulArray) that made it crash, or have I really encounter a bug to crash the compiler?

PS: not sure I understand this trace, however I noticed there is no %5.

@davidedelvento you can specify the argument types in lld, but you need to use extern int enzyme_const etc

And the issue here is that enzyme_out only applies to immutable data-types aka float, and not mutable memory, like float*.

Here you use enzyme_out on mul_array, which is a mutable pointer and an error.

Thanks for your patience in explaining me even these basics! I understand now and the problem is solved.
For the sake of others who may find this while having a similar problem here is the working code (with a bit sloppy variable names)

#include <stdio.h>

extern int enzyme_dup;
extern int enzyme_dupnoneed;
extern int enzyme_out;
extern int enzyme_const;

double __enzyme_autodiff(void*, ...);

double sumAndMul(double* array, int size, double* mul) {
  double sum = 0;
  for(int i=0; i<size; i++) {
    sum += array[i];
  }
  return sum * mul[0];
}

double sumAndMulArray(double* array, int size, double* mul) {
  double sum = 0;
  for(int i=0; i<size; i++) {
    sum += array[i] * mul[i];
  }
  return sum;
}

int main() {
    double array[4] = {1.0, 2.0, 3.0, 4.0};
    // double d_array[4] = {1.0, 2.0, 3.0, 4.0};
    double d_array[4] = {.0, .0, .0, .0};
    int size = 4;
    double mul[1] = {3.0};
    double d_mul[1] = {.0};

    double tot = sumAndMul(array, size, mul);
    __enzyme_autodiff((void*)sumAndMul,
                     enzyme_dup,   array, d_array,
                     enzyme_const, size,
                     enzyme_dup,   mul, d_mul);

    printf("d_array=%f %f %f %f, mul=%f\n", d_array[0],  d_array[1],  d_array[2],  d_array[3], mul[0]);
    printf("tot=%f, der=%f\n", tot, d_mul[0]);
    printf("=====\n");

    for(int i=0; i<size; i++)
        d_array[i] = .0;
    double mul_array[4] = {3.0, 4.0, 5.0, 6.0};
    double dmul_array[4] = {.0, .0, .0, .0};

    tot = sumAndMulArray(array, size, mul_array);
    __enzyme_autodiff((void*)sumAndMulArray,
                     enzyme_dup,   array, d_array,
                     enzyme_const, size,
                     enzyme_dup,   mul_array, dmul_array);

    printf("d_array=%f %f %f %f\n", d_array[0],  d_array[1],  d_array[2],  d_array[3]);
    printf("mul_array=%f %f %f %f\n", mul_array[0],  mul_array[1],  mul_array[2],  mul_array[3]);
    printf("dmul_array=%f %f %f %f\n", dmul_array[0],  dmul_array[1],  dmul_array[2],  dmul_array[3]);
    printf("tot=%f\n", tot);
}

Compiling and running it produces the expected output, that is differentiation wrt both the data array and the mul array. Of course the original example (with scalar mul) works just fine with an enzyme_out option, but as soon as that becomes an array one has to use the duplicate approach.