unbind info

Author: wsmoses

src/absint.jl

@@ -5,6 +5,14 @@
 
 const JL_MAX_TAGS = 64 # see `enum jl_small_typeof_tags` in julia.h
 
+function unbind(@nospecialize(val))
+   if val isa Core.Binding
+       return val.value
+   else
+       return val
+   end
+end
+
 function absint(@nospecialize(arg::LLVM.Value), partial::Bool = false, istracked::Bool=false, typetag::Bool=false)::Tuple{Bool, Any}
     if (value_type(arg) == LLVM.PointerType(LLVM.StructType(LLVMType[]), Tracked)) || (value_type(arg) == LLVM.PointerType(LLVM.StructType(LLVMType[]), Derived)) || istracked
         ce, _ = get_base_and_offset(arg; offsetAllowed = false, inttoptr = true)
@@ -455,12 +463,14 @@ function abs_typeof(
                 nm == "jl_gc_alloc_typed" ||
                 nm == "ijl_gc_alloc_typed"
             vals = absint(operands(arg)[3], partial, false, #=typetag=#true)
+	    @assert !(vals[2] isa Core.Binding)
             return (vals[1], vals[2], vals[1] ? GPUCompiler.BITS_REF : nothing)
         end
         # Type tag is arg 3
         if nm == "jl_alloc_genericmemory_unchecked" ||
 		nm == "ijl_alloc_genericmemory_unchecked"
 	    vals = absint(operands(arg)[3], partial, true, #=typetag=#true)
+	    @assert !(vals[2] isa Core.Binding)
             return (vals[1], vals[2], vals[1] ? GPUCompiler.MUT_REF : nothing)
         end
         # Type tag is arg 1
@@ -475,11 +485,13 @@ function abs_typeof(
                 nm == "jl_alloc_genericmemory" ||
                 nm == "ijl_alloc_genericmemory"
             vals = absint(operands(arg)[1], partial, false, #=typetag=#true)
+	    @assert !(vals[2] isa Core.Binding)
             return (vals[1], vals[2], vals[1] ? GPUCompiler.MUT_REF : nothing)
         end
 
         if nm == "jl_new_structt" || nm == "ijl_new_structt"
             vals = absint(operands(arg)[1], partial, false, #=typetag=#true)
+	    @assert !(vals[2] isa Core.Binding)
             return (vals[1], vals[2], vals[1] ? GPUCompiler.MUT_REF : nothing)
         end
 
@@ -498,6 +510,7 @@ function abs_typeof(
             if nm == "jl_new_structv" || nm == "ijl_new_structv"
                 @assert index == 2
                 vals = absint(operands(arg)[index], partial, false, #=typetag=#true)
+	    	@assert !(vals[2] isa Core.Binding)
                 return (vals[1], vals[2], vals[1] ? GPUCompiler.MUT_REF : nothing)
             end
 
@@ -531,9 +544,11 @@ function abs_typeof(
             if nm == "jl_f__apply_iterate" || nm == "ijl_f__apply_iterate"
                 index += 1
                 legal, iterfn = absint(operands(arg)[index])
+	    	iterfn = unbind(iterfn)
                 index += 1
                 if legal && iterfn == Base.iterate
                     legal0, combfn = absint(operands(arg)[index])
+		    combfn = unbind(combfn)
                     index += 1
                     if legal0 && combfn == Core.apply_type && partial
                         return (true, Type, GPUCompiler.BITS_REF)
@@ -871,6 +886,7 @@ function abs_typeof(
 
     legal, val = absint(arg, partial)
     if legal
+	val = unbind(val)
         return (true, Core.Typeof(val), GPUCompiler.BITS_REF)
     end
     return (false, nothing, nothing)

src/compiler.jl

@@ -1694,8 +1694,10 @@ function shadow_alloc_rewrite(V::LLVM.API.LLVMValueRef, gutils::API.EnzymeGradie
 				index += 1
 				found = Any[]
 				legal, Ty = absint(operands(arg)[index], partial)
+				Ty = unbind(Ty)
 				if legal && Ty == NTuple
 				   legal, Ty = absint(operands(arg)[index+2])
+				   Ty = unbind(Ty)
 				   if legal
 					# count should represent {the total size in bytes, the aligned size of each element}
 					B = LLVM.IRBuilder()
@@ -5400,6 +5402,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT
 @static if VERSION < v"1.11-"
 else    
                     legal2, obj = absint(inst)
+		    obj = unbind(obj)
 		    if legal2 && is_memory_instance(obj)
                         metadata(inst)["nonnull"] = MDNode(LLVM.Metadata[])
                     end
@@ -5629,6 +5632,7 @@ end
                         string(cur)
                     slegal, foundv = absint(cur)
                     if slegal
+		    	foundv = unbind(foundv)
                         resstr *= "of type " * string(foundv)
                     end
                     emit_error(builder, user, resstr, EnzymeMutabilityException)
@@ -6464,6 +6468,7 @@ const DumpLLVMCall = Ref(false)
         end
         reinsert_gcmarker!(llvm_f)
 
+	Enzyme.Compiler.JIT.prepare!(mod)
 	if DumpLLVMCall[]
 	   API.EnzymeDumpModuleRef(mod.ref)
 	end

src/compiler/validation.jl

@@ -278,7 +278,7 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr
 			b = IRBuilder()
 			position!(b, inst)
 			ccall(:jl_, Cvoid, (Any,), (gname, obj, obj0, ptr, string(addr)))
-			newf = unsafe_to_llvm(b, obj; insert_name_if_not_exists=gname) 
+			newf = unsafe_to_llvm(b, obj0; insert_name_if_not_exists=gname) 
 			replace_uses!(inst, newf)
 			LLVM.API.LLVMInstructionEraseFromParent(inst)
 			continue
@@ -835,7 +835,7 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr
 	    if isa(flib, LLVM.ConstantExpr)
 		legal, flib2 = absint(flib)
 		if legal
-		   flib = flib2
+		    flib = unbind(flib2)
 		end
             end
             if isa(flib, GlobalRef) && isdefined(flib.mod, flib.name)
@@ -986,6 +986,7 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr
                 iteroff = 2
 
                 legal, iterlib = absint(operands(inst)[iteroff+1])
+		iterlib = unbind(iterlib)
                 if legal && iterlib == Base.iterate
                     legal, GT, byref = abs_typeof(operands(inst)[4+1], true)
                     funcoff = 3
@@ -1075,6 +1076,7 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr
                         push!(tys, typ)
                     end
                     legal, flib = absint(operands(inst)[offset+1])
+		    flib = unbind(flib)
                     if legal && isa(flib, Core.MethodInstance)
                         if !Base.isvarargtype(flib.specTypes.parameters[end])
                             @assert length(tys) == length(flib.specTypes.parameters)
@@ -1229,6 +1231,7 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr
                     push!(tys, typ)
                 end
                 legal, flib = absint(operands(inst)[offset+1])
+		flib = unbind(flib)
                 if legal && isa(flib, Core.MethodInstance)
                     if !Base.isvarargtype(flib.specTypes.parameters[end])
                         if length(tys) != length(flib.specTypes.parameters)

src/errors.jl

@@ -1070,8 +1070,13 @@ function julia_error(
             end
 	    legal, obj = absint(val)
 	    if legal
-	        println(io, "\nValue of type: ", Core.Typeof(val))
-		println(io ,  " of value    : ", val)
+		obj0 = obj
+		obj = unbind(obj)
+	        println(io, "\nValue of type: ", Core.Typeof(obj))
+		println(io ,  " of value    : ", obj)
+		if obj0 isa Core.Binding
+		println(io ,  " binding     : ", obj0)	    
+		end
 		println(io)
 	    end
 	    if !isa(val, LLVM.Argument) && !isa(val, LLVM.GlobalVariable) 
@@ -1297,9 +1302,10 @@ function julia_error(
                 end
 
                 legal2, obj = absint(cur)
-
+		obj0 = obj
                 # Only do so for the immediate operand/etc to a phi, since otherwise we will make multiple
                 if legal2
+		   obj = unbind(obj)
 		   if is_memory_instance(obj)
 			return make_batched(ncur, prevbb)
 		   end
@@ -1337,9 +1343,6 @@ function julia_error(
                     end
                     end
 
-		    if is_memory_instance(obj)
-                        return make_batched(ncur, prevbb)
-                    end
                 end
 
 @static if VERSION < v"1.11-"
@@ -1348,6 +1351,7 @@ else
                     larg, off = get_base_and_offset(operands(cur)[1])
                     if isa(larg, LLVM.LoadInst)
                         legal2, obj = absint(larg)
+			obj = unbind(obj)
 			if legal2 && is_memory_instance(obj)
                             return make_batched(ncur, prevbb)
                         end
@@ -1356,7 +1360,10 @@ else
 end
 
                 badval = if legal2
-                    string(obj) * " of type" * " " * string(TT)
+                    sv = string(obj) * " of type" * " " * string(TT)
+		    if obj0 isa Core.Binding
+			sv = sv *" binded at "*string(obj0)
+		    end
                 else
                     "Unknown object of type" * " " * string(TT)
                 end

src/jlrt.jl

@@ -454,7 +454,7 @@ function emit_apply_type!(B::LLVM.IRBuilder, @nospecialize(Ty::Type), args::Vect
     for arg in args
         slegal, foundv = absint(arg)
         if slegal
-            push!(found, foundv)
+	    push!(found, unbind(foundv))
         else
             legal = false
             break
@@ -509,7 +509,7 @@ function emit_tuple!(B::LLVM.IRBuilder, args::Vector{LLVM.Value})::LLVM.Value
     for arg in args
         slegal, foundv = absint(arg)
         if slegal
-            push!(found, foundv)
+	    push!(found, unbind(foundv))
         else
             legal = false
             break
@@ -871,6 +871,7 @@ function emit_layout_of_type!(B::LLVM.IRBuilder, @nospecialize(ty::LLVM.Value))
 	ls = get_layout_struct()
 	lptr = LLVM.PointerType(ls, 10)
 	if legal
+		JTy = unbind(JTy)
 		return LLVM.const_inttoptr(LLVM.ConstantInt(Base.reinterpret(UInt, JTy.layout)), lptr)
 	end
 	@assert !isa(ty, LLVM.ConstantExpr)
@@ -889,6 +890,7 @@ end
 function emit_type_layout_elsz!(B::LLVM.IRBuilder, @nospecialize(ty::LLVM.Value))
 	legal, JTy = absint(ty)
 	if legal
+	    JTy = unbind(JTy)
 	    @assert JTy isa Type
 	    res = Compiler.datatype_layoutsize(JTy)
 	    return LLVM.ConstantInt(res)

src/llvm/transforms.jl

@@ -442,7 +442,7 @@ function memcpy_alloca_to_loadstore(mod::LLVM.Module)
 	@static if VERSION < v"1.11-"
 	else    
 			    legal2, obj = absint(src)
-			    if legal2 && is_memory_instance(obj) 
+			    if legal2 && is_memory_instance(unbind(obj)) 
 				metadata(src)["nonnull"] = MDNode(LLVM.Metadata[])
 			    end
 	end

src/rules/jitrules.jl

@@ -2361,7 +2361,9 @@ function common_apply_iterate_fwd(offset, B, orig, gutils, normalR, shadowR)
     end
 
     v, isiter = absint(operands(orig)[offset+1])
+    isiter = unbind(isiter)
     v2, istup = absint(operands(orig)[offset+2])
+    istup = unbind(istup)
 
     width = get_width(gutils)
 
@@ -2513,6 +2515,8 @@ function common_apply_iterate_augfwd(offset, B, orig, gutils, normalR, shadowR,
 
     v, isiter = absint(operands(orig)[offset+1])
     v2, istup = absint(operands(orig)[offset+2])
+    isiter = unbind(isiter)
+    istup = unbind(istup)
 
     width = get_width(gutils)