[BUG] CuTeDSL 4.8.0 nvpu-optimization fails without a diagnostic on a sum of 10 dynamic extents #3699

import argparse                                                                              
                                                                                                 
import cutlass                                                                               
import torch                                                                                 
from cutlass import cute                                                                     
from cutlass.cute.runtime import from_dlpack                                                                                                                                             
                                                                                                 
@cute.kernel                                                                                 
def kernel(out: cute.Tensor):                                                                
    tidx, _, _ = cute.arch.thread_idx()                                                      
    if tidx == 0:                                                                            
        out[0] = cutlass.Int32(1)                                                            
                                                                                            
@cute.jit                                                                                    
def host(out: cute.Tensor, tensors: tuple):                                                  
    total = sum(cute.size(t, mode=[0]) for t in tensors)                                     
    kernel(out).launch(grid=(total, 1, 1), block=(32, 1, 1))                                 
                                                                                                                                                                                               
def main():                                                                                  
    parser = argparse.ArgumentParser()                                                       
    parser.add_argument("--n", type=int, default=10, help="number of summed sizes")          
    args = parser.parse_args()                                                               
                                                                                           
    out = from_dlpack(torch.zeros(1, dtype=torch.int32, device="cuda"))                      
    tensors = tuple(                                                                         
        from_dlpack(torch.zeros(8, dtype=torch.int32, device="cuda")).mark_layout_dynamic()  
        for _ in range(args.n)                                                               
    )                                                                                        
    cute.compile(host, out, tensors)                                                         
    print(f"n={args.n}: compiled OK")                                                        
                                                                                                 
                                                                                                 
if __name__ == "__main__":                                                                   
    main()

"""Save a failed compilation's partially rewritten IR and retry the full pipeline."""

import json
import os
from pathlib import Path

# Disable caching before importing the DSL so the compiler hook always runs.
os.environ["CUTE_DSL_NO_CACHE"] = "1"

from cutlass._mlir import ir, passmanager
from cutlass.base_dsl.compiler import Compiler

from repro_tuple_add_chain import main

original_compile = Compiler.compile
output = Path(os.environ.get("REPRO_DUMP_DIR", "tuple-retry-dump"))


def compile_with_retry(self, module, pipeline, *args, **kwargs):
    """Retry the identical full pipeline on serialized failure-state MLIR."""
    output.mkdir(parents=True, exist_ok=False)
    (output / "pipeline.txt").write_text(pipeline)
    (output / "before.mlir").write_text(str(module))
    pm = passmanager.PassManager.parse(pipeline)
    (output / "expanded-pipeline.txt").write_text(str(pm))
    pm.enable_verifier(kwargs.get("enable_verifier", False))
    pm.enable_timing()
    diagnostics = []

    def record_diagnostic(diagnostic):
        """Retain native diagnostics without changing pass behavior."""
        diagnostics.append(str(diagnostic))
        return True

    handler = module.context.attach_diagnostic_handler(record_diagnostic)
    try:
        pm.run(module.operation)
    except ir.MLIRError as error:
        (output / "first-error.txt").write_text(str(error))
        (output / "after-failure.mlir").write_text(str(module))
        print(f"First pipeline failed: {error}")
    else:
        raise RuntimeError("Expected the first pipeline to fail; use 4.8.0 and --n 10")
    finally:
        handler.detach()
        (output / "first-diagnostics.json").write_text(json.dumps(diagnostics, indent=2))
        # Destruction prints native pass timings, including the last pass reached.
        del pm

    replay = ir.Module.parse((output / "after-failure.mlir").read_text())
    replay.operation.verify()
    result = original_compile(self, replay, pipeline, *args, **kwargs)
    (output / "after-retry.mlir").write_text(str(result))
    print("Full pipeline succeeded on reparsed after-failure.mlir")
    return result


if __name__ == "__main__":
    Compiler.compile = compile_with_retry
    try:
        main()
    finally:
        Compiler.compile = original_compile