Separated out memory instance rewrite pass (#2561)

* Separated out memory instance rewrite pass

* try ci

* fix

* fix

* fix

Author: wsmoses

Project.toml

@@ -41,7 +41,7 @@ BFloat16s = "0.2, 0.3, 0.4, 0.5"
 CEnum = "0.4, 0.5"
 ChainRulesCore = "1"
 EnzymeCore = "0.8.13"
-Enzyme_jll = "0.0.195"
+Enzyme_jll = "0.0.196"
 GPUArraysCore = "0.1.6, 0.2"
 GPUCompiler = "1.6"
 LLVM = "6.1, 7, 8, 9"

src/absint.jl

@@ -255,8 +255,9 @@ function should_recurse(@nospecialize(typ2), @nospecialize(arg_t::LLVM.LLVMType)
     end
 end
 
-function get_base_and_offset(@nospecialize(larg::LLVM.Value); offsetAllowed::Bool = true, inttoptr::Bool = false)::Tuple{LLVM.Value, Int}
+function get_base_and_offset(@nospecialize(larg::LLVM.Value); offsetAllowed::Bool = true, inttoptr::Bool = false, inst::Union{LLVM.Instruction, Nothing} = nothing)::Tuple{LLVM.Value, Int}
     offset = 0
+    pinst = isa(larg, LLVM.Instruction) ? larg::LLVM.Instruction : inst
     while true
         if isa(larg, LLVM.ConstantExpr)
             if opcode(larg) == LLVM.API.LLVMBitCast || opcode(larg) == LLVM.API.LLVMAddrSpaceCast || opcode(larg) == LLVM.API.LLVMPtrToInt
@@ -267,6 +268,24 @@ function get_base_and_offset(@nospecialize(larg::LLVM.Value); offsetAllowed::Boo
                 larg = operands(larg)[1]
                 continue
             end
+	    if opcode(larg) == LLVM.API.LLVMGetElementPtr && pinst isa LLVM.Instruction
+		    b = LLVM.IRBuilder()
+		    position!(b, pinst)
+		    offty = LLVM.IntType(8 * sizeof(Int))
+		    offset2 = API.EnzymeComputeByteOffsetOfGEP(b, larg, offty)
+		    if isa(offset2, LLVM.ConstantInt)
+			val = convert(Int, offset2)
+			if offsetAllowed || val == 0
+			    offset += val
+			    larg = operands(larg)[1]
+			    continue
+			else
+			    break
+			end
+		    else
+			break
+		    end
+		end
         end
         if isa(larg, LLVM.BitCastInst) || isa(larg, LLVM.AddrSpaceCastInst) || isa(larg, LLVM.IntToPtrInst)
             larg = operands(larg)[1]

src/compiler/optimize.jl

@@ -488,6 +488,8 @@ function optimize!(mod::LLVM.Module, tm::LLVM.TargetMachine)
         gvn!(pm) # Extra
         LLVM.run!(pm, mod)
     end
+
+    rewrite_generic_memory!(mod)
 
     ModulePassManager() do pm
         add_library_info!(pm, triple(mod))

src/errors.jl

@@ -524,6 +524,30 @@ function julia_error(
                 return seen[cur]
             end
 
+@static if VERSION < v"1.11-"
+else   
+	if isa(cur, LLVM.LoadInst) && isa(value_type(cur), LLVM.PointerType) && LLVM.addrspace(value_type(operands(cur)[1])) == Derived
+                    larg, off = get_base_and_offset(operands(cur)[1]; inst=ncur, inttoptr=true)
+		    if isa(larg, LLVM.ConstantInt) && off == sizeof(Int)
+			ptr = reinterpret(Ptr{Cvoid}, convert(UInt, larg))
+			obj = Base.unsafe_pointer_to_objref(ptr)
+                        if obj isa Memory && obj == typeof(obj).instance
+                            return make_batched(ncur, prevbb)
+                        end
+		    end
+                end
+	if isa(cur, LLVM.ConstantExpr) && isa(value_type(cur), LLVM.PointerType) && LLVM.addrspace(value_type(cur)) == Derived
+		larg, off = get_base_and_offset(cur; inst=first(instructions(position(prevbb))), inttoptr=true)
+		if isa(larg, LLVM.ConstantInt) && (off == sizeof(Int) || off == 0)
+			ptr = reinterpret(Ptr{Cvoid}, convert(UInt, larg))
+			obj = Base.unsafe_pointer_to_objref(ptr)
+                        if obj isa Memory && obj == typeof(obj).instance
+                            return make_batched(ncur, prevbb)
+                        end
+		    end
+                end
+end
+
             legal, TT, byref = abs_typeof(cur, true)
 
             if legal

src/llvm/transforms.jl

@@ -2297,6 +2297,33 @@ function checkNoAssumeFalse(mod::LLVM.Module, shouldshow::Bool = false)
     end
 end
 
+function rewrite_generic_memory!(mod::LLVM.Module)
+@static if VERSION < v"1.11-"
+else    
+    for f in functions(mod), bb in blocks(f)
+    iter = LLVM.API.LLVMGetFirstInstruction(bb)
+    while iter != C_NULL
+        inst = LLVM.Instruction(iter)
+        iter = LLVM.API.LLVMGetNextInstruction(iter)
+        if !isa(inst, LLVM.LoadInst)
+	   continue
+	end
+	
+	if isa(operands(inst)[1], LLVM.ConstantExpr)
+                    legal2, obj = absint(inst)
+                    if legal2 && obj isa Memory && obj == typeof(obj).instance
+			b = LLVM.IRBuilder()
+			position!(b, inst)
+                       replace_uses!(inst, unsafe_to_llvm(b, obj))
+		       LLVM.API.LLVMInstructionEraseFromParent(inst)
+		       continue
+		    end
+		end
+    end
+    end
+end
+end
+
 function removeDeadArgs!(mod::LLVM.Module, tm::LLVM.TargetMachine)
     # We need to run globalopt first. This is because remove dead args will otherwise
     # take internal functions and replace their args with undef. Then on LLVM up to