reverse-engineering-anything/src/gpu/GpuElfKernels.ts

import { readCudaTarget } from "./CudaTarget.js";
import { readAmdGpuNoteTargetKey } from "./AmdGpuTarget.js";
import type { GpuPayloadTarget } from "./GpuObject.js";
import { err, ok, type Result } from "../domain/result.js";
import {
  CudaMetadataBudget,
  readCudaInfo,
  readCudaParameters,
} from "./CudaInfo.js";
import type { GpuKernel, GpuLocation } from "../domain/gpu/gpuCodeObjects.js";
import {
  BinaryReader,
  BinaryFormatFailure,
  BinaryResourceFailure,
} from "../native/metadata/BinaryContainer.js";
import {
  readElfSectionDirectory,
  type NativeSection,
  type NativeSectionDirectory,
} from "../native/metadata/NativeSectionDirectory.js";
import { readElfString, readElfSymbol } from "../native/metadata/ElfSymbol.js";
import {
  gpuCount,
  readGpuNativeDirectory,
  rejectGpu,
  type GpuSource,
} from "./GpuSource.js";
import {
  AMD_METADATA_BYTES,
  amdCount,
  amdRecord,
  amdString,
  readAmdMetadata,
} from "./AmdGpuMetadata.js";
const sectionBytes = (
  source: GpuSource,
  section: NativeSection,
  budget?: CudaMetadataBudget,
): Promise<Buffer> => {
  if (!section.file_backed || (section.flags & 0x800n) !== 0n)
    rejectGpu("GPU metadata section has no supported uncompressed file range");
  budget?.read(gpuCount(section.size));
  return source.read(gpuCount(section.offset), gpuCount(section.size));
};
const site = (
  source: GpuSource,
  section: NativeSection,
  offset = 0,
  size = gpuCount(section.size),
): GpuLocation => ({
  container_id: source.containerId,
  offset: gpuCount(section.offset) + offset,
  size,
  section_index: section.index,
  section_name: section.name,
});
const blankKernel = (
  source: GpuSource,
  id: string,
  name: string,
  location: GpuLocation,
): GpuKernel => ({
  id: `${source.containerId!}/${id}`,
  object_id: source.containerId!,
  representation: "cuda-elf-metadata",
  name,
  symbol: name,
  parameter_count: null,
  parameter_bytes: null,
  parameters: [],
  registers: {
    physical: null,
    sgpr: null,
    vgpr: null,
    agpr: null,
    virtual: [],
  },
  shared_bytes: null,
  private_bytes: null,
  location,
  sources: { parameters: null, registers: null, shared: null, private: null },
  limitations: [],
});
export const readGpuElfKernels = async (
  source: GpuSource,
  kind: "cubin" | "amdgpu",
  onKernel: (kernel: GpuKernel) => void,
  preparedAmd?: AmdElfMetadata,
  onTarget?: (target: GpuPayloadTarget) => void,
): Promise<{
  limitations: string[];
  coverage: "complete" | "partial";
}> => {
  if (kind === "amdgpu" && preparedAmd !== undefined)
    return amdKernels(source, preparedAmd, onKernel, onTarget);
  const directory = await readGpuNativeDirectory(
    source,
    readElfSectionDirectory(source.size),
  );
  if (directory.machine !== (kind === "cubin" ? 190 : 224))
    rejectGpu("GPU entry kind disagrees with ELF machine");
  if (kind === "cubin") {
    onTarget?.(readCudaTarget(source, directory));
    const coverage = await cudaKernels(source, directory, onKernel);
    return {
      coverage,
      limitations: [
        ...(directory.sections.some((s) => s.name === null)
          ? [
              "Some ELF section names are unknown; named metadata coverage is partial",
            ]
          : []),
        ...(directory.sections.some((s) => s.type === 2 || s.type === 11)
          ? []
          : [
              "No CUBIN symbol table was available; kernel declaration coverage is unknown",
            ]),
      ],
    };
  }
  const metadata = await readAmdElfMetadata(source, directory);
  const parsed = amdKernels(
    source,
    metadata.ok ? metadata.value : metadata.error.partial,
    onKernel,
    onTarget,
  );
  if (!metadata.ok) throw metadata.error.cause;
  return parsed;
};
const indexCudaMetadata = async (
  source: GpuSource,
  directory: NativeSectionDirectory,
) => {
  const infosByText = new Map<number, NativeSection[]>(),
    sharedByText = new Map<number, NativeSection[]>();
  const registerFacts = new Map<string, Map<number, GpuLocation>>();
  const budget = new CudaMetadataBudget(),
    seenMetadata = new Set<string>();
  for (const row of directory.sections) {
    const map =
      row.type === 8 && row.name?.startsWith(".nv.shared.") === true
        ? sharedByText
        : row.type === 0x70000000 && row.name?.startsWith(".nv.info.") === true
          ? infosByText
          : undefined;
    const global = row.type === 0x70000000 && row.name === ".nv.info";
    if (map === undefined && !global) continue;
    // The same scoped byte declaration may be aliased by many ELF headers.
    const key = `${row.type}:${row.flags}:${row.name}:${row.offset}:${row.size}:${row.link}:${row.info}`;
    if (seenMetadata.has(key)) continue;
    budget.retain(key);
    seenMetadata.add(key);
    if (map !== undefined) {
      budget.retain({ text: row.info, section: row.index });
      const values = map.get(row.info) ?? [];
      values.push(row);
      map.set(row.info, values);
    }
    if (global) {
      for (const record of readCudaInfo(
        await sectionBytes(source, row, budget),
        directory.little,
        budget,
      )) {
        if (record.format !== 4 || record.attribute !== 0x2f) continue;
        if (record.data.length !== 8)
          rejectGpu("CUDA register count record has an unsupported extent");
        const r = new BinaryReader(record.data, directory.little),
          key = `${row.link}:${r.u32(0)}`;
        const values = registerFacts.get(key) ?? new Map<number, GpuLocation>();
        const value = r.u32(4);
        if (!values.has(value)) {
          const location = site(source, row, record.offset, record.length);
          budget.retain({ symbol: key, value, location });
          values.set(value, location);
          registerFacts.set(key, values);
        }
      }
    }
  }
  const allNamesKnown = directory.sections.every((row) => row.name !== null);
  return {
    source,
    directory,
    infosByText,
    sharedByText,
    registerFacts,
    allNamesKnown,
    budget,
    parameters: new Map<number, ReturnType<typeof readCudaParameters>>(),
  };
};
type CudaMetadataIndex = Awaited<ReturnType<typeof indexCudaMetadata>>;

const attachCudaResources = (
  kernel: GpuKernel,
  text: NativeSection,
  index: CudaMetadataIndex,
  symbolKey: string,
): void => {
  const { source, directory } = index;
  const registerRows = index.registerFacts.get(symbolKey);
  const registerValues = new Set(registerRows?.keys());
  const sectionRegisters = text.info >>> 24;
  kernel.registers.physical =
    registerValues.size === 1
      ? [...registerValues][0]!
      : registerValues.size === 0 && sectionRegisters !== 0
        ? sectionRegisters
        : null;
  if (registerValues.size === 1)
    kernel.sources.registers = registerRows!.values().next().value!;
  else if (sectionRegisters !== 0)
    kernel.sources.registers = {
      container_id: source.containerId,
      offset: text.header_offset,
      size: directory.bits === 64 ? 64 : 40,
      section_index: text.index,
      section_name: text.name,
    };
  if (
    registerValues.size > 1 ||
    (kernel.registers.physical !== null &&
      sectionRegisters !== 0 &&
      kernel.registers.physical !== sectionRegisters)
  ) {
    kernel.registers.physical = null;
    kernel.limitations.push(
      "Conflicting CUDA register declarations; physical count remains unknown",
    );
  }
  const shared = index.sharedByText.get(text.index) ?? [];
  kernel.shared_bytes =
    shared.length === 1
      ? gpuCount(shared[0]!.size)
      : shared.length === 0 && index.allNamesKnown
        ? 0
        : null;
  if (shared.length === 1)
    kernel.sources.shared = {
      container_id: source.containerId,
      offset: shared[0]!.header_offset,
      size: directory.bits === 64 ? 64 : 40,
      section_index: shared[0]!.index,
      section_name: shared[0]!.name,
    };
};

const cudaKernels = async (
  source: GpuSource,
  directory: NativeSectionDirectory,
  onKernel: (kernel: GpuKernel) => void,
): Promise<"complete" | "partial"> => {
  const index = await indexCudaMetadata(source, directory);
  let complete = index.allNamesKnown;
  let symbolTables = 0;
  for (const section of directory.sections) {
    if (section.type !== 2 && section.type !== 11) continue;
    symbolTables++;
    const width = gpuCount(section.entry_size),
      minimum = directory.bits === 64 ? 24 : 16;
    if (width < minimum || section.size % BigInt(width) !== 0n)
      rejectGpu("CUBIN symbol table width is invalid");
    const namesSection = directory.sections[section.link];
    if (namesSection === undefined || namesSection.type !== 3)
      rejectGpu("CUBIN symbol string table is invalid");
    const names = await sectionBytes(source, namesSection!, index.budget);
    const bytes = await sectionBytes(source, section, index.budget);
    const symbolCount = bytes.length / width;
    let extendedIndices: Buffer | null = null;
    for (let offset = 0; offset < bytes.length; offset += width) {
      index.budget.visit();
      const symbol = readElfSymbol(
        bytes.subarray(offset, offset + width),
        directory.bits,
        directory.little,
      );
      // STO_CUDA_ENTRY distinguishes launchable kernels from ordinary device functions.
      if (symbol.type !== 2 || (symbol.other & 0x10) === 0) continue;
      let sectionIndex = symbol.section_index;
      if (sectionIndex === 0xffff) {
        if (extendedIndices === null) {
          const matches = directory.sections.filter(
            (row) => row.type === 18 && row.link === section.index,
          );
          if (matches.length !== 1)
            rejectGpu(
              "CUBIN extended symbol section-index table is unavailable",
            );
          const extended = matches[0]!;
          if (
            extended.entry_size !== 4n ||
            extended.size !== BigInt(symbolCount) * 4n
          )
            rejectGpu(
              "CUBIN extended symbol section-index table has an invalid extent",
            );
          extendedIndices = await sectionBytes(source, extended, index.budget);
        }
        sectionIndex = new BinaryReader(extendedIndices, directory.little).u32(
          (offset / width) * 4,
        );
      }
      const text = directory.sections[sectionIndex];
      if (text === undefined || text.type !== 1 || (text.flags & 4n) === 0n)
        rejectGpu("CUBIN kernel symbol has no executable section");
      const kernel = blankKernel(
        source,
        `sym:${section.index}:${offset / width}`,
        readElfString(names, symbol.name_offset),
        site(source, section, offset, width),
      );
      attachCudaResources(
        kernel,
        text!,
        index,
        `${section.index}:${offset / width}`,
      );
      const infos = index.infosByText.get(text!.index) ?? [];
      if (infos.length === 1) {
        const info = infos[0]!;
        let decoded = index.parameters.get(info.index);
        if (decoded === undefined) {
          decoded = readCudaParameters(
            await sectionBytes(source, info, index.budget),
            directory.little,
            index.budget,
          );
          index.budget.retain({ parameter_section: info.index });
          index.parameters.set(info.index, decoded);
        }
        if (decoded.limitations.length > 0) complete = false;
        kernel.sources.parameters = site(source, infos[0]!);
        kernel.parameter_bytes = decoded.parameter_bytes;
        kernel.parameter_count = decoded.parameter_count;
        kernel.parameters = decoded.parameters;
        kernel.limitations.push(...decoded.limitations);
      } else {
        complete = false;
        kernel.limitations.push(
          "Exactly one matching CUDA nv.info parameter section was not available",
        );
      }
      kernel.limitations.push(
        "CUDA metadata fields are declarations from ELF symbols, section headers and nv.info; SASS and runtime behavior were not inspected",
      );
      onKernel(kernel);
    }
  }
  return complete && symbolTables > 0 ? "complete" : "partial";
};
export interface AmdElfMetadata {
  readonly target: string | null;
  readonly target_key: string | null;
  readonly limitations: string[];
  readonly notes: {
    section: NativeSection;
    noteOffset: number;
    desc: number;
    size: number;
    kernels: unknown[];
  }[];
}
/** One bounded parsed owner for target discovery and later selected kernel projection. */
export const readAmdElfMetadata = async (
  source: GpuSource,
  directory: NativeSectionDirectory,
): Promise<
  Result<
    AmdElfMetadata,
    { kind: "incomplete-amd-metadata"; cause: unknown; partial: AmdElfMetadata }
  >
> => {
  const targets = new Set<string>(),
    limitations: string[] = [];
  const notes: AmdElfMetadata["notes"] = [];
  let retainedBytes = 0;
  const observed = (): AmdElfMetadata => {
    const target = targets.size === 1 ? [...targets][0]! : null;
    // LLVM notes use both the empty and explicit `unknown` environment forms.
    const target_key = target === null ? null : readAmdGpuNoteTargetKey(target);
    return { target, target_key, limitations, notes };
  };
  try {
    if (directory.machine !== 224)
      rejectGpu("AMDGPU metadata requires an AMDGPU ELF machine");
    for (const section of directory.sections) {
      if (section.type !== 7) continue;
      const sectionSize = gpuCount(section.size);
      if (sectionSize > AMD_METADATA_BYTES - retainedBytes)
        throw new BinaryResourceFailure(
          "structure",
          AMD_METADATA_BYTES,
          "AMDGPU ELF note discovery exceeds the 8 MiB retained metadata budget",
        );
      retainedBytes += sectionSize;
      const bytes = await sectionBytes(source, section),
        r = new BinaryReader(bytes, directory.little);
      for (let p = 0; p < bytes.length;) {
        if (bytes.length - p < 12) rejectGpu("Truncated ELF note header");
        const namesize = r.u32(p),
          size = r.u32(p + 4),
          type = r.u32(p + 8);
        const nameOffset = p + 12,
          desc = nameOffset + Math.ceil(namesize / 4) * 4,
          end = desc + Math.ceil(size / 4) * 4;
        r.range(nameOffset, namesize, "ELF note vendor");
        r.range(desc, size, "ELF note descriptor");
        if (end > bytes.length)
          rejectGpu("ELF note alignment extends outside its section");
        const vendor = bytes
          .subarray(nameOffset, nameOffset + namesize)
          .toString("latin1")
          .replace(/\0$/u, "");
        if (vendor === "AMDGPU" && type === 32) {
          const metadata = amdRecord(
            readAmdMetadata(bytes.subarray(desc, desc + size)),
          );
          const kernels = metadata?.["amdhsa.kernels"];
          if (!Array.isArray(kernels))
            return rejectGpu("AMDGPU metadata has no valid kernel array");
          const target = amdString(metadata!["amdhsa.target"]);
          if (target !== null) targets.add(target);
          notes.push({ section, noteOffset: p, desc, size, kernels });
        } else if (vendor === "AMD" && (type === 10 || type === 11))
          limitations.push(
            "Legacy AMDGPU YAML metadata is outside MessagePack coverage",
          );
        p = end;
      }
    }
    if (notes.length === 0)
      throw new BinaryFormatFailure(
        "unsupported",
        "No AMDGPU v3+ MessagePack kernel metadata was available",
      );
    return ok(observed());
  } catch (cause) {
    return err({ kind: "incomplete-amd-metadata", cause, partial: observed() });
  }
};

const amdKernels = (
  source: GpuSource,
  prepared: AmdElfMetadata,
  onKernel: (kernel: GpuKernel) => void,
  onTarget?: (target: GpuPayloadTarget) => void,
): { limitations: string[]; coverage: "complete" } => {
  onTarget?.({
    value: prepared.target,
    key: prepared.target_key,
    scope: "exact",
  });
  for (const {
    section,
    noteOffset,
    desc,
    size,
    kernels: rows,
  } of prepared.notes) {
    for (let index = 0; index < rows.length; index++) {
      const value = amdRecord(rows[index]);
      const name = value === null ? null : amdString(value[".name"]);
      if (value === null || name === null)
        rejectGpu("AMDGPU kernel metadata has no name");
      const kernel = blankKernel(
        source,
        `note:${section.index}:${noteOffset}:${index}`,
        name!,
        site(source, section, desc, size),
      );
      kernel.representation = "amdgpu-metadata";
      kernel.sources = {
        parameters: kernel.location,
        registers: kernel.location,
        shared: kernel.location,
        private: kernel.location,
      };
      kernel.symbol = amdString(value![".symbol"]);
      kernel.parameter_bytes = amdCount(value![".kernarg_segment_size"]);
      kernel.shared_bytes = amdCount(value![".group_segment_fixed_size"]);
      kernel.private_bytes = amdCount(value![".private_segment_fixed_size"]);
      kernel.registers.sgpr = amdCount(value![".sgpr_count"]);
      kernel.registers.vgpr = amdCount(value![".vgpr_count"]);
      kernel.registers.agpr = amdCount(value![".agpr_count"]);
      const args = value![".args"];
      if (Array.isArray(args)) {
        kernel.parameter_count = args.length;
        kernel.parameters = args.map((argument, ordinal) => {
          const row = amdRecord(argument);
          if (row === null) rejectGpu("Malformed AMDGPU argument metadata");
          const kind = amdString(row![".value_kind"]);
          return {
            ordinal,
            name: amdString(row![".name"]),
            offset: amdCount(row![".offset"]),
            size: amdCount(row![".size"]),
            type: amdString(row![".type_name"]),
            hidden: kind === null ? null : kind.startsWith("hidden_"),
          };
        });
      }
      kernel.limitations.push(
        "AMDGPU parameter count includes hidden ABI arguments when present; missing fields remain unknown. Kernel metadata is a declaration, not runtime verification.",
      );
      onKernel(kernel);
    }
  }
  return {
    limitations: prepared.limitations,
    coverage: "complete",
  };
};