diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp index b3c66aa4893..996bbe357f8 100644 --- a/enzyme/Enzyme/GradientUtils.cpp +++ b/enzyme/Enzyme/GradientUtils.cpp @@ -5441,6 +5441,65 @@ static bool allNullOrUndef(Value *C, const DataLayout &dl, TypeTree TT) { return false; } +/// Return a copy of the scalar constant `C` with every byte that `TT` types as +/// a float (or as something other than an integer or a pointer) replaced by +/// zero, keeping the remaining bytes as they were. This is the constant-folded +/// equivalent of the alloca/store/store-zero/load sequence invertPointerM emits +/// for a value that is only partially float, and it lets the shadow of a +/// constant stay a constant. Returns null if `C` cannot be reinterpreted +/// bitwise, in which case the caller falls back to emitting that sequence. +static Constant *nullFloatBytesOfConstant(Constant *C, const TypeTree &TT, + const DataLayout &DL) { + auto *Ty = C->getType(); + + APInt bytes(8, 0); + if (auto CI = dyn_cast(C)) + bytes = CI->getValue(); + else if (auto CFP = dyn_cast(C)) + bytes = CFP->getValueAPF().bitcastToAPInt(); + else + return nullptr; + + // Only a whole number of bytes can be masked byte-wise, and only a type + // whose value fills its in-memory footprint can be rebuilt from those bytes + // (this rules out i1, i24, x86_fp80, ...). + unsigned bits = bytes.getBitWidth(); + size_t size = bits / 8; + size_t storeSize = (DL.getTypeStoreSizeInBits(Ty) + 7) / 8; + if (size == 0 || size * 8 != bits || size != storeSize) + return nullptr; + + APInt mask(bits, 0); + for (size_t i = 0; i < size;) { + auto CT = TT[{(int)i}]; + size_t chunk = 1; + if (CT == BaseType::Pointer) { + // A pointer is carried through unchanged, like an integer. + i += DL.getPointerSize(0); + continue; + } else if (auto flt = CT.isFloat()) { + chunk = (DL.getTypeSizeInBits(flt) + 7) / 8; + } else if (CT == BaseType::Integer) { + i++; + continue; + } + // Float, unknown, and anything all get a zero derivative. Byte `j` of the + // memory image is bits [8j, 8j+8) of the APInt on a little endian target + // and bits [8(size-1-j), 8(size-j)) on a big endian one. + size_t end = std::min(i + chunk, size); + size_t loByte = DL.isLittleEndian() ? i : size - end; + size_t hiByte = DL.isLittleEndian() ? end : size - i; + mask |= APInt::getBitsSet(bits, loByte * 8, hiByte * 8); + i += chunk; + } + bytes &= ~mask; + + if (Ty->isIntegerTy()) + return ConstantInt::get(Ty, bytes); + return ConstantFP::get(Ty->getContext(), + APFloat(Ty->getFltSemantics(), bytes)); +} + Value *GradientUtils::invertPointerM(Value *const oval, IRBuilder<> &BuilderM) { return invertPointerM(oval, BuilderM, TR.query(oval)); } @@ -5462,6 +5521,19 @@ Value *GradientUtils::invertPointerM(Value *const oval, IRBuilder<> &BuilderM, #endif auto &DL = oldFunc->getParent()->getDataLayout(); + + // A constant can arrive here with no type information at all: the + // GlobalVariable case inverts an initializer with the tree it holds for the + // pointer, which says nothing about the pointee's contents, and the + // aggregate cases below then slice that empty tree per element. Every byte + // stays Unknown, which anyFloat() reads as "may be a float" while allFloat() + // reads as "not all float" -- the one combination that sends a compile-time + // constant down the runtime partially-float path. Type analysis already + // knows the answer for constants, so ask it rather than guessing from + // nothing. + if (!TT.isKnown() && isa(oval)) + TT = TR.query(oval); + if (isa(oval) || isa(oval) || isa(oval) || isa(oval) || isa(oval)) { @@ -5570,7 +5642,18 @@ Value *GradientUtils::invertPointerM(Value *const oval, IRBuilder<> &BuilderM, if (TT.anyFloat(oval, DL)) { if (TT.allFloat(oval, DL, /*anythingIsFloat*/ true)) return Constant::getNullValue(getShadowType(oval->getType())); - else { + // The shadow of a constant has to stay a constant: this value may be an + // element of a constant aggregate whose inversion above feeds + // ConstantArray/ConstantStruct/ConstantVector::get. Zero the float bytes + // by folding rather than with the alloca/store/load below, which would + // hand back a LoadInst and trip the cast in those callers. + Constant *folded = nullptr; + if (auto C = dyn_cast(oval)) + folded = nullFloatBytesOfConstant(C, TT, DL); + if (folded) { + auto rule = [&folded]() { return folded; }; + return applyChainRule(oval->getType(), BuilderM, rule); + } else { IRBuilder<> bb(inversionAllocs); if (auto arg = dyn_cast(oval)) { arg = getNewFromOriginal(arg); diff --git a/enzyme/test/Enzyme/ReverseMode/partialfloatglobal.ll b/enzyme/test/Enzyme/ReverseMode/partialfloatglobal.ll new file mode 100644 index 00000000000..6336c28c0a1 --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/partialfloatglobal.ll @@ -0,0 +1,32 @@ +; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme-preopt=false -enzyme -S | FileCheck %s; fi +; RUN: %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="enzyme" -S | FileCheck %s + +; @tab holds two packed doubles as raw bytes, so its ConstantDataArray elements +; are i8. A single byte overlaps a double without covering all of it, so +; inverting that element used to take the generic partially-float path in +; invertPointerM, which emits an alloca/store/load and hands back a LoadInst. +; The ConstantDataArray case then did cast on it and asserted with +; "cast() argument of incompatible type!". + +target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +@tab = global [16 x i8] c"UUUUUU\D5?\00\00\00\00\00\00\D0?" +@enzyme_const = external global ptr +@enzyme_dup = external global ptr + +define void @f(ptr %a, i64 %n, ptr %b, i64 %m, i32 %c) { + store i32 0, ptr @tab, align 4 + ret void +} + +define {} @entry() { + %r = tail call {} (...) @__enzyme_autodiff(ptr @f, ptr @enzyme_dup, ptr null, ptr null, ptr @enzyme_const, i64 0, ptr @enzyme_dup, ptr null, ptr null, ptr @enzyme_const, i64 0, ptr @enzyme_const, i32 0) + ret {} %r +} + +declare {} @__enzyme_autodiff(...) + +; Every byte of @tab is part of a double, so the whole shadow folds to zero and +; stays a constant initializer. +; CHECK: @tab_shadow = global [16 x i8] zeroinitializer