test: llvm.inline_asm with Int32._mlir_type matching cvt_i8_bf16 pattern
This commit is contained in:
90
tests/unit/test_ptx_llvm.py
Normal file
90
tests/unit/test_ptx_llvm.py
Normal file
@@ -0,0 +1,90 @@
|
||||
"""Test: try llvm.inline_asm matching cvt_i8_bf16 pattern exactly."""
|
||||
import torch
|
||||
import cutlass.cute as cute
|
||||
import cutlass.torch as cutlass_torch
|
||||
from cutlass.cutlass_dsl import dsl_user_op
|
||||
from cutlass._mlir.dialects import llvm
|
||||
from cutlass.cute.typing import Float32, Int32
|
||||
import cutlass
|
||||
|
||||
|
||||
# Approach 1: llvm.inline_asm with Int32._mlir_type (matching cvt_i8_bf16 pattern)
|
||||
@dsl_user_op
|
||||
def f32_to_i32_rni_v1(x: Float32, *, loc=None, ip=None) -> Int32:
|
||||
val_i32 = llvm.inline_asm(
|
||||
Int32._mlir_type(),
|
||||
[Float32(x).ir_value(loc=loc, ip=ip)],
|
||||
"cvt.rni.s32.f32 $0, $1;",
|
||||
"=r,f",
|
||||
has_side_effects=False,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return Int32(val_i32)
|
||||
|
||||
|
||||
# Approach 2: llvm.inline_asm without asm_dialect (default)
|
||||
@dsl_user_op
|
||||
def f32_to_i32_rni_v2(x: Float32, *, loc=None, ip=None) -> Int32:
|
||||
val_i32 = llvm.inline_asm(
|
||||
Int32._mlir_type(),
|
||||
[Float32(x).ir_value(loc=loc, ip=ip)],
|
||||
"cvt.rni.s32.f32 $0, $1;",
|
||||
"=r,f",
|
||||
has_side_effects=False,
|
||||
is_align_stack=False,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return Int32(val_i32)
|
||||
|
||||
|
||||
# Approach 3: Using multi-line block like cvt_i8_bf16
|
||||
@dsl_user_op
|
||||
def f32_to_i32_rni_v3(x: Float32, *, loc=None, ip=None) -> Int32:
|
||||
val_i32 = llvm.inline_asm(
|
||||
Int32._mlir_type(),
|
||||
[Float32(x).ir_value(loc=loc, ip=ip)],
|
||||
"{\n\tcvt.rni.s32.f32 $0, $1;\n\t}",
|
||||
"=r,f",
|
||||
has_side_effects=False,
|
||||
is_align_stack=False,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return Int32(val_i32)
|
||||
|
||||
|
||||
KERNELS = {
|
||||
"v1_mlir_type": (f32_to_i32_rni_v1, None),
|
||||
"v2_no_asm_dialect": (f32_to_i32_rni_v2, None),
|
||||
"v3_multiline": (f32_to_i32_rni_v3, None),
|
||||
}
|
||||
|
||||
import sys
|
||||
approach = sys.argv[1] if len(sys.argv) > 1 else "v1_mlir_type"
|
||||
func = KERNELS[approach][0]
|
||||
|
||||
|
||||
@cute.kernel
|
||||
def test_k(inp: cute.Tensor, out: cute.Tensor):
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
if tidx == Int32(0):
|
||||
x = cute.arch.load(inp.iterator, Float32)
|
||||
r = func(x)
|
||||
cute.arch.store(out.iterator, r)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
x = torch.tensor([3.7], dtype=torch.float32, device='cuda')
|
||||
o = torch.zeros(1, dtype=torch.int32, device='cuda')
|
||||
xc = cutlass_torch.from_dlpack(x).mark_layout_dynamic(leading_dim=0)
|
||||
oc = cutlass_torch.from_dlpack(o).mark_layout_dynamic(leading_dim=0)
|
||||
print(f"Approach: {approach}")
|
||||
print("Compiling...")
|
||||
compiled = cute.compile(test_k, xc, oc)
|
||||
print("Running...")
|
||||
compiled(xc, oc)
|
||||
print(f"Result: {o.item()} (expected 4)")
|
||||
19
tests/unit/test_ptx_llvm_runner.py
Normal file
19
tests/unit/test_ptx_llvm_runner.py
Normal file
@@ -0,0 +1,19 @@
|
||||
"""Runner: test each llvm.inline_asm approach in separate process."""
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
for approach in ["v1_mlir_type", "v2_no_asm_dialect", "v3_multiline"]:
|
||||
print(f"\n{'='*50}")
|
||||
print(f"Testing: {approach}")
|
||||
print(f"{'='*50}")
|
||||
result = subprocess.run(
|
||||
[sys.executable, "tests/unit/test_ptx_llvm.py", approach],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
)
|
||||
print(result.stdout)
|
||||
if result.stderr:
|
||||
# Only show last 500 chars of stderr (LLVM ERROR is usually at the end)
|
||||
print(f"STDERR (last 500): ...{result.stderr[-500:]}")
|
||||
print(f"Exit code: {result.returncode}")
|
||||
Reference in New Issue
Block a user