[BUG] Improper hash calculation with deferred static evaluation

[Kenny-Vilella]

Bug Description

It comes from MjWarp, in some special edge case we launch the wrong kernel because the hash assigned to the kernel does not include critical information from deferred static evaluation.

Simple repro:

import warp as wp

wp.init()

_KERNEL_CACHE = {}

def cache_kernel(func):
    def wrapper(*args):
        # Hash lists by converting to tuple
        key = hash(tuple(tuple(a) if isinstance(a, list) else a for a in args))
        if key not in _KERNEL_CACHE:
            _KERNEL_CACHE[key] = func(*args)
        return _KERNEL_CACHE[key]
    return wrapper


@cache_kernel
def make_kernel(collision_types):
    """Kernel factory - the wp.static(collision_types[i]) references loop var 'i'."""

    def my_kernel(result: wp.array(dtype=int)):
        tid = wp.tid()
        if tid == 0:
            # Loop over collision types - 'i' is a loop variable
            for i in range(wp.static(len(collision_types))):
                # wp.static references 'i', so it's DEFERRED (not in initial hash)
                result[i] = wp.static(collision_types[i])

    return wp.kernel(module="unique", enable_backward=False)(my_kernel)

# Test 1: Create kernel with types [100, 200]
types_1 = [100, 200]
result_1 = wp.zeros(2, dtype=int)
wp.launch(make_kernel(types_1), dim=1, inputs=[], outputs=[result_1])
wp.synchronize()
print(f"  Result: {result_1.numpy().tolist()}")
print(f"  Expected: {types_1}")

# Test 2: Create kernel with types [999, 888] - DIFFERENT types, SAME length!
types_2 = [999, 888]
result_2 = wp.zeros(2, dtype=int)
wp.launch(make_kernel(types_2), dim=1, inputs=[], outputs=[result_2])
wp.synchronize()
print(f"  Result: {result_2.numpy().tolist()}")
print(f"  Expected: {types_2}")


# Test 3: Create kernel with types [999, 888, 100] - DIFFERENT length
types_3 = [999, 888, 100]
result_3 = wp.zeros(3, dtype=int)
wp.launch(make_kernel(types_3), dim=1, inputs=[], outputs=[result_3])
wp.synchronize()
print(f"  Result: {result_3.numpy().tolist()}")
print(f"  Expected: {types_3}")

System Information

No response

[Kenny-Vilella]

@shi-eric The issue has not been totally fixed, apparently it just use the last iteration of the loop instead of all values.

Another repro:

import warp as wp

wp.init()


def make_kernel(collision_types):
    """Kernel factory — each combination of collision_types should produce a unique kernel."""

    def my_kernel(result: wp.array(dtype=int)):
        tid = wp.tid()
        if tid == 0:
            for i in range(wp.static(len(collision_types))):
                result[i] = wp.static(collision_types[i])

    return wp.kernel(module="unique", enable_backward=False)(my_kernel)


# --- Test 1: [100, 999] — last element is 999 ---
k1 = make_kernel([100, 999])
r1 = wp.zeros(2, dtype=int)
wp.launch(k1, dim=1, inputs=[], outputs=[r1])
print(f"  Result:   {r1.numpy().tolist()}")
print(f"  Expected: [100, 999]")

# --- Test 2: [200, 999] — DIFFERENT first element, SAME last element ---
k2 = make_kernel([200, 999])
r2 = wp.zeros(2, dtype=int)
wp.launch(k2, dim=1, inputs=[], outputs=[r2])
print(f"  Result:   {r2.numpy().tolist()}")
print(f"  Expected: [200, 999]")

# --- Test 3: [300, 999] — yet another first element, same last ---
k3 = make_kernel([300, 999])
r3 = wp.zeros(2, dtype=int)
wp.launch(k3, dim=1, inputs=[], outputs=[r3])
print(f"  Result:   {r3.numpy().tolist()}")
print(f"  Expected: [300, 999]")

# --- Test 4: [100, 888] — same first element as Test 1, DIFFERENT last ---
k4 = make_kernel([100, 888])
r4 = wp.zeros(2, dtype=int)
wp.launch(k4, dim=1, inputs=[], outputs=[r4])
print(f"  Result:   {r4.numpy().tolist()}")
print(f"  Expected: [100, 888]")

[shi-eric]

Added 1ff6807 to fix the above issue, but we know there's another multi-process hashing-related bug floating that's unresolved and being investigated

[Kenny-Vilella]

The last bug seems to be a MjWarp issue, PR in here.

[Kenny-Vilella]

PR is in, everything should be fixed now.
@shi-eric Do you confirm that we are all good?

[shi-eric]

PR is in, everything should be fixed now. @shi-eric Do you confirm that we are all good?

I think it's fixed! There were two jobs that ran today using an MJWarp commit that included your change and both passed, compared to before the fix when the failure rate was 100% or close to it. Thank you!