| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from cutlass import BFloat16, Constexpr, Float32, Uint16, Uint32, cute |
| 4 | from cutlass._mlir import ir |
| 5 | from cutlass._mlir.dialects import llvm, vector |
| 6 | from cutlass.cutlass_dsl import T, dsl_user_op |
| 7 | |
| 8 | |
| 9 | @dsl_user_op |
| 10 | def fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32: |
| 11 | out = llvm.inline_asm( |
| 12 | T.i32(), |
| 13 | [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], |
| 14 | "cvt.rn.bf16x2.f32 $0, $2, $1;", |
| 15 | "=r,f,f", |
| 16 | has_side_effects=False, |
| 17 | is_align_stack=False, |
| 18 | ) |
| 19 | return Uint32(out) |
| 20 | |
| 21 | |
| 22 | @dsl_user_op |
| 23 | def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]: |
| 24 | if isinstance(data, Uint32): |
| 25 | out = llvm.inline_asm( |
| 26 | llvm.StructType.get_literal([T.f32(), T.f32()]), |
| 27 | [data.ir_value(loc=loc, ip=ip)], |
| 28 | "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;", |
| 29 | "=f,=f,r", |
| 30 | has_side_effects=False, |
| 31 | is_align_stack=False, |
| 32 | loc=loc, |
| 33 | ip=ip, |
| 34 | ) |
| 35 | return ( |
| 36 | Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)), |
| 37 | Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)), |
| 38 | ) |
| 39 | |
| 40 | elif isinstance(data, (cute.Tensor, cute.TensorSSA)): |
| 41 | # recast to 32-bit registers if needed |
| 42 | if data.element_type == BFloat16: |
| 43 | data = cute.recast_tensor(data, Uint32) |
| 44 | assert data.element_type == Uint32 |
| 45 | |
| 46 | # NOTE: the output is always 1D |
| 47 | size = cute.size(data.shape) |
| 48 | out = cute.make_rmem_tensor(size * 2, Float32) |
| 49 | for i in range(size): |
| 50 | out[i * 2], out[i * 2 + 1] = bf16x2_to_fp32x2(data[i]) |
| 51 | return out |
| 52 | |
| 53 | else: |
| 54 | raise ValueError(f"Unsupported type {type(data)}") |
| 55 | |
| 56 | |
| 57 | @dsl_user_op |
| 58 | def fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA: |
| 59 | # there is only fp8->fp16 conversion, hence we need to go |
| 60 | # round trip through fp16. |
| 61 | out = llvm.inline_asm( |
| 62 | llvm.StructType.get_literal([T.i32()] * 2), |
| 63 | [x.ir_value(loc=loc, ip=ip)], |
| 64 | "{\n\t" |
| 65 | ".reg .b16 x0, x1;\n\t" |
| 66 | ".reg .b16 t00, t01, t10, t11;\n\t" |
| 67 | "mov.b32 {x0, x1}, $2;\n\t" |
| 68 | "cvt.rn.f16x2.e4m3x2 $0, x0;\n\t" |
| 69 | "cvt.rn.f16x2.e4m3x2 $1, x1;\n\t" |
| 70 | "mov.b32 {t00, t01}, $0;\n\t" |
| 71 | "mov.b32 {t10, t11}, $1;\n\t" |
| 72 | "cvt.rn.bf16.f16 t00, t00;\n\t" |
| 73 | "cvt.rn.bf16.f16 t01, t01;\n\t" |
| 74 | "cvt.rn.bf16.f16 t10, t10;\n\t" |
| 75 | "cvt.rn.bf16.f16 t11, t11;\n\t" |
| 76 | "mov.b32 $0, {t00, t01};\n\t" |
| 77 | "mov.b32 $1, {t10, t11};\n\t" |
| 78 | "}\n", |
| 79 | "=r,=r,r", |
| 80 | has_side_effects=False, |
| 81 | is_align_stack=False, |
| 82 | ) |
| 83 | vec = vector.from_elements( |
| 84 | ir.VectorType.get([2], T.i32(), loc=loc), |
| 85 | [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(2)], |
| 86 | loc=loc, |
| 87 | ip=ip, |
| 88 | ) |
| 89 | return cute.TensorSSA(vec, 2, Uint32) |
| 90 | |
| 91 | |
| 92 | @dsl_user_op |
| 93 | def fp32x2_to_fp8x2(a0: Float32, a1: Float32, *, loc=None, ip=None) -> Uint16: |
| 94 | out = llvm.inline_asm( |
| 95 | T.i16(), |
| 96 | [ |
| 97 | a0.ir_value(loc=loc, ip=ip), |
| 98 | a1.ir_value(loc=loc, ip=ip), |
| 99 | ], |
| 100 | "cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;", |
| 101 | "=h,f,f", |
| 102 | has_side_effects=False, |
| 103 | is_align_stack=False, |
| 104 | ) |
| 105 | return Uint16(out) |
| 106 | |
| 107 | |
| 108 | @dsl_user_op |
| 109 | def fp8x4_to_fp16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA: |
| 110 | out = llvm.inline_asm( |
| 111 | llvm.StructType.get_literal([T.i32()] * 2), |
| 112 | [x.ir_value(loc=loc, ip=ip)], |
| 113 | "{\n\t" |
| 114 | ".reg .b16 lo, hi;\n\t" |
| 115 | "mov.b32 {lo, hi}, $2;\n\t" |
| 116 | "cvt.rn.f16x2.e4m3x2 $0, lo;\n\t" |
| 117 | "cvt.rn.f16x2.e4m3x2 $1, hi;\n\t" |
| 118 | "}\n", |
| 119 | "=r,=r,r", |
| 120 | has_side_effects=False, |
| 121 | is_align_stack=False, |
| 122 | loc=loc, |
| 123 | ip=ip, |
| 124 | ) |
| 125 | vec = vector.from_elements( |
| 126 | ir.VectorType.get([2], T.i32(), loc=loc), |
| 127 | [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(2)], |
| 128 | loc=loc, |
| 129 | ip=ip, |
| 130 | ) |
| 131 | return cute.TensorSSA(vec, 2, Uint32) |
| 132 | |
| 133 | |
| 134 | @dsl_user_op |
| 135 | def fp32x4_to_fp8x4( |
| 136 | a0: Float32, |
| 137 | a1: Float32, |
| 138 | a2: Float32, |
| 139 | a3: Float32, |
| 140 | *, |
| 141 | loc=None, |
| 142 | ip=None, |
| 143 | ) -> Uint32: |
| 144 | # Pack four FP32 values into one b32 of four e4m3 bytes, byte order |
| 145 | # {a0, a1, a2, a3} from low to high address. |
| 146 | out = llvm.inline_asm( |
| 147 | T.i32(), |
| 148 | [ |
| 149 | a0.ir_value(loc=loc, ip=ip), |
| 150 | a1.ir_value(loc=loc, ip=ip), |
| 151 | a2.ir_value(loc=loc, ip=ip), |
| 152 | a3.ir_value(loc=loc, ip=ip), |
| 153 | ], |
| 154 | "{\n\t" |
| 155 | ".reg .b16 t0, t1;\n\t" |
| 156 | "cvt.rn.satfinite.e4m3x2.f32 t0, $2, $1;\n\t" |
| 157 | "cvt.rn.satfinite.e4m3x2.f32 t1, $4, $3;\n\t" |
| 158 | "mov.b32 $0, {t0, t1};\n\t" |
| 159 | "}\n", |
| 160 | "=r,f,f,f,f", |
| 161 | has_side_effects=False, |
| 162 | is_align_stack=False, |
| 163 | ) |
| 164 | return Uint32(out) |
| 165 | |
| 166 | |
| 167 | @dsl_user_op |
| 168 | def fp32x8_to_fp4x8( |
| 169 | vals: cute.Tensor, |
| 170 | offset: Constexpr[int], |
| 171 | *, |
| 172 | loc=None, |
| 173 | ip=None, |
| 174 | ) -> Uint32: |
| 175 | # Pack eight scaled FP32 values into four E2M1x2 bytes, returned as one b32. |
| 176 | assert vals.element_type is Float32 |
| 177 | out = llvm.inline_asm( |
| 178 | T.i32(), |
| 179 | [vals[offset + i].ir_value(loc=loc, ip=ip) for i in range(8)], |
| 180 | "{\n\t" |
| 181 | ".reg .b8 x0, x1, x2, x3;\n\t" |
| 182 | "cvt.rn.satfinite.e2m1x2.f32 x0, $2, $1;\n\t" |
| 183 | "cvt.rn.satfinite.e2m1x2.f32 x1, $4, $3;\n\t" |
| 184 | "cvt.rn.satfinite.e2m1x2.f32 x2, $6, $5;\n\t" |
| 185 | "cvt.rn.satfinite.e2m1x2.f32 x3, $8, $7;\n\t" |
| 186 | "mov.b32 $0, {x0, x1, x2, x3};\n\t" |
| 187 | "}\n", |
| 188 | "=r,f,f,f,f,f,f,f,f", |
| 189 | has_side_effects=False, |
| 190 | is_align_stack=False, |
| 191 | ) |
| 192 | return Uint32(out) |