Advice in debugging memory leak with Enzyme

[elite-sheep]

Enzyme: 0.0.29
LLVM: 12.1

Hi, I am building a differentiable renderer with Enzyme. However, I spotted a very bad memory leak issue with Enzyme when I differentiate the following function:

    void d_pixelColor(const Scene &scene, Scene &d_scene,
                      const Array2i &pixel_idx, RndSampler *sampler,
                      int n_sample, Spectrum d_res)
    {
        Spectrum res;
        for (int i = 0; i < n_sample; i++)
        {
            [[maybe_unused]] Spectrum _d_res = d_res / n_sample;
            __enzyme_autodiff((void *)Li,
                              enzyme_dup, &scene, &d_scene,
                              enzyme_const, &pixel_idx,
                              enzyme_const, sampler,
                              enzyme_const, n_sample,
                              enzyme_dupnoneed, &res, &_d_res);
        }
    }

And the function to be differentiated is:

Spectrum Direct::Li(const Scene &scene, RndSampler *sampler, const Array2i &pixel_idx) const
{
    Spectrum ret = Spectrum::Zero();
    ::Li(scene, pixel_idx, sampler, 1, ret);
    return ret;
}

    Spectrum radiance(const Scene &scene, RndSampler *sampler, Ray &ray,
                      const Array2i &pixel_idx)
    {
        Intersection its;
        Float throughput = 1.;
        Spectrum ret = Spectrum::Zero();
        // Perform the first intersection
        scene.rayIntersect(ray, true, its, IntersectionMode::EMaterial);
        if (its.isValid())
        {
            throughput *= scene.camera.evalFilter(pixel_idx.x(), pixel_idx.y(), its);
            if (its.isEmitter())
                ret += its.Le(-ray.dir) * its.J;
            // Direct illumination
            DirectSamplingRecord dRec(its);
            auto value = scene.sampleEmitterDirect(sampler->next2D(), dRec);
            if (!value.isZero())
            {
#ifdef MIS
                auto bsdf_val = its.evalBSDF(its.toLocal(dRec.dir));
                Float bsdf_pdf = its.pdfBSDF(its.toLocal(dRec.dir)) * dRec.G;
                bsdf_val = mueller::to_world_mueller(its, bsdf_val, -its.toLocal(dRec.dir), its.wi);

                Float mis_weight = square(dRec.pdf) / (square(dRec.pdf) + square(bsdf_pdf));
                ret += bsdf_val * value* mis_weight;
#else
                auto bsdf_val = its.evalBSDF(its.toLocal(dRec.dir));
                bsdf_val = mueller::to_world_mueller(its, bsdf_val, -its.toLocal(dRec.dir), its.wi);
                ret += bsdf_val * value * its.J;
#endif
            }

#ifdef MIS
            Float bsdf_pdf, bsdf_eta;
            Vector wo;
            auto bsdf_weight = its.sampleBSDF(sampler->next3D(), wo, bsdf_pdf, bsdf_eta);
            if (!bsdf_weight.isZero())
            {
                Ray ray(its.p, its.toWorld(wo));
                Intersection its_light;
                if (scene.rayIntersect(ray, true, its_light, IntersectionMode::EMaterial))
                {
                    if (its_light.isEmitter())
                    {
                        Float G = geometric(ray.org, its_light.p, its_light.geoFrame.n);
                        bsdf_pdf *= detach(G); // area measure
                        auto le = its_light.Le(-ray.dir) * its_light.J * G;
                        if (!le.isZero())
                        {
                            bsdf_weight = its.evalBSDF(its.toLocal(its_light.p - its.p).normalized());
                            bsdf_weight = mueller::to_world_mueller(its, bsdf_weight, -wo, its.wi);
                            Float pdf_nee = scene.pdfEmitterSample(its_light);
                            Float mis_weight = square(bsdf_pdf) / (square(pdf_nee) + square(bsdf_pdf));
                            ret += bsdf_weight * le * mis_weight * its.J / bsdf_pdf;
                        }
                    }
                }
            }
#endif
        }

#if defined(PSDR_POLARIZED)
       const Camera* camera = &scene.camera;
       Matrix4x4 transform = camera->cam_to_world;
       Vector3 ray_d = ray.dir;
       Vector3 cur_basis = mueller::stokes_basis(-ray_d);
       Vector4 affine_vertical = transform * Vector4(0.0, 1.0, 0.0, 0.0);
       Vector3 vertical(affine_vertical[0], affine_vertical[1], affine_vertical[2]);
       Vector3 target_basis = (ray_d).cross(vertical);
       Spectrum rotater = mueller::rotate_stoke_basis(-ray_d, cur_basis, target_basis);
       ret = rotater * ret;
#endif

        return throughput * ret;
    }

The story is I call d_pixelColor from Python(Binded using PyBind11) but there is a very bad memory leak issue. I need some advice in debugging the memory leak with the Enzyme generated code. Can anyone please share some ideas?

[elite-sheep]

Attach the output by -enzyme-print. The first one is the one with memory leak. https://drive.google.com/file/d/1BxkbKpG_ljpQvC_gFk9fO-nWA301I0GR/view?usp=sharing

[wsmoses]

Fantastic! For the sake of testing, can you provide a zip/repo and build instructions on this.

Additionally the result of adding -mllvm -print-before=enzyme would also be appreciated.

[wsmoses]

Also if you're able to reduce the code above to a smaller case that still has the leak, that will also make it a lot easier to debug.

Regardless, I personally am attacking several paper deadlines until 4/1 so I may not get a chance to look into until then (if someone else has time beforehand free).