TheRock/rocm-libraries/shared/stinkytofu/python_module

StinkyTofu Python Module

A Python interface for generating GPU assembly code for AMD GPUs using the StinkyTofu IR.

Features

• Instance-based design - No global state, each builder is independent
• Explicit module management - Create named IR modules and add instructions explicitly
• Architecture-specific - Target different GPU architectures (e.g., gfx1250)
• Thread-safe - Multiple instances can be used concurrently

Building

Prerequisites

• CMake 3.16+
• Python 3.6+
• nanobind: pip install nanobind

Build Instructions

From the repository root:

cd shared/stinkytofu
mkdir -p build && cd build
cmake .. -DSTINKYTOFU_BUILD_PYTHON=ON
cmake --build . --target stinkytofu_python -j12

The Python module will be built to build/lib/stinkytofu.*.so.

Add to Python Path

export PYTHONPATH=/path/to/build/lib:$PYTHONPATH

Examples

Basic Usage

from stinkytofu import StinkyAsmIR, vgpr, sgpr

# Create builder for gfx1250
st = StinkyAsmIR([12, 5, 0])

# Create a module to hold instructions
module = st.createIRList("my_kernel")

# Create instructions and add to module
module.add(st.VAddU32(vgpr(0), vgpr(1), vgpr(2), "add"))
module.add(st.VMulF32(vgpr(3), vgpr(0), vgpr(4), "multiply"))
module.add(st.SBarrier("sync"))

# Emit assembly
print(module.emitAssembly())

Register Ranges

from stinkytofu import StinkyAsmIR, vgpr, acc

st = StinkyAsmIR([12, 5, 0])
module = st.createIRList("mfma_kernel")

# Use count parameter for register ranges
# createMFMA for Matrix-FMA instructions
module.add(st.createMFMA(
    instType="f32",  # input type
    accType="f32",   # accumulator type
    m=16, n=16, k=4, # dimensions
    blocks=1,        # number of blocks
    mfma1k=False,    # not 1k variant
    acc=acc(0, 4),   # dst: a[0:3] - 4 registers
    a=vgpr(0, 4),    # matrix A: v[0:3] - 4 registers
    b=vgpr(4, 4),    # matrix B: v[4:7] - 4 registers
    acc2=acc(0, 4),  # input acc: a[0:3] - 4 registers
    comment="MFMA"
))

print(module.emitAssembly())

Composite Instructions (Architecture-Aware Lowering)

Some instructions may not be supported on all architectures. StinkyTofu automatically handles this by lowering unsupported instructions into equivalent sequences:

from stinkytofu import StinkyAsmIR, vgpr

st = StinkyAsmIR([12, 5, 0])  # gfx1250 supports v_pk_add_f32
module = st.createIRList("composite_example")

# VAddPKF32 demonstrates the composite instruction pattern:
# - On architectures with v_pk_add_f32: returns 1 instruction
# - On older architectures: returns 2 v_add_f32 instructions (lowered)
insts = st.VAddPKF32(vgpr(0, 2), vgpr(2, 2), vgpr(4, 2), "packed add")
print(f"Generated {len(insts)} instruction(s)")  # Prints "1" on gfx1250

# module.add() accepts lists of instructions
module.add(insts)

print(module.emitAssembly())
# Output on gfx1250: v_pk_add_f32 v[0:1], v[2:3], v[4:5] // packed add

Note: All instruction creation methods return a list of instructions. Most of the time this list contains a single instruction, but for composite/lowered instructions, it may contain multiple instructions depending on the target architecture. The module.add() method accepts these lists directly.

Multiple Modules

from stinkytofu import StinkyAsmIR, vgpr

st = StinkyAsmIR([12, 5, 0])

# Create multiple kernels
kernel_a = st.createIRList("kernel_a")
kernel_b = st.createIRList("kernel_b")

# Add different instructions to each
kernel_a.add(st.VAddU32(vgpr(0), vgpr(1), vgpr(2)))
kernel_b.add(st.VMulF32(vgpr(0), vgpr(1), vgpr(2)))

# Emit separately
print("Kernel A:")
print(kernel_a.emitAssembly())

print("\nKernel B:")
print(kernel_b.emitAssembly())

Conditional Assembly

from stinkytofu import StinkyAsmIR, vgpr

st = StinkyAsmIR([12, 5, 0])
module = st.createIRList("conditional_kernel")

# Create instruction but don't add yet
inst = st.VAddU32(vgpr(0), vgpr(1), vgpr(2), "conditional add")

# Add conditionally
use_optimization = True
if use_optimization:
    module.add(inst)

module.add(st.SBarrier("sync"))
print(module.emitAssembly())

More examples in examples/basic_usage.py.

Testing

The test suite uses pytest for better test organization and reporting.

Quick Start

# Install dependencies
pip install -r python_module/tests/requirements.txt

# Run all tests
PYTHONPATH=build/lib:$PYTHONPATH pytest python_module/tests/ -v

Detailed Testing Guide

For comprehensive testing documentation including:

• Running specific tests and test categories
• Using fixtures and parametrized tests
• Coverage reports and parallel execution
• Debugging and CI/CD integration

See tests/testing.md for the complete testing guide.

API Reference

Main Classes

• StinkyAsmIR(arch) - Create builder for target architecture

  • arch: List [major, minor, stepping], e.g., [12, 5, 0] for gfx1250

• IRListModule - Container for instructions

  • add(inst) - Add instruction to module (returns inst for chaining)
  • emitAssembly(emit_cycle_info=False, emit_user_comments=True) - Emit assembly code
  • getName() - Get module name

Register Helpers

• vgpr(idx, count=1) - Vector general purpose register
• sgpr(idx, count=1) - Scalar general purpose register
• acc(idx, count=1) - Accumulator register

Instruction Methods (on StinkyAsmIR)

Vector ALU:

• VAddU32(dst, src0, src1, comment="")
• VMulF32(dst, src0, src1, comment="")

Scalar:

• SAbsI32(dst, src, comment="")
• SBarrier(comment="")

Matrix (Generic Creation):

• createMFMA(output_type, input_type, M, N, K, dst, src0, src1, src2, comment="")
• createMXMFMA(output_type, input_type, M, N, K, dst, src0, src1, src2, comment="")
• createSMFMA(output_type, input_type, M, N, K, dst, src0, src1, src2, index, comment="")

See docs/mfma-instructions.md for detailed MFMA documentation.

All instruction methods return a list of StinkyInstruction* pointers.

License

MIT License - See LICENSE file for details.

Copyright (C) 2025 Advanced Micro Devices, Inc.