Compilation hangs when sorting about enzyme

[unrealwill]

I am sorting integers. They should be inactive.

[dan-zheng]

I'm trying to sort vector of integers inside the function I want to differentiate.
It should be a no operation in the backward pass.

Note: in general, I believe the derivative of sort is a permutation, not a no-op. It involves shuffling derivative values.
Edit: my bad, the issue is totally unrelated - thanks for explaining!

Differentiable sort example from Swift

import _Differentiation

// Workaround because `sort` is defined as a default implementation on `Collection` protocols, not on `Array`.
func _sort<T: RandomAccessCollection & MutableCollection>(_ x: inout T)
where T.Element: Comparable {
  x.sort()
}

// Stub defining `sort` on `Array` via the `Collection` protocol default implementation.
extension Array where Element: Comparable {
  mutating func sort() {
    _sort(&self)
  }
}

// Register a derivative for `Array.sort`.
extension Array where Element: Comparable & Differentiable {
  @derivative(of: sort)
  mutating func vjpSort() -> (
    value: (), pullback: (inout TangentVector) -> Void
  ) {
    // Original computation: sort while obtaining index permutation.
    // https://stackoverflow.com/a/17554343/5365899
    let sortedIndices = indices.sorted(by: { i, j in self[i] < self[j] })
    var sorted = self
    for i in indices {
      sorted[i] = self[sortedIndices[i]]
    }
    self = sorted

    // Pullback: apply sorted permutation to tangent array.
    func pullback(_ v: inout TangentVector) {
      var sortedBase = v.base
      for i in sortedBase.indices {
        sortedBase[i] = v.base[sortedIndices[i]]
      }
      v.base = sortedBase
    }
    return ((), pullback)
  }
}

// Usage example.

let array: [Float] = (0..<10).map(Float.init).shuffled()
let pb = pullback(at: array) { array -> [Float] in
  var tmp = array
  tmp.sort() // differentiable sort invocation
  return tmp
}
let v = (0..<10).map(Float.init)
print("   Array:", array)
print("Gradient:", pb(.init(v)))
//    Array: [5.0, 6.0, 2.0, 9.0, 0.0, 8.0, 3.0, 1.0, 4.0, 7.0]
// Gradient: [4.0, 7.0, 2.0, 6.0, 8.0, 0.0, 1.0, 9.0, 5.0, 3.0]

[wsmoses]

Pardon the delay (just got to computer), will look into this after your other bug report.

Yes this should be inactive for many reasons including the type and also the fact it is not used.

My guess as to what is occurring here is that Enzyme's Type Analysis is being too aggressive and attempting to assess types down a recursive call chain (which will eventually terminate due to limits on how big it is allowed to go, but clearly take way too long and should be made less aggressive).

[wsmoses]

Ok I've confirmed this indeed to be the case. Specifying a smaller limit via -mllvm -enzyme-max-type-offset=20 if you call with clang or -enzyme-max-type-offset=20 if you call with opt remedies the long compile time (the default is 500, which is indeed likely way larger than it should be).

On my machine, however, all of the levels of introsort/etc that it analyzes/differentiates through provides a similar recursive issue (this time err'ing with Enzyme's cache analysis stack overflowing). I'm now looking into the second piece, but thought I'd provide an update regardless.

[wsmoses]

Also FWIW, answering your original question -- yes Enzyme does have a way to mark functions as inactive by definition (by specifying their LLVM attributes to include enzyme_inactive). This is currently used by the Julia frontend (and perhaps Rust, in progress), though we haven't added a binding for it to clang.

[unrealwill]

Thanks for the quick response :) (even more so on a sunday :)
I don't quite know the procedure to add some LLVM attributes in a c++ code (maybe some pragma ...)

[wsmoses]

I just added a quick and dirty "this is inactive" call mechanism both in #129 and also my dev branch.

Calling Enzyme as follows:

~/git/Enzyme/build/bin/clang sorter.cpp -lstdc++ -lm -Xclang -load -Xclang Enzyme/ClangEnzyme-7.so -O2 -o test2 -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions -mllvm -enzyme-max-type-offset=20 

with the following modification will no longer result in a compiler crash. I'm going to leave this issue open regardless for me to have the issue fixed without manual intervention, but hopefully this unblocks you.

void __enzyme_call_inactive(void*, ...);

double foo( double* __restrict__ parts,int n, int* cellId)
 {
    vector< int > sorted(n);
    for( int i = 0 ; i < n ; i++)
    {
        sorted[i] = cellId[i];
    }

    //sort( sorted.begin(),sorted.end());
    __enzyme_call_inactive((void*)std::sort<typeof(sorted.begin())>, sorted.begin(), sorted.end());

    double out = parts[0];

    return out;
}

[unrealwill]

Thanks, but now the Option Parser can't seem to parse the -mllvm -enzyme-max-type-offset=20

clang test2.cpp -lstdc++ -lm -Xclang -load -Xclang ~/Enzyme/enzyme/build/Enzyme/ClangEnzyme-7.so -O2 -o test2 -fno-vectorize -fno-slp-vectorize -fno-unroll-loops -fno-exceptions -mllvm -enzyme-max-type-offset=20 
clang (LLVM option parsing): for the -enzyme-max-type-offset option: Cannot find option named '20'!

I have done
git checkout fortranrb
and remove then rebuild the build directories (llvm and enzyme) following the installation procedure.

[wsmoses]

It's likely that you may have needed to git fetch origin or the likes during checkout. Regardless, #120 has been merged into master which has the cache optimization that you desire and the following commit in master (21bb9d0) should ensure the command line flag functions as expected.

So in short, repull master and use that and let us know if it still errs?

[unrealwill]

Just did that, and it works fine :) 👍
I'll let you close the issue when you see fit as it seems you wanted for enzyme to be able to do so without manual intervention or extra-flags.

[wsmoses]

The original sort code without special handling now "works" though type analysis is still quite slow (3m50s on my system without special speedup flags).

That said the odd segfault during compilation error is now gone so I will close the issue for now (and look into speeding up type analysis).