Fix shuffle canonicalization (rust-lang#819)

* Fix shuffle canonicalization

* Add assertion

Co-authored-by: Manuel Drehwald <[email protected]>

Author: wsmoses

enzyme/Enzyme/AdjointGenerator.h

@@ -11697,6 +11697,7 @@ class AdjointGenerator
         if (subretused) {
           Value *dcall = nullptr;
           assert(returnIdx);
+          assert(augmentcall);
           dcall = (returnIdx.getValue() < 0)
                       ? augmentcall
                       : BuilderZ.CreateExtractValue(

enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp

@@ -2032,7 +2032,7 @@ void TypeAnalyzer::visitShuffleVectorInst(ShuffleVectorInst &I) {
       if (direction & DOWN) {
         result |= TypeTree(BaseType::Anything)
                       .Only(-1)
-                      .ShiftIndices(dl, 0, size, size * i);
+                      .ShiftIndices(dl, 0, size, newOff);
       }
     } else {
       if ((size_t)mask[i] < numFirst) {

enzyme/test/TypeAnalysis/canonwidth.ll

@@ -0,0 +1,18 @@
+; RUN: %opt < %s %loadEnzyme -print-type-analysis -type-analysis-func=err -o /dev/null | FileCheck %s
+
+define internal void @err() {
+bb:
+  %a3 = alloca <4 x i32>, align 4
+  %wide.load.le = load <4 x i32>, <4 x i32>* %a3, align 4
+  %a4 = icmp slt <4 x i32> %wide.load.le, <i32 1, i32 1, i32 1, i32 1>
+  %shift = shufflevector <4 x i1> %a4, <4 x i1> zeroinitializer, <4 x i32> <i32 undef, i32 undef, i32 3, i32 undef>
+  ret void
+}
+
+; CHECK: err - {} |
+; CHECK-NEXT: bb
+; CHECK-NEXT:   %a3 = alloca <4 x i32>, align 4: {[-1]:Pointer}
+; CHECK-NEXT:   %wide.load.le = load <4 x i32>, <4 x i32>* %a3, align 4: {}
+; CHECK-NEXT:   %a4 = icmp slt <4 x i32> %wide.load.le, <i32 1, i32 1, i32 1, i32 1>: {[-1]:Integer}
+; CHECK-NEXT:   %shift = shufflevector <4 x i1> %a4, <4 x i1> zeroinitializer, <4 x i32> <i32 undef, i32 undef, i32 3, i32 undef>: {[-1]:Anything}
+; CHECK-NEXT:   ret void: {}