submission 154437

"""
v635: Protocol ID_SWEEP - Black-box reverse engineering of cuBLASLt algo config attributes
Sweep attribute IDs 0-60 to find valid (potentially undocumented) attributes
"""
import sys
import ctypes
from ctypes import c_void_p, c_int, c_uint32, c_uint64, c_size_t, POINTER, byref, create_string_buffer

print("=" * 70, file=sys.stderr)
print("PROTOCOL: ID_SWEEP - cuBLASLt Attribute Discovery", file=sys.stderr)
print("=" * 70, file=sys.stderr)

# Load libraries
try:
    cublaslt = ctypes.CDLL("/usr/local/cuda/lib64/libcublasLt.so")
    cudart = ctypes.CDLL("libcudart.so")
    print("[+] Libraries loaded successfully", file=sys.stderr)
except Exception as e:
    print(f"[-] Library load failed: {e}", file=sys.stderr)
    sys.exit(1)

# cuBLAS status codes
CUBLAS_STATUS_SUCCESS = 0
CUBLAS_STATUS_NOT_SUPPORTED = 8
CUBLAS_STATUS_INVALID_VALUE = 7

# Data types
CUDA_R_16F = 2   # float16
CUDA_R_32F = 0   # float32

# Matmul operation types
CUBLASLT_EPILOGUE_DEFAULT = 1

# Order types
CUBLASLT_ORDER_ROW = 1
CUBLASLT_ORDER_COL = 0

# =============================================================================
# STEP 1: Setup - Create handle and descriptors
# =============================================================================
print("\n[STEP 1] Creating cuBLASLt handle and descriptors...", file=sys.stderr)

# cublasLtCreate
cublaslt.cublasLtCreate.argtypes = [POINTER(c_void_p)]
cublaslt.cublasLtCreate.restype = c_int

handle = c_void_p()
status = cublaslt.cublasLtCreate(byref(handle))
if status != CUBLAS_STATUS_SUCCESS:
    print(f"[-] cublasLtCreate failed: {status}", file=sys.stderr)
    sys.exit(1)
print(f"[+] Handle created: {hex(handle.value)}", file=sys.stderr)

# Create matrix layout descriptors for a simple 128x128x128 FP16 GEMM
M, N, K = 128, 128, 128

# cublasLtMatrixLayoutCreate
cublaslt.cublasLtMatrixLayoutCreate.argtypes = [POINTER(c_void_p), c_int, c_uint64, c_uint64, c_uint64]
cublaslt.cublasLtMatrixLayoutCreate.restype = c_int

layout_a = c_void_p()
layout_b = c_void_p()
layout_c = c_void_p()

# A: M x K, row-major (ld = K)
status = cublaslt.cublasLtMatrixLayoutCreate(byref(layout_a), CUDA_R_16F, M, K, K)
if status != 0:
    print(f"[-] Layout A creation failed: {status}", file=sys.stderr)
else:
    print(f"[+] Layout A created", file=sys.stderr)

# B: K x N, row-major (ld = N)
status = cublaslt.cublasLtMatrixLayoutCreate(byref(layout_b), CUDA_R_16F, K, N, N)
if status != 0:
    print(f"[-] Layout B creation failed: {status}", file=sys.stderr)
else:
    print(f"[+] Layout B created", file=sys.stderr)

# C: M x N, row-major (ld = N)
status = cublaslt.cublasLtMatrixLayoutCreate(byref(layout_c), CUDA_R_16F, M, N, N)
if status != 0:
    print(f"[-] Layout C creation failed: {status}", file=sys.stderr)
else:
    print(f"[+] Layout C created", file=sys.stderr)

# Create matmul descriptor
cublaslt.cublasLtMatmulDescCreate.argtypes = [POINTER(c_void_p), c_int, c_int]
cublaslt.cublasLtMatmulDescCreate.restype = c_int

matmul_desc = c_void_p()
# computeType = CUBLAS_COMPUTE_16F = 64, scaleType = CUDA_R_16F = 2
CUBLAS_COMPUTE_16F = 64
status = cublaslt.cublasLtMatmulDescCreate(byref(matmul_desc), CUBLAS_COMPUTE_16F, CUDA_R_16F)
if status != 0:
    print(f"[-] MatmulDesc creation failed: {status}", file=sys.stderr)
else:
    print(f"[+] MatmulDesc created", file=sys.stderr)

# Create preference descriptor
cublaslt.cublasLtMatmulPreferenceCreate.argtypes = [POINTER(c_void_p)]
cublaslt.cublasLtMatmulPreferenceCreate.restype = c_int

preference = c_void_p()
status = cublaslt.cublasLtMatmulPreferenceCreate(byref(preference))
if status != 0:
    print(f"[-] Preference creation failed: {status}", file=sys.stderr)
else:
    print(f"[+] Preference created", file=sys.stderr)

# =============================================================================
# STEP 2: Get a valid algorithm via heuristics
# =============================================================================
print("\n[STEP 2] Getting valid algorithm via heuristics...", file=sys.stderr)

# cublasLtMatmulAlgoGetHeuristic
# Returns algorithms in a result array

# Define the result structure (cublasLtMatmulHeuristicResult_t)
# This is an opaque struct, but we know it contains cublasLtMatmulAlgo_t at offset 0
# Size is approximately 64 bytes based on CUDA headers

class CublasLtMatmulHeuristicResult(ctypes.Structure):
    _fields_ = [
        ("algo", ctypes.c_ubyte * 64),  # cublasLtMatmulAlgo_t (opaque, ~64 bytes)
        ("workspaceSize", c_size_t),
        ("state", c_int),
        ("wavesCount", ctypes.c_float),
        ("reserved", ctypes.c_ubyte * 16),
    ]

cublaslt.cublasLtMatmulAlgoGetHeuristic.argtypes = [
    c_void_p,  # handle
    c_void_p,  # matmulDesc
    c_void_p,  # layout A
    c_void_p,  # layout B
    c_void_p,  # layout C
    c_void_p,  # layout D (same as C for our case)
    c_void_p,  # preference
    c_int,     # requestedAlgoCount
    POINTER(CublasLtMatmulHeuristicResult),  # results array
    POINTER(c_int),  # returnedAlgoCount
]
cublaslt.cublasLtMatmulAlgoGetHeuristic.restype = c_int

# Request up to 10 algorithms
MAX_ALGOS = 10
results = (CublasLtMatmulHeuristicResult * MAX_ALGOS)()
returned_count = c_int(0)

status = cublaslt.cublasLtMatmulAlgoGetHeuristic(
    handle,
    matmul_desc,
    layout_a,
    layout_b,
    layout_c,
    layout_c,  # D = C
    preference,
    MAX_ALGOS,
    results,
    byref(returned_count)
)

if status != CUBLAS_STATUS_SUCCESS:
    print(f"[-] GetHeuristic failed: {status}", file=sys.stderr)
    print("[!] Trying alternative approach - creating algo directly...", file=sys.stderr)

    # Alternative: Try to create algorithm directly with cublasLtMatmulAlgoInit
    # This might not be available, so we'll handle failure gracefully
    algo_buffer = (ctypes.c_ubyte * 64)()
    have_algo = False
else:
    print(f"[+] GetHeuristic returned {returned_count.value} algorithms", file=sys.stderr)
    if returned_count.value > 0:
        algo_buffer = results[0].algo
        have_algo = True
        print(f"[+] Using first algorithm, workspaceSize={results[0].workspaceSize}", file=sys.stderr)
    else:
        print("[-] No algorithms returned", file=sys.stderr)
        have_algo = False

# =============================================================================
# STEP 3: The Sweep - Probe all attribute IDs
# =============================================================================
print("\n[STEP 3] Sweeping attribute IDs 0-60...", file=sys.stderr)

if not have_algo:
    print("[-] Cannot sweep without valid algorithm object", file=sys.stderr)
else:
    # cublasLtMatmulAlgoConfigSetAttribute
    cublaslt.cublasLtMatmulAlgoConfigSetAttribute.argtypes = [
        ctypes.POINTER(ctypes.c_ubyte * 64),  # algo pointer
        c_int,      # attribute ID
        c_void_p,   # data pointer
        c_size_t,   # data size
    ]
    cublaslt.cublasLtMatmulAlgoConfigSetAttribute.restype = c_int

    # cublasLtMatmulAlgoConfigGetAttribute
    cublaslt.cublasLtMatmulAlgoConfigGetAttribute.argtypes = [
        ctypes.POINTER(ctypes.c_ubyte * 64),  # algo pointer
        c_int,      # attribute ID
        c_void_p,   # data pointer
        c_size_t,   # data size
        POINTER(c_size_t),  # size written
    ]
    cublaslt.cublasLtMatmulAlgoConfigGetAttribute.restype = c_int

    # Known public attributes (for reference)
    KNOWN_ATTRS = {
        0: "ALGO_CONFIG_ID",
        1: "ALGO_CONFIG_TILE_ID",
        2: "ALGO_CONFIG_SPLITK_NUM",
        3: "ALGO_CONFIG_REDUCTION_SCHEME",
        4: "ALGO_CONFIG_CTA_SWIZZLING",
        5: "ALGO_CONFIG_CUSTOM_OPTION",
        6: "ALGO_CONFIG_STAGES_ID",
        7: "ALGO_CONFIG_INNER_SHAPE_ID",
        8: "ALGO_CONFIG_CLUSTER_SHAPE_ID",
    }

    valid_get_attrs = []
    valid_set_attrs = []

    print("\n--- GET Attribute Sweep ---", file=sys.stderr)
    for attr_id in range(61):
        # Try GET first with uint32
        value_u32 = c_uint32(0)
        size_written = c_size_t(0)

        status = cublaslt.cublasLtMatmulAlgoConfigGetAttribute(
            byref(algo_buffer),
            attr_id,
            byref(value_u32),
            ctypes.sizeof(value_u32),
            byref(size_written)
        )

        if status == CUBLAS_STATUS_SUCCESS:
            name = KNOWN_ATTRS.get(attr_id, "UNKNOWN")
            valid_get_attrs.append((attr_id, name, value_u32.value, size_written.value))
            print(f"    [GET] ID={attr_id:2d} ({name:30s}): value={value_u32.value}, size={size_written.value}", file=sys.stderr)
        elif status != CUBLAS_STATUS_NOT_SUPPORTED and status != CUBLAS_STATUS_INVALID_VALUE:
            print(f"    [GET] ID={attr_id:2d}: unexpected status {status}", file=sys.stderr)

    print("\n--- SET Attribute Sweep (uint32 payload=1) ---", file=sys.stderr)
    # Make a copy of the algo to avoid corrupting it
    algo_copy = (ctypes.c_ubyte * 64)()
    ctypes.memmove(algo_copy, algo_buffer, 64)

    for attr_id in range(61):
        # Restore original algo
        ctypes.memmove(algo_copy, algo_buffer, 64)

        # Try SET with uint32 value = 1
        value_u32 = c_uint32(1)

        status = cublaslt.cublasLtMatmulAlgoConfigSetAttribute(
            byref(algo_copy),
            attr_id,
            byref(value_u32),
            ctypes.sizeof(value_u32)
        )

        if status == CUBLAS_STATUS_SUCCESS:
            name = KNOWN_ATTRS.get(attr_id, "UNKNOWN")
            valid_set_attrs.append((attr_id, name, "u32"))
            print(f"    [SET] ID={attr_id:2d} ({name:30s}): SUCCESS with u32", file=sys.stderr)
        elif status == CUBLAS_STATUS_INVALID_VALUE:
            # Value rejected but attribute might be valid
            print(f"    [SET] ID={attr_id:2d}: INVALID_VALUE (attr exists but value rejected)", file=sys.stderr)
            valid_set_attrs.append((attr_id, KNOWN_ATTRS.get(attr_id, "UNKNOWN"), "exists_invalid_val"))

    # =============================================================================
    # STEP 4: Analysis
    # =============================================================================
    print("\n" + "=" * 70, file=sys.stderr)
    print("ANALYSIS RESULTS", file=sys.stderr)
    print("=" * 70, file=sys.stderr)

    print("\n[GET] Valid readable attributes:", file=sys.stderr)
    for attr_id, name, value, size in valid_get_attrs:
        marker = " <-- UNDOCUMENTED" if attr_id > 8 else ""
        print(f"    ID={attr_id:2d} {name:35s} value={value:10d} size={size}{marker}", file=sys.stderr)

    print(f"\n[SET] Valid writable attributes:", file=sys.stderr)
    for attr_id, name, status in valid_set_attrs:
        marker = " <-- UNDOCUMENTED CANDIDATE" if attr_id > 8 else ""
        print(f"    ID={attr_id:2d} {name:35s} status={status}{marker}", file=sys.stderr)

    # Highlight candidates
    undocumented_candidates = [x for x in valid_get_attrs if x[0] > 8]
    if undocumented_candidates:
        print(f"\n[!!!] UNDOCUMENTED ATTRIBUTE CANDIDATES (ID > 8):", file=sys.stderr)
        for attr_id, name, value, size in undocumented_candidates:
            print(f"    ID={attr_id}: current_value={value}, size={size} bytes", file=sys.stderr)
    else:
        print(f"\n[!] No undocumented attributes found in range 9-60", file=sys.stderr)

# Cleanup
cublaslt.cublasLtMatmulPreferenceDestroy.argtypes = [c_void_p]
cublaslt.cublasLtMatmulDescDestroy.argtypes = [c_void_p]
cublaslt.cublasLtMatrixLayoutDestroy.argtypes = [c_void_p]
cublaslt.cublasLtDestroy.argtypes = [c_void_p]

cublaslt.cublasLtMatmulPreferenceDestroy(preference)
cublaslt.cublasLtMatmulDescDestroy(matmul_desc)
cublaslt.cublasLtMatrixLayoutDestroy(layout_a)
cublaslt.cublasLtMatrixLayoutDestroy(layout_b)
cublaslt.cublasLtMatrixLayoutDestroy(layout_c)
cublaslt.cublasLtDestroy(handle)

print("\n[+] Cleanup complete", file=sys.stderr)
print("=" * 70, file=sys.stderr)

# Standard kernel for benchmark
import torch
torch._C._set_sm_carveout_experimental(142)

def custom_kernel(data):
    with torch.inference_mode():
        a, b, _, _, sfa_perm, sfb_perm, _ = data
        return torch._scaled_mm(
            a[:, :, 0],
            b[:, :, 0].t(),
            sfa_perm[:, :, :, :, :, 0].permute(2, 4, 0, 1, 3).reshape(-1),
            sfb_perm[:, :, :, :, :, 0].permute(2, 4, 0, 1, 3).reshape(-1),
            bias=None,
            out_dtype=torch.float16
        ).unsqueeze(-1)