Function build_cutlass

tests/python/relax/test_codegen_cutlass.py:114–124

(mod, assert_all_bindings_fused=True, num_final_bindings=1)

def build_cutlass(mod, assert_all_bindings_fused=True, num_final_bindings=1):
    mod = partition_for_cutlass(mod)

    if assert_all_bindings_fused:
        assert len(mod["main"].body.blocks[0].bindings) == num_final_bindings, (
            "Not all bindings are fused. " + str(mod["main"])
        )

    codegen_pass = relax.transform.RunCodegen({"cutlass": {"sm": 80, "find_first_valid": True}})
    mod = codegen_pass(mod)
    return mod


def get_result_with_relax_cutlass_offload(