Unexpected result from LoopRecord #1137
-
SummaryHi, i want to implement a resampling integrator that selects 1 path from N traced paths. The result i produced seem to be different with and without System configurationSystem information: OS: Windows-10 Dr.Jit: 0.4.4 (DEBUG) DescriptionMy implementation is laid inside Here is my code: import mitsuba as mi
import drjit as dr
mi.set_variant('cuda_ad_rgb')
dr.set_flag(dr.JitFlag.LoopRecord, False)
def mis_weight(pdf_a, pdf_b):
a2 = dr.sqr(pdf_a)
b2 = dr.sqr(pdf_b)
w = a2 / (a2 + b2)
return dr.detach(dr.select(dr.isfinite(w), w, 0))
def luminance(L: mi.Color3f) -> mi.Float:
return dr.dot(L, mi.Vector3f(0.212671, 0.715160, 0.072169))
class RestirSample:
p_hat: mi.Color3f # target function
def __init__(self):
self.p_hat = mi.Color3f(0.0)
def assign(self, s: 'RestirSample', active: mi.Bool):
self.p_hat[active] = mi.Color3f(s.p_hat)
DRJIT_STRUCT = {'p_hat': mi.Color3f}
class RestirReservoir:
z: RestirSample
w_sum: mi.Float # accumulated sample picking-weight
W: mi.Float # weight
c: mi.UInt32 # confidence weight
def __init__(self):
self.z = RestirSample()
self.w_sum = mi.Float(0)
self.W = mi.Float(0)
self.c = mi.UInt32(0)
def update(self, sample: mi.Float, snew: RestirSample, wnew: mi.Float, active: mi.Bool):
# update the reservoir
active = mi.Bool(active)
if dr.shape(active)[-1] == 1:
dr.make_opaque(active)
self.w_sum += dr.select(active, wnew, 0)
self.c += dr.select(active, 1, 0)
active_selected = sample < (wnew / (self.w_sum + 1e-9))
self.z.assign(snew, active & active_selected)
def update_weight(self):
self.W = dr.select(luminance(self.z.p_hat) > 1e-6 , self.w_sum / luminance(self.z.p_hat), 0)
def merge(self, sample: mi.Float, r: "RestirReservoir", wnew, active: mi.Bool):
active = mi.Bool(active)
c0 = mi.UInt(
self.c
) # We have to construct a new UInt so we don't overwrite `self.M`
self.update(sample, r.z, wnew, active)
self.c = dr.select(active, c0 + r.c, c0)
DRJIT_STRUCT = {'z': RestirSample, 'w_sum': mi.Float, 'W': mi.Float, 'c': mi.UInt32}
class ReSamplingIntegrator(mi.SamplingIntegrator):
def __init__(self, props=mi.Properties()):
super().__init__(props)
self.max_depth = props.get('max_depth', -1)
self.rr_depth = props.get('rr_depth', 100)
self.hide_emitters = props.get('hide_emitters', False)
self.n_init_paths = props.get("n_init_paths", 16)
self.use_nee = True
self.film_size: None | mi.Vector2u = None
def render(self, scene: mi.Scene, sensor: mi.Sensor, seed: int = 0, spp: int = 1,develop: bool = True,evaluate: bool = True,):
film = sensor.film()
film_size = film.crop_size()
if self.film_size is None:
self.film_size = film_size
wavefront_size = film_size.x * film_size.y * spp
self.wavefront_size = wavefront_size
sampler = sensor.sampler()
sampler.set_sample_count(spp)
sampler.set_samples_per_wavefront(spp)
sampler.seed(0, wavefront_size)
idx = dr.arange(mi.UInt, wavefront_size)
pos = mi.Vector2u()
pos.x = idx // spp % film_size.x
pos.y = idx // film_size.x // spp
self.layer = idx % spp
self.spp = spp
sample_pos = (mi.Point2f(pos) + sampler.next_2d()) / mi.Point2f(
film.crop_size()
)
res = self.sample_initial(scene, sampler, sensor, sample_pos, self.n_init_paths)
rad = res
film.prepare(self.aov_names())
block: mi.ImageBlock = film.create_block()
aovs = [rad.x, rad.y, rad.z, mi.Float(1)]
block.put(pos, aovs)
film.put_block(block)
img = film.develop()
dr.eval(img)
return img
def sample_initial(self, scene: mi.Scene, sampler: mi.Sampler, sensor: mi.Sensor, sample_pos: mi.Point2f, m_random_walk: mi.UInt32) -> RestirSample:
init_reservoir = RestirReservoir()
i = mi.UInt32(0)
active_init = mi.Bool(True)
loop_init = mi.Loop(name=f"Init)",
state=lambda: (i, active_init, sampler, init_reservoir))
while loop_init(active_init):
active_init &= i < m_random_walk
restir_sample = RestirSample()
reservoir = RestirReservoir()
ray, ray_weight = sensor.sample_ray(0.0, 0.0, sample_pos, mi.Point2f(0.5))
ray = mi.Ray3f(ray)
depth = mi.UInt32(0) # Depth of current vertex
L = mi.Spectrum(0) # Radiance accumulator
throughput = mi.Spectrum(1) # Path throughput weight
η = mi.Float(1) # Index of refraction
active = mi.Bool(True)
active &= active_init
si = dr.zeros(mi.SurfaceInteraction3f)
needs_intersection = mi.Bool(True)
last_scatter_event = dr.zeros(mi.Interaction3f)
last_scatter_direction_pdf = mi.Float(1.0)
depth = mi.UInt32(0)
valid_ray = mi.Bool(False)
specular_chain = mi.Bool(True)
loop = mi.Loop(name=f"Subsurface Scattering RIS)",
state=lambda: (sampler, active, depth, ray, si,
throughput, L, needs_intersection,
last_scatter_event, specular_chain, η,
last_scatter_direction_pdf, valid_ray,
restir_sample, reservoir))
while loop(active):
# if still exist at least 1 throughput not equal to zeros
# keep the loop going
active &= dr.any(dr.neq(throughput, 0.0))
active &= depth < self.max_depth
#--------------------- Perform russian roulette --------------------
q = dr.minimum(dr.max(throughput) * dr.sqr(η), 0.99)
perform_rr = (depth > self.rr_depth) # start rr from this depth level
active &= (sampler.next_1d() < q) | ~perform_rr
throughput[perform_rr] = throughput * dr.rcp(q)
active_surface = active
#--------------------- Surface Interactions --------------------
intersect = active_surface & needs_intersection
si[intersect] = scene.ray_intersect(ray, intersect)
# ----------------- Intersection with emitters -----------------
ray_from_camera = active_surface & dr.eq(depth, 0)
count_direct = ray_from_camera | specular_chain
emitter = si.emitter(scene)
active_e = active_surface & dr.neq(emitter, None) & ~(dr.eq(depth, 0) & self.hide_emitters)
# Get the PDF of sampling this emitter using next event estimation
ds = mi.DirectionSample3f(scene, si, last_scatter_event)
if self.use_nee:
emitter_pdf = scene.pdf_emitter_direction(last_scatter_event, ds, active_e)
else:
emitter_pdf = 0.0
emitted = emitter.eval(si, active_e)
contrib = dr.select(count_direct, throughput * emitted,
throughput * mis_weight(last_scatter_direction_pdf, emitter_pdf) * emitted)
L[active_e] += dr.detach(contrib)
restir_sample.p_hat[active_e] = contrib
reservoir.update(sampler.next_1d(), restir_sample, luminance(contrib), active_e)
active_surface &= si.is_valid()
ctx = mi.BSDFContext()
bsdf = si.bsdf(ray) #drjit::DiffArray<drjit::LLVMArray<mitsuba.BSDF const*> >
# ---------------------- Emitter sampling ----------------------
if self.use_nee:
active_e_surface = active_surface & mi.has_flag(bsdf.flags(), mi.BSDFFlags.Smooth) & (depth + 1 < self.max_depth)
active_e = active_e_surface
emitted, ds = self.sample_emitter(si, active_e_surface,
scene, sampler, active_e)
# Query the BSDF for that emitter-sampled direction
bsdf_val, bsdf_pdf = bsdf.eval_pdf(ctx, si, dr.detach(si.to_local(ds.d)), active_e_surface)
nee_weight = bsdf_val
nee_directional_pdf = dr.select(ds.delta, 0.0, bsdf_pdf)
contrib = throughput * nee_weight * mis_weight(ds.pdf, nee_directional_pdf) * emitted
L[active_e] += dr.detach(contrib)
restir_sample.p_hat = contrib
reservoir.update(sampler.next_1d(), restir_sample, luminance(contrib), active_e)
# ------------------------ BSDF sampling -----------------------
bs, bsdf_weight = bsdf.sample(ctx, si,
sampler.next_1d(),
sampler.next_2d(),
active_surface)
active_surface &= bs.pdf > 0
bsdf_eval = bsdf.eval(ctx, si, bs.wo, active_surface)
throughput[active_surface] *= bsdf_weight
η[active_surface] *= bs.eta
bsdf_ray = si.spawn_ray(dr.detach(si.to_world(bs.wo)))
ray[active_surface] = bsdf_ray
needs_intersection |= active_surface
non_null_bsdf = active_surface & ~mi.has_flag(bs.sampled_type, mi.BSDFFlags.Null)
depth[non_null_bsdf] += 1
last_scatter_event[non_null_bsdf] = si
last_scatter_direction_pdf[non_null_bsdf] = bs.pdf
valid_ray |= non_null_bsdf
specular_chain |= non_null_bsdf & mi.has_flag(bs.sampled_type, mi.BSDFFlags.Delta)
specular_chain &= ~(active_surface & mi.has_flag(bs.sampled_type, mi.BSDFFlags.Smooth))
active &= active_surface
i += mi.UInt32(1)
reservoir.update_weight()
w = (1 / m_random_walk) * luminance(reservoir.z.p_hat) * reservoir.W
init_reservoir.merge(sampler.next_1d(), reservoir, w, active_init)
active_init &= i < m_random_walk
init_reservoir.update_weight()
dr.eval(init_reservoir)
return init_reservoir.z.p_hat * init_reservoir.W
def sample_emitter(self, si, active_surface, scene, sampler,
active, adj_emitted=None, δL=None, mode=None):
active = mi.Bool(active)
ref_interaction = dr.zeros(mi.Interaction3f)
ref_interaction[active_surface] = si
ds, emitter_val = scene.sample_emitter_direction(ref_interaction,
sampler.next_2d(active),
False, active)
ds = dr.detach(ds)
invalid = dr.eq(ds.pdf, 0.0)
emitter_val[invalid] = 0.0
active &= ~invalid
return emitter_val, ds
mi.register_integrator("resampler", lambda props: ReSamplingIntegrator(props))
scene = mi.load_file('./scenes/cornell-box/cornell-box.xml')
img = mi.render(scene, spp = 1)
mi.util.write_bitmap("./output/test/my_render.png", img)Here my scene description file: I want to ask how to overcome this problem since recording loop save me a lot of overhead computation cost. |
Beta Was this translation helpful? Give feedback.
Answered by
NhatMinh2208
Apr 18, 2024
Replies: 2 comments 2 replies
-
|
I also notice that if i add new attribute (w) to my custom DrJit struct here: (This is the only few lines i add, i do not touch on any other part of the program) and the picture also change, with or without LoopRecord on (which is not supposed to). |
Beta Was this translation helpful? Give feedback.
0 replies
-
|
Thanks for the complete reproducer, however I can't seem to reproduce this. |
Beta Was this translation helpful? Give feedback.
2 replies
-
|
Thanks @njroussel for replying my question. |
Beta Was this translation helpful? Give feedback.
-
|
I just checked it, the loop seem to stuck forever cause i forgot |
Beta Was this translation helpful? Give feedback.
Answer selected by
NhatMinh2208
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
I just checked it, the loop seem to stuck forever cause i forgot
active &= depth < self.max_depthThank you @njroussel i will double check the result again.