|
|
@@ -0,0 +1,5524 @@
|
|
|
+"""
|
|
|
+Copyright (c) 2025 by FlashInfer team.
|
|
|
+
|
|
|
+Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
+you may not use this file except in compliance with the License.
|
|
|
+You may obtain a copy of the License at
|
|
|
+
|
|
|
+ http://www.apache.org/licenses/LICENSE-2.0
|
|
|
+
|
|
|
+Unless required by applicable law or agreed to in writing, software
|
|
|
+distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
+See the License for the specific language governing permissions and
|
|
|
+limitations under the License.
|
|
|
+
|
|
|
+Common utilities for FP4 quantization kernels.
|
|
|
+
|
|
|
+This module contains shared PTX intrinsics, helper functions, and reduction
|
|
|
+utilities used by the fused frontend kernels.
|
|
|
+"""
|
|
|
+
|
|
|
+import functools
|
|
|
+import math
|
|
|
+import operator
|
|
|
+from typing import Callable, Tuple
|
|
|
+
|
|
|
+import cutlass
|
|
|
+import cutlass.cute as cute
|
|
|
+import torch
|
|
|
+import torch.nn.functional as F
|
|
|
+from cutlass import Float32, Int32, Int64, Uint8, Uint16, Uint32, Uint64
|
|
|
+from cutlass.cutlass_dsl import T, dsl_user_op
|
|
|
+from cutlass._mlir import ir
|
|
|
+from cutlass._mlir.dialects import llvm
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# Constants
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+FLOAT4_E2M1_MAX = 6.0 # Maximum value representable in FP4 E2M1
|
|
|
+FLOAT8_E4M3_MAX = 448.0 # Maximum value representable in FP8 E4M3
|
|
|
+SF_VEC_SIZE = 16 # Elements per scale factor block
|
|
|
+COPY_BITS = 128 # 128-bit vectorized loads
|
|
|
+_FP4_MAG_LUT = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
|
|
|
+
|
|
|
+
|
|
|
+def align_up(value: int, alignment: int) -> int:
|
|
|
+ return ((value + alignment - 1) // alignment) * alignment
|
|
|
+
|
|
|
+
|
|
|
+def make_swizzle_indices(rows_padded: int, cols_padded: int, device: torch.device) -> torch.Tensor:
|
|
|
+ """Pre-compute swizzle gather indices for FP8 block-scale factors."""
|
|
|
+ numel = rows_padded * cols_padded
|
|
|
+ idx = torch.arange(numel, device=device)
|
|
|
+ idx_5d = idx.reshape(rows_padded // 128, 4, 32, cols_padded // 4, 4)
|
|
|
+ return idx_5d.permute(0, 3, 2, 1, 4).contiguous().reshape(-1)
|
|
|
+
|
|
|
+
|
|
|
+def swizzle_block_scale(scale: torch.Tensor) -> torch.Tensor:
|
|
|
+ if scale.ndim == 2:
|
|
|
+ scale = scale.unsqueeze(0)
|
|
|
+ squeeze_batch = True
|
|
|
+ elif scale.ndim == 3:
|
|
|
+ squeeze_batch = False
|
|
|
+ else:
|
|
|
+ raise ValueError(f"scale must be 2D or 3D, got {tuple(scale.shape)}")
|
|
|
+
|
|
|
+ batch, rows, cols = scale.shape
|
|
|
+ rows_padded = align_up(rows, 128)
|
|
|
+ cols_padded = align_up(cols, 4)
|
|
|
+
|
|
|
+ padded = torch.zeros((batch, rows_padded, cols_padded), dtype=scale.dtype, device=scale.device)
|
|
|
+ padded[:, :rows, :cols] = scale
|
|
|
+ swizzled = padded.reshape(batch, rows_padded // 128, 4, 32, cols_padded // 4, 4)
|
|
|
+ swizzled = swizzled.permute(0, 1, 4, 3, 2, 5).contiguous()
|
|
|
+ swizzled = swizzled.reshape(batch, rows_padded, cols_padded)
|
|
|
+ return swizzled[0] if squeeze_batch else swizzled
|
|
|
+
|
|
|
+
|
|
|
+def as_grouped_scale_view(scale_storage: torch.Tensor, rows: int, cols: int) -> torch.Tensor:
|
|
|
+ batch = scale_storage.shape[0]
|
|
|
+ rows_padded = align_up(rows, 128)
|
|
|
+ cols_padded = align_up(cols // SF_VEC_SIZE, 4)
|
|
|
+ sf = scale_storage.view(torch.float8_e4m3fn)
|
|
|
+ sf = sf.view(batch, rows_padded // 128, cols_padded // 4, 32, 4, 4)
|
|
|
+ return sf.permute(3, 4, 1, 5, 2, 0)
|
|
|
+
|
|
|
+
|
|
|
+def as_grouped_scale_view_mx(scale_storage: torch.Tensor, rows: int, cols: int) -> torch.Tensor:
|
|
|
+ """Grouped UE8M0 scale view for vec32 (MXFP8 w4a8) block scales."""
|
|
|
+ batch = scale_storage.shape[0]
|
|
|
+ rows_padded = align_up(rows, 128)
|
|
|
+ cols_padded = align_up(cols // 32, 4)
|
|
|
+ sf = scale_storage.view(torch.float8_e8m0fnu)
|
|
|
+ sf = sf.view(batch, rows_padded // 128, cols_padded // 4, 32, 4, 4)
|
|
|
+ return sf.permute(3, 4, 1, 5, 2, 0)
|
|
|
+
|
|
|
+
|
|
|
+def _fp4_quantize_values(x: torch.Tensor) -> torch.Tensor:
|
|
|
+ sign = torch.sign(x)
|
|
|
+ x = torch.abs(x.clone())
|
|
|
+ x[(x >= 0.0) & (x <= 0.25)] = 0.0
|
|
|
+ x[(x > 0.25) & (x < 0.75)] = 0.5
|
|
|
+ x[(x >= 0.75) & (x <= 1.25)] = 1.0
|
|
|
+ x[(x > 1.25) & (x < 1.75)] = 1.5
|
|
|
+ x[(x >= 1.75) & (x <= 2.5)] = 2.0
|
|
|
+ x[(x > 2.5) & (x < 3.5)] = 3.0
|
|
|
+ x[(x >= 3.5) & (x <= 5.0)] = 4.0
|
|
|
+ x[x > 5.0] = 6.0
|
|
|
+ return x * sign
|
|
|
+
|
|
|
+
|
|
|
+def fp4_quantize_values_torch(x: torch.Tensor) -> torch.Tensor:
|
|
|
+ """Pure-Torch FP4 E2M1 quantization with kernel-matching tie-breaking."""
|
|
|
+ return _fp4_quantize_values(x)
|
|
|
+
|
|
|
+
|
|
|
+def _fp4_encode_nibbles(values: torch.Tensor) -> torch.Tensor:
|
|
|
+ mags = values.abs()
|
|
|
+ idx = torch.zeros_like(values, dtype=torch.uint8)
|
|
|
+ for code, mag in enumerate(_FP4_MAG_LUT):
|
|
|
+ idx = torch.where(mags == mag, torch.full_like(idx, code), idx)
|
|
|
+ sign_bit = (values < 0).to(torch.uint8) << 3
|
|
|
+ return idx | sign_bit
|
|
|
+
|
|
|
+
|
|
|
+def pack_grouped_fp4_values(values: torch.Tensor) -> torch.Tensor:
|
|
|
+ """Pack grouped FP4 values `[G, R, C]` into uint8 `[R, C/2, G]`."""
|
|
|
+ nibbles = _fp4_encode_nibbles(values)
|
|
|
+ packed = nibbles.view(values.shape[0], values.shape[1], values.shape[2] // 2, 2)
|
|
|
+ packed = packed[..., 0] | (packed[..., 1] << 4)
|
|
|
+ return packed.permute(1, 2, 0).contiguous()
|
|
|
+
|
|
|
+
|
|
|
+def quantize_grouped_nvfp4_torch(
|
|
|
+ input_tensor: torch.Tensor,
|
|
|
+ row_counts: torch.Tensor,
|
|
|
+ global_scale: torch.Tensor,
|
|
|
+) -> tuple[torch.Tensor, torch.Tensor]:
|
|
|
+ """Pure-Torch grouped NVFP4 quantization with reciprocal global scales."""
|
|
|
+ num_groups, rows, cols = input_tensor.shape
|
|
|
+ if cols % SF_VEC_SIZE != 0:
|
|
|
+ raise ValueError(f"cols must be divisible by {SF_VEC_SIZE}, got {cols}")
|
|
|
+ if global_scale.numel() == 1:
|
|
|
+ global_scale = global_scale.expand(num_groups).contiguous()
|
|
|
+
|
|
|
+ quantized = torch.zeros((num_groups, rows, cols), dtype=torch.float32, device=input_tensor.device)
|
|
|
+ scales = torch.zeros((num_groups, rows, cols // SF_VEC_SIZE), dtype=torch.float32, device=input_tensor.device)
|
|
|
+ for group_idx in range(num_groups):
|
|
|
+ valid_rows = int(row_counts[group_idx].item())
|
|
|
+ if valid_rows == 0:
|
|
|
+ continue
|
|
|
+ x = input_tensor[group_idx, :valid_rows].float()
|
|
|
+ sliced = x.view(valid_rows, cols // SF_VEC_SIZE, SF_VEC_SIZE)
|
|
|
+ block_max = sliced.abs().amax(dim=-1, keepdim=True).to(torch.float32)
|
|
|
+ scale = (global_scale[group_idx] * (block_max / FLOAT4_E2M1_MAX)).to(torch.float8_e4m3fn).to(torch.float32)
|
|
|
+ output_scale = 1.0 / (scale * (1.0 / global_scale[group_idx]))
|
|
|
+ clipped = torch.clamp(sliced * output_scale, -FLOAT4_E2M1_MAX, FLOAT4_E2M1_MAX).view(valid_rows, cols)
|
|
|
+ quantized[group_idx, :valid_rows] = _fp4_quantize_values(clipped)
|
|
|
+ scales[group_idx, :valid_rows] = scale.squeeze(-1)
|
|
|
+
|
|
|
+ packed = pack_grouped_fp4_values(quantized)
|
|
|
+ swizzled = swizzle_block_scale(scales.to(torch.float8_e4m3fn))
|
|
|
+ scale_view = as_grouped_scale_view(swizzled.view(torch.uint8), rows, cols)
|
|
|
+ return packed, scale_view
|
|
|
+
|
|
|
+
|
|
|
+def silu_mul_quantize_grouped_nvfp4_torch(
|
|
|
+ input_tensor: torch.Tensor,
|
|
|
+ row_counts: torch.Tensor,
|
|
|
+ global_scale: torch.Tensor,
|
|
|
+) -> tuple[torch.Tensor, torch.Tensor]:
|
|
|
+ cols = input_tensor.shape[-1] // 2
|
|
|
+ left = input_tensor[..., :cols].float()
|
|
|
+ right = input_tensor[..., cols:].float()
|
|
|
+ activated = (F.silu(left) * right).to(input_tensor.dtype).to(torch.float32)
|
|
|
+ return quantize_grouped_nvfp4_torch(activated, row_counts, global_scale)
|
|
|
+
|
|
|
+
|
|
|
+def relu2_quantize_grouped_nvfp4_torch(
|
|
|
+ input_tensor: torch.Tensor,
|
|
|
+ row_counts: torch.Tensor,
|
|
|
+ global_scale: torch.Tensor,
|
|
|
+) -> tuple[torch.Tensor, torch.Tensor]:
|
|
|
+ """Pure-Torch grouped NVFP4 quantization for ReLU^2 FC1 outputs."""
|
|
|
+ activated = torch.square(F.relu(input_tensor.float()))
|
|
|
+ activated = activated.to(input_tensor.dtype).to(torch.float32)
|
|
|
+ return quantize_grouped_nvfp4_torch(activated, row_counts, global_scale)
|
|
|
+
|
|
|
+
|
|
|
+MX_SF_VEC_SIZE = 32 # Elements per UE8M0 scale block (MXFP8 w4a8 activations)
|
|
|
+_INV_FLOAT8_E4M3_MAX = 1.0 / FLOAT8_E4M3_MAX
|
|
|
+
|
|
|
+
|
|
|
+def pow2_ceil_ue8m0_torch(scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
|
+ """Bit-exact Torch replica of the ``pow2_ceil_ue8m0`` device intrinsic.
|
|
|
+
|
|
|
+ Rounds positive fp32 scales UP to a power of two by bumping the exponent
|
|
|
+ whenever the mantissa is nonzero. Returns ``(rounded_fp32, ue8m0_u8)``;
|
|
|
+ a zero scale maps to (0.0, byte 0).
|
|
|
+ """
|
|
|
+ bits = scale.to(torch.float32).contiguous().view(torch.int32)
|
|
|
+ mant = bits & 0x007FFFFF
|
|
|
+ bumped = torch.where(mant != 0, (bits + 0x00800000) & 0x7F800000, bits)
|
|
|
+ rounded = bumped.view(torch.float32)
|
|
|
+ byte = ((bumped >> 23) & 0xFF).to(torch.uint8)
|
|
|
+ return rounded, byte
|
|
|
+
|
|
|
+
|
|
|
+def _ue8m0_output_scale_torch(byte: torch.Tensor) -> torch.Tensor:
|
|
|
+ """Torch replica of ``ue8m0_to_output_scale``: 2^(127-byte), 0 for byte 0."""
|
|
|
+ inv_bits = (254 - byte.to(torch.int32)).clamp(min=0) << 23
|
|
|
+ inv = inv_bits.view(torch.float32)
|
|
|
+ return torch.where(byte == 0, torch.zeros_like(inv), inv)
|
|
|
+
|
|
|
+
|
|
|
+def quant_dequant_mxfp8_torch(x: torch.Tensor) -> torch.Tensor:
|
|
|
+ """Per-32-block MXFP8 quantize-dequantize roundtrip (oracle helper).
|
|
|
+
|
|
|
+ Matches the in-kernel ``quantize_block_fp8_mx`` numerics bit-for-bit:
|
|
|
+ UE8M0 ceil block scale, E4M3 RN-saturating payload.
|
|
|
+ """
|
|
|
+ orig_shape = x.shape
|
|
|
+ cols = orig_shape[-1]
|
|
|
+ if cols % MX_SF_VEC_SIZE != 0:
|
|
|
+ raise ValueError(f"last dim must be divisible by {MX_SF_VEC_SIZE}, got {cols}")
|
|
|
+ blocked = x.to(torch.float32).reshape(-1, cols // MX_SF_VEC_SIZE, MX_SF_VEC_SIZE)
|
|
|
+ block_max = blocked.abs().amax(dim=-1, keepdim=True)
|
|
|
+ rounded, byte = pow2_ceil_ue8m0_torch(block_max * _INV_FLOAT8_E4M3_MAX)
|
|
|
+ inv = _ue8m0_output_scale_torch(byte)
|
|
|
+ payload = (
|
|
|
+ (blocked * inv)
|
|
|
+ .clamp(-FLOAT8_E4M3_MAX, FLOAT8_E4M3_MAX)
|
|
|
+ .to(torch.float8_e4m3fn)
|
|
|
+ .to(torch.float32)
|
|
|
+ )
|
|
|
+ return (payload * rounded).reshape(orig_shape)
|
|
|
+
|
|
|
+
|
|
|
+def quantize_grouped_mxfp8_torch(
|
|
|
+ input_tensor: torch.Tensor,
|
|
|
+ row_counts: torch.Tensor,
|
|
|
+) -> tuple[torch.Tensor, torch.Tensor]:
|
|
|
+ """Pure-Torch grouped MXFP8 quantization (no global scale).
|
|
|
+
|
|
|
+ Mirrors ``quantize_grouped_nvfp4_torch`` for the w4a8 activation path:
|
|
|
+ returns ``(payload, scale_view)`` where ``payload`` is E4M3 bytes in
|
|
|
+ ``[R, C, G]`` order and ``scale_view`` is the swizzled UE8M0 grouped
|
|
|
+ view (vec32 columns).
|
|
|
+ """
|
|
|
+ num_groups, rows, cols = input_tensor.shape
|
|
|
+ if cols % MX_SF_VEC_SIZE != 0:
|
|
|
+ raise ValueError(f"cols must be divisible by {MX_SF_VEC_SIZE}, got {cols}")
|
|
|
+
|
|
|
+ payload = torch.zeros((num_groups, rows, cols), dtype=torch.uint8, device=input_tensor.device)
|
|
|
+ scales = torch.zeros(
|
|
|
+ (num_groups, rows, cols // MX_SF_VEC_SIZE), dtype=torch.uint8, device=input_tensor.device
|
|
|
+ )
|
|
|
+ for group_idx in range(num_groups):
|
|
|
+ valid_rows = int(row_counts[group_idx].item())
|
|
|
+ if valid_rows == 0:
|
|
|
+ continue
|
|
|
+ x = input_tensor[group_idx, :valid_rows].float()
|
|
|
+ blocked = x.view(valid_rows, cols // MX_SF_VEC_SIZE, MX_SF_VEC_SIZE)
|
|
|
+ block_max = blocked.abs().amax(dim=-1, keepdim=True)
|
|
|
+ rounded, byte = pow2_ceil_ue8m0_torch(block_max * _INV_FLOAT8_E4M3_MAX)
|
|
|
+ inv = _ue8m0_output_scale_torch(byte)
|
|
|
+ q = (
|
|
|
+ (blocked * inv)
|
|
|
+ .clamp(-FLOAT8_E4M3_MAX, FLOAT8_E4M3_MAX)
|
|
|
+ .to(torch.float8_e4m3fn)
|
|
|
+ .view(torch.uint8)
|
|
|
+ .view(valid_rows, cols)
|
|
|
+ )
|
|
|
+ payload[group_idx, :valid_rows] = q
|
|
|
+ scales[group_idx, :valid_rows] = byte.squeeze(-1)
|
|
|
+
|
|
|
+ payload_view = payload.permute(1, 2, 0).contiguous()
|
|
|
+ swizzled = swizzle_block_scale(scales)
|
|
|
+ scale_view = as_grouped_scale_view_mx(swizzled, rows, cols)
|
|
|
+ return payload_view, scale_view
|
|
|
+
|
|
|
+
|
|
|
+def silu_mul_quantize_grouped_mxfp8_torch(
|
|
|
+ input_tensor: torch.Tensor,
|
|
|
+ row_counts: torch.Tensor,
|
|
|
+) -> tuple[torch.Tensor, torch.Tensor]:
|
|
|
+ """SiLU(left)*right then grouped MXFP8 quantization (w4a8 FC2 input)."""
|
|
|
+ cols = input_tensor.shape[-1] // 2
|
|
|
+ left = input_tensor[..., :cols].float()
|
|
|
+ right = input_tensor[..., cols:].float()
|
|
|
+ activated = (F.silu(left) * right).to(input_tensor.dtype).to(torch.float32)
|
|
|
+ return quantize_grouped_mxfp8_torch(activated, row_counts)
|
|
|
+
|
|
|
+
|
|
|
+def relu2_quantize_grouped_mxfp8_torch(
|
|
|
+ input_tensor: torch.Tensor,
|
|
|
+ row_counts: torch.Tensor,
|
|
|
+) -> tuple[torch.Tensor, torch.Tensor]:
|
|
|
+ """ReLU^2 then grouped MXFP8 quantization (w4a8 FC2 input, non-gated)."""
|
|
|
+ activated = torch.square(F.relu(input_tensor.float()))
|
|
|
+ activated = activated.to(input_tensor.dtype).to(torch.float32)
|
|
|
+ return quantize_grouped_mxfp8_torch(activated, row_counts)
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# Architecture Detection
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@functools.lru_cache(maxsize=16)
|
|
|
+def get_sm_version(device: int | torch.device | str | None = None) -> int:
|
|
|
+ """Get the SM version of a CUDA device.
|
|
|
+
|
|
|
+ Args:
|
|
|
+ device: CUDA device to query. Can be an int (device index), torch.device,
|
|
|
+ device string (e.g., 'cuda:0'), or None to use current device.
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ SM version as an integer (e.g., 120 on the current target).
|
|
|
+ """
|
|
|
+ if not torch.cuda.is_available():
|
|
|
+ return 80
|
|
|
+ if device is None:
|
|
|
+ device = torch.cuda.current_device()
|
|
|
+ props = torch.cuda.get_device_properties(device)
|
|
|
+ return props.major * 10 + props.minor
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# PTX Intrinsics - Cluster Operations
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def set_block_rank(
|
|
|
+ smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: Int32, *, loc=None, ip=None
|
|
|
+) -> Int32:
|
|
|
+ """Map smem pointer to address at another CTA rank in the cluster."""
|
|
|
+ smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
|
|
|
+ return Int32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()],
|
|
|
+ "mapa.shared::cluster.u32 $0, $1, $2;",
|
|
|
+ "=r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def store_shared_remote(
|
|
|
+ val: Float32,
|
|
|
+ smem_ptr: cute.Pointer,
|
|
|
+ mbar_ptr: cute.Pointer,
|
|
|
+ peer_cta_rank_in_cluster: Int32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> None:
|
|
|
+ """Store Float32 value to shared memory on a remote CTA in the cluster."""
|
|
|
+ remote_smem_ptr_i32 = set_block_rank(
|
|
|
+ smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
|
|
|
+ ).ir_value()
|
|
|
+ remote_mbar_ptr_i32 = set_block_rank(
|
|
|
+ mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip
|
|
|
+ ).ir_value()
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [remote_smem_ptr_i32, val.ir_value(loc=loc, ip=ip), remote_mbar_ptr_i32],
|
|
|
+ "st.async.shared::cluster.mbarrier::complete_tx::bytes.f32 [$0], $1, [$2];",
|
|
|
+ "r,f,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def elem_pointer(x: cute.Tensor, coord, *, loc=None, ip=None) -> cute.Pointer:
|
|
|
+ """Get pointer to element at coordinate in tensor."""
|
|
|
+ return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip)
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# PTX Intrinsics - 128-bit Vectorized Global Loads/Stores
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_global_v4_u32(
|
|
|
+ base_ptr: Int64, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32, Uint32, Uint32]:
|
|
|
+ """Load 128 bits (4 x uint32) from global memory."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]),
|
|
|
+ [Int64(base_ptr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.global.v4.u32 {$0, $1, $2, $3}, [$4];",
|
|
|
+ "=r,=r,=r,=r,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+ v0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ v1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ v2 = llvm.extractvalue(T.i32(), result, [2], loc=loc, ip=ip)
|
|
|
+ v3 = llvm.extractvalue(T.i32(), result, [3], loc=loc, ip=ip)
|
|
|
+
|
|
|
+ return Uint32(v0), Uint32(v1), Uint32(v2), Uint32(v3)
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# PTX Intrinsics - Non-Coherent Global Loads
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_global_nc_u32(base_ptr: Int64, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Load 32 bits from global memory via non-coherent cache (.nc)."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Int64(base_ptr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.global.nc.u32 $0, [$1];",
|
|
|
+ "=r,l",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_global_b16(base_ptr: Int64, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Load one 16-bit global-memory payload into a 32-bit register.
|
|
|
+
|
|
|
+ PTX permits a bit-typed load into a wider bit register. This is the form
|
|
|
+ emitted by FlashInfer for the DSV4 RoPE gather and avoids rounding the
|
|
|
+ address down to a word followed by a runtime halfword select.
|
|
|
+ """
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Int64(base_ptr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.global.b16 $0, [$1];",
|
|
|
+ "=r,l",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def prefetch_global_l2(base_ptr: Int64, *, loc=None, ip=None) -> None:
|
|
|
+ """Prefetch a global memory line into L2."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [Int64(base_ptr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "prefetch.global.L2 [$0];",
|
|
|
+ "l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_global_nc_v2_u32(
|
|
|
+ base_ptr: Int64, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Load 64 bits (2 x uint32) from global memory via non-coherent cache (.nc)."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Int64(base_ptr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.global.nc.v2.u32 {$0, $1}, [$2];",
|
|
|
+ "=r,=r,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ return (
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)),
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_global_nc_v4_u32(
|
|
|
+ base_ptr: Int64, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32, Uint32, Uint32]:
|
|
|
+ """Load 128 bits (4 x uint32) from global memory via non-coherent cache (.nc)."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]),
|
|
|
+ [Int64(base_ptr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.global.nc.v4.u32 {$0, $1, $2, $3}, [$4];",
|
|
|
+ "=r,=r,=r,=r,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ return (
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)),
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)),
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), result, [2], loc=loc, ip=ip)),
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), result, [3], loc=loc, ip=ip)),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_global_u64(base_ptr: Int64, value: Uint64, *, loc=None, ip=None):
|
|
|
+ """Store 64 bits to global memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(base_ptr).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint64(value).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.global.u64 [$0], $1;",
|
|
|
+ "l,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_global_u8(base_ptr: Int64, value: Uint8, *, loc=None, ip=None):
|
|
|
+ """Store 8 bits to global memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(base_ptr).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint8(value).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.global.u8 [$0], $1;",
|
|
|
+ "l,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_global_f32(base_ptr: Int64, value: Float32, *, loc=None, ip=None):
|
|
|
+ """Store 32-bit float to global memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(base_ptr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(value).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.global.f32 [$0], $1;",
|
|
|
+ "l,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_global_v2_f32(base_ptr: Int64, v0: Float32, v1: Float32, *, loc=None, ip=None):
|
|
|
+ """Store 64 bits (2 x float32) to global memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(base_ptr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v1).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.global.v2.f32 [$0], {$1, $2};",
|
|
|
+ "l,f,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_global_v4_f32(
|
|
|
+ base_ptr: Int64,
|
|
|
+ v0: Float32,
|
|
|
+ v1: Float32,
|
|
|
+ v2: Float32,
|
|
|
+ v3: Float32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+):
|
|
|
+ """Store 128 bits (4 x float32) to global memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(base_ptr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.global.v4.f32 [$0], {$1, $2, $3, $4};",
|
|
|
+ "l,f,f,f,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_global_v4_f32(
|
|
|
+ base_ptr: Int64,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """Load 128 bits (4 x float32) from global memory."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [Int64(base_ptr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.global.v4.f32 {$0, $1, $2, $3}, [$4];",
|
|
|
+ "=f,=f,=f,=f,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Float32(r0), Float32(r1), Float32(r2), Float32(r3)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_global_v4_u32(
|
|
|
+ base_ptr: Int64,
|
|
|
+ v0: Uint32,
|
|
|
+ v1: Uint32,
|
|
|
+ v2: Uint32,
|
|
|
+ v3: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+):
|
|
|
+ """Store 128 bits (4 x uint32) to global memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(base_ptr).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(v0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(v1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(v2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(v3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.global.v4.u32 [$0], {$1, $2, $3, $4};",
|
|
|
+ "l,r,r,r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def smem_ptr_to_addr(ptr: cute.Pointer, *, loc=None, ip=None) -> Int32:
|
|
|
+ """Convert a generic/smem pointer to a shared-memory u32 address for PTX."""
|
|
|
+ generic_addr = llvm.ptrtoint(T.i64(), ptr.llvm_ptr, loc=loc, ip=ip)
|
|
|
+ return Int32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [generic_addr],
|
|
|
+ "cvta.to.shared.u32 $0, $1;",
|
|
|
+ "=r,l",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def shared_ptr_to_u32(ptr: cute.Pointer, *, loc=None, ip=None) -> Int32:
|
|
|
+ """Convert an address-space-3 shared-memory pointer to a u32 address."""
|
|
|
+ return Int32(llvm.ptrtoint(T.i32(), ptr.llvm_ptr, loc=loc, ip=ip))
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ldmatrix_m8n8x4_b16(smem_addr: Int32, *, loc=None, ip=None) -> Tuple[Uint32, Uint32, Uint32, Uint32]:
|
|
|
+ """Issue `ldmatrix.sync.aligned.m8n8.x4.shared.b16` from a shared-memory byte address."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]),
|
|
|
+ [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];",
|
|
|
+ "=r,=r,=r,=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.i32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.i32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Uint32(r0), Uint32(r1), Uint32(r2), Uint32(r3)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_shared_u16_offset(
|
|
|
+ smem_addr: Int32,
|
|
|
+ byte_offset: int,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Uint16:
|
|
|
+ """Load a shared halfword at a compile-time byte displacement.
|
|
|
+
|
|
|
+ Keeping the result as i16 lets NVVM select a PTX ``%rs`` operand and
|
|
|
+ avoids a redundant ``cvt.u16.u32`` before packed FP8 decode. Embedding the
|
|
|
+ displacement also avoids materializing one address register per unrolled
|
|
|
+ K step.
|
|
|
+ """
|
|
|
+ offset = int(byte_offset)
|
|
|
+ address = f"[$1+{offset}]" if offset else "[$1]"
|
|
|
+ return Uint16(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i16(),
|
|
|
+ [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ f"ld.shared.u16 $0, {address};",
|
|
|
+ "=h,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_shared_u8_offset(
|
|
|
+ smem_addr: Int32,
|
|
|
+ byte_offset: int,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Uint32:
|
|
|
+ """Load one shared byte at a compile-time displacement, zero-extended.
|
|
|
+
|
|
|
+ This is the native ``ld.shared.u8`` form used for UE8M0 scale tables. A
|
|
|
+ common row base plus an immediate displacement keeps the unrolled scale
|
|
|
+ cache free of word-alignment, shift, and mask instructions.
|
|
|
+ """
|
|
|
+ offset = int(byte_offset)
|
|
|
+ address = f"[$1+{offset}]" if offset else "[$1]"
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ f"ld.shared.u8 $0, {address};",
|
|
|
+ "=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ldmatrix_m8n8x2_b16(smem_addr: Int32, *, loc=None, ip=None) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Issue `ldmatrix.sync.aligned.m8n8.x2.shared.b16` from a shared-memory byte address."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ldmatrix.sync.aligned.m8n8.x2.shared.b16 {$0, $1}, [$2];",
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(r0), Uint32(r1)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ldmatrix_m8n8x4_left_half_b16(smem_addr: Int32, *, loc=None, ip=None) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Issue `ldmatrix.sync.aligned.m8n8.x4.shared.b16` and return the left half fragment pair."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, _, $1, _}, [$2];",
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(r0), Uint32(r1)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ldmatrix_m8n8x4_right_half_b16(smem_addr: Int32, *, loc=None, ip=None) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Issue `ldmatrix.sync.aligned.m8n8.x4.shared.b16` and return the right half fragment pair."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {_, $0, _, $1}, [$2];",
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(r0), Uint32(r1)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ldmatrix_m8n8x4_trans_b16(
|
|
|
+ smem_addr: Int32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32, Uint32, Uint32]:
|
|
|
+ """Issue `ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16` from a shared-memory byte address."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]),
|
|
|
+ [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];",
|
|
|
+ "=r,=r,=r,=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.i32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.i32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Uint32(r0), Uint32(r1), Uint32(r2), Uint32(r3)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ldmatrix_m8n8x4_trans_left_half_b16(smem_addr: Int32, *, loc=None, ip=None) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Issue transposed `ldmatrix` and return the left half fragment pair."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {$0, $1, _, _}, [$2];",
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(r0), Uint32(r1)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ldmatrix_m8n8x4_trans_right_half_b16(smem_addr: Int32, *, loc=None, ip=None) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Issue transposed `ldmatrix` and return the right half fragment pair."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {_, _, $0, $1}, [$2];",
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(r0), Uint32(r1)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_shared_v4_u32(smem_addr: Int32, *, loc=None, ip=None) -> Tuple[Uint32, Uint32, Uint32, Uint32]:
|
|
|
+ """Load 128 bits (4 x uint32) from shared memory. smem_addr is a u32 shared-memory address."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]),
|
|
|
+ [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.shared.v4.u32 {$0, $1, $2, $3}, [$4];",
|
|
|
+ "=r,=r,=r,=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.i32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.i32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Uint32(r0), Uint32(r1), Uint32(r2), Uint32(r3)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_shared_v2_u32(smem_addr: Int32, *, loc=None, ip=None) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Load 64 bits (2 x uint32) from shared memory."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Int32(smem_addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.shared.v2.u32 {$0, $1}, [$2];",
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(r0), Uint32(r1)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_shared_u64(smem_addr: Int32, value: Uint64, *, loc=None, ip=None):
|
|
|
+ """Store 64 bits to shared memory. smem_addr is a u32 shared-memory address."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint64(value).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.shared.u64 [$0], $1;",
|
|
|
+ "r,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_shared_v4_u32(
|
|
|
+ smem_addr: Int32,
|
|
|
+ v0: Uint32,
|
|
|
+ v1: Uint32,
|
|
|
+ v2: Uint32,
|
|
|
+ v3: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+):
|
|
|
+ """Store 128 bits (4 x uint32) to shared memory. smem_addr is a u32 shared-memory address."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(v0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(v1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(v2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(v3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.shared.v4.u32 [$0], {$1, $2, $3, $4};",
|
|
|
+ "r,r,r,r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cp_async_bulk_g2s_mbar(
|
|
|
+ smem_dst_u32: Int32,
|
|
|
+ gmem_src_i64: Int64,
|
|
|
+ nbytes: Int32,
|
|
|
+ mbar_u32: Int32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> None:
|
|
|
+ """Emit cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.
|
|
|
+
|
|
|
+ CTA-scope bulk g2s gather completing on a CTA-scope mbarrier. Mirrors
|
|
|
+ FlashInfer cp_async_bulk_g2s (sparse_mla_sm120/arch/cp_async.cuh:64-72)
|
|
|
+ exactly:
|
|
|
+ [$0]=smem dst (u32 shared addr), [$1]=gmem src (i64 generic),
|
|
|
+ $2=bytes (u32), [$3]=mbar (u32 shared addr).
|
|
|
+ Used by the IO warp of the DSV4/GLM sparse-MLA decode port to gather a
|
|
|
+ whole token row (e.g. 656B or 576B+8B) in one transaction.
|
|
|
+ """
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_dst_u32).ir_value(loc=loc, ip=ip),
|
|
|
+ Int64(gmem_src_i64).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(nbytes).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(mbar_u32).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes"
|
|
|
+ " [$0], [$1], $2, [$3];",
|
|
|
+ "r,l,r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def create_l2_evict_first_policy(*, loc=None, ip=None) -> Uint64:
|
|
|
+ """Create an L2 evict-first policy for one-pass streaming cache rows."""
|
|
|
+ return Uint64(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i64(),
|
|
|
+ [],
|
|
|
+ "createpolicy.fractional.L2::evict_first.b64 $0, 1.0;",
|
|
|
+ "=l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cp_async_bulk_g2s_mbar_l2hint(
|
|
|
+ smem_dst_u32: Int32,
|
|
|
+ gmem_src_i64: Int64,
|
|
|
+ nbytes: Int32,
|
|
|
+ mbar_u32: Int32,
|
|
|
+ cache_policy: Uint64,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> None:
|
|
|
+ """Bulk global-to-shared copy with an explicit L2 cache policy."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_dst_u32).ir_value(loc=loc, ip=ip),
|
|
|
+ Int64(gmem_src_i64).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(nbytes).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(mbar_u32).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint64(cache_policy).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint"
|
|
|
+ " [$0], [$1], $2, [$3], $4;",
|
|
|
+ "r,l,r,r,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_shared_u8(smem_addr: Int32, value: Uint8, *, loc=None, ip=None):
|
|
|
+ """Store 8 bits to shared memory. smem_addr is a u32 shared-memory address."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint8(value).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.shared.u8 [$0], $1;",
|
|
|
+ "r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_shared_u16(smem_addr: Int32, value: Uint32, *, loc=None, ip=None):
|
|
|
+ """Store the low 16 bits of a register to shared memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(value).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.shared.u16 [$0], $1;",
|
|
|
+ "r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def quantize_scaled_store_shared_v2_e4m3(
|
|
|
+ smem_addr: Int32,
|
|
|
+ scaled_value0: Float32,
|
|
|
+ scaled_value1: Float32,
|
|
|
+ inv_scale: Float32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+):
|
|
|
+ """Scale, quantize, and vector-store two contiguous E4M3 bytes.
|
|
|
+
|
|
|
+ ``scaled_value*`` already include the per-candidate V scale used by the
|
|
|
+ preceding atomic-max pass. Keeping that multiplication in ordinary SSA
|
|
|
+ lets LLVM retain/reuse the products across the barrier, as FlashInfer does.
|
|
|
+ The E4M3 results never widen through ``Uint32`` and are written with the
|
|
|
+ same half-register ``st.shared.v2.b8`` form.
|
|
|
+ """
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(scaled_value0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(scaled_value1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(inv_scale).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 q0, q1;
|
|
|
+ .reg .f32 zero, tmp;
|
|
|
+ mov.f32 zero, 0f00000000;
|
|
|
+ mul.f32 tmp, $1, $3;
|
|
|
+ cvt.rn.satfinite.e4m3x2.f32 q0, zero, tmp;
|
|
|
+ mul.f32 tmp, $2, $3;
|
|
|
+ cvt.rn.satfinite.e4m3x2.f32 q1, zero, tmp;
|
|
|
+ st.shared.v2.b8 [$0], {q0, q1};
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "r,f,f,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cp_async_u32_shared_global(smem_addr: Int32, gmem_addr: Int64, *, loc=None, ip=None):
|
|
|
+ """4-byte `cp.async.ca.shared.global` copy."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int64(gmem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "cp.async.ca.shared.global [$0], [$1], 4;",
|
|
|
+ "r,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cp_async_u64_shared_global(smem_addr: Int32, gmem_addr: Int64, *, loc=None, ip=None):
|
|
|
+ """8-byte `cp.async.ca.shared.global` copy."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int64(gmem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "cp.async.ca.shared.global [$0], [$1], 8;",
|
|
|
+ "r,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cp_async4_shared_global(smem_addr: Int32, gmem_addr: Int64, *, loc=None, ip=None):
|
|
|
+ """16-byte `cp.async.cg.shared.global` copy."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int64(gmem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "cp.async.cg.shared.global [$0], [$1], 16;",
|
|
|
+ "r,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cp_async4_shared_global_pred(
|
|
|
+ smem_addr: Int32, gmem_addr: Int64, pred: Int32, *, loc=None, ip=None
|
|
|
+):
|
|
|
+ """Predicated 16-byte `cp.async.cg.shared.global` copy."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(pred).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int64(gmem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "{ .reg .pred p; setp.ne.b32 p, $0, 0; @p cp.async.cg.shared.global [$1], [$2], 16; }",
|
|
|
+ "r,r,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cp_async4_ca_shared_global_pred(
|
|
|
+ smem_addr: Int32, gmem_addr: Int64, pred: Int32, *, loc=None, ip=None
|
|
|
+):
|
|
|
+ """Predicated 16-byte `cp.async.ca.shared.global` copy."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(pred).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int64(gmem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "{ .reg .pred p; setp.ne.b32 p, $0, 0; @p cp.async.ca.shared.global [$1], [$2], 16; }",
|
|
|
+ "r,r,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def get_ptr_as_int64(tensor: cute.Tensor, offset, *, loc=None, ip=None) -> Int64:
|
|
|
+ """Get the memory address of tensor[offset] as Int64."""
|
|
|
+ elem_ptr = tensor.iterator + offset
|
|
|
+ ptr_int = llvm.ptrtoint(T.i64(), elem_ptr.llvm_ptr, loc=loc, ip=ip)
|
|
|
+ return Int64(ptr_int)
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# PTX Intrinsics - Global Atomics
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def atomic_add_global_i32(addr: Int64, val: Int32, *, loc=None, ip=None) -> Int32:
|
|
|
+ """Global memory int32 atomic add. Returns old value."""
|
|
|
+ return Int32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "atom.global.add.s32 $0, [$1], $2;",
|
|
|
+ "=r,l,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def red_add_global_release_i32(addr: Int64, val: Int32, *, loc=None, ip=None):
|
|
|
+ """No-return global int32 add with a GPU-scope release fence."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "fence.acq_rel.gpu;\n"
|
|
|
+ "red.relaxed.gpu.global.add.s32 [$0], $1;",
|
|
|
+ "l,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def red_add_global_f32(addr: Int64, val: Float32, *, loc=None, ip=None):
|
|
|
+ """No-return global fp32 add reduction (relaxed device scope)."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "red.global.add.f32 [$0], $1;",
|
|
|
+ "l,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_bf16x2_to_f16x2(packed: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Convert a u32 holding two bf16 (lo, hi) into an f16x2 u32 (lo, hi)."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(packed).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 blo, bhi, hlo, hhi;
|
|
|
+ .reg .f32 flo, fhi;
|
|
|
+ mov.b32 {blo, bhi}, $1;
|
|
|
+ cvt.f32.bf16 flo, blo;
|
|
|
+ cvt.f32.bf16 fhi, bhi;
|
|
|
+ cvt.rn.f16.f32 hlo, flo;
|
|
|
+ cvt.rn.f16.f32 hhi, fhi;
|
|
|
+ mov.b32 $0, {hlo, hhi};
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def red_add_global_bf16x2(addr: Int64, packed: Uint32, *, loc=None, ip=None):
|
|
|
+ """No-return global atomic add of a packed bf16x2 value (2 contiguous bf16).
|
|
|
+
|
|
|
+ Used by the W4A16 TC-decode FC2 epilogue to fold the per-route partial
|
|
|
+ outputs into the per-token output without a separate top-k-sum launch.
|
|
|
+ The address must be 4-byte aligned and cover two consecutive bf16 lanes.
|
|
|
+ """
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(packed).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "red.relaxed.gpu.global.add.noftz.bf16x2 [$0], $1;",
|
|
|
+ "l,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def red_max_global_f32_nonnegative(addr: Int64, val: Float32, *, loc=None, ip=None):
|
|
|
+ """No-return global max reduction for a non-negative fp32 value.
|
|
|
+
|
|
|
+ Non-negative IEEE-754 float ordering matches unsigned integer bit ordering,
|
|
|
+ so this emits a no-return u32 max reduction on the value's bit pattern.
|
|
|
+ Callers must initialize the destination with a non-negative fp32 value.
|
|
|
+ """
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "{ .reg .u32 vi; mov.b32 vi, $1; red.relaxed.gpu.global.max.u32 [$0], vi; }",
|
|
|
+ "l,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def atomic_cas_global_i32(addr: Int64, compare: Int32, value: Int32, *, loc=None, ip=None) -> Int32:
|
|
|
+ """Global memory int32 atomic compare-and-swap. Returns old value."""
|
|
|
+ return Int32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(compare).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(value).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "atom.global.cas.b32 $0, [$1], $2, $3;",
|
|
|
+ "=r,l,r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def atomic_add_shared_i32(addr: Int32, val: Int32, *, loc=None, ip=None) -> Int32:
|
|
|
+ """Shared-memory int32 atomic add (CTA-scope). Returns old value.
|
|
|
+
|
|
|
+ Uses ``atom.shared.add.s32`` with a 32-bit shared-memory address
|
|
|
+ (the native address width for smem on NVIDIA GPUs).
|
|
|
+ """
|
|
|
+ return Int32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [
|
|
|
+ Int32(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "atom.shared.add.s32 $0, [$1], $2;",
|
|
|
+ "=r,r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def atomic_max_shared_f32(smem_addr: Int32, val: Float32, *, loc=None, ip=None):
|
|
|
+ """No-return shared atomic max for a NON-NEGATIVE fp32 ``val``.
|
|
|
+
|
|
|
+ Relies on the IEEE-754 ordering of non-negative floats matching the signed
|
|
|
+ int ordering of their bit patterns, so the smem max can be done with a
|
|
|
+ single s32 reduction. All callers consume the reduced shared slot after a
|
|
|
+ barrier; none needs the old or resulting value in the issuing lane.
|
|
|
+ """
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "{.reg .s32 vi; mov.b32 vi, $1; red.shared.max.s32 [$0], vi;}",
|
|
|
+ "r,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def atomic_max_shared_f32_offset(
|
|
|
+ smem_addr: Int32,
|
|
|
+ byte_offset: int,
|
|
|
+ val: Float32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+):
|
|
|
+ """No-return nonnegative FP32 shared max at a constant displacement."""
|
|
|
+ offset = int(byte_offset)
|
|
|
+ address = f"[$0+{offset}]" if offset else "[$0]"
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(smem_addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ f"{{.reg .s32 vi; mov.b32 vi, $1; red.shared.max.s32 {address}, vi;}}",
|
|
|
+ "r,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def red_add_shared_i32(addr: Int32, val: Int32, *, loc=None, ip=None):
|
|
|
+ """No-return shared-memory int32 add.
|
|
|
+
|
|
|
+ Use this when the old value is not needed. NVCC lowers equivalent C++
|
|
|
+ histogram increments to cheaper no-return shared atomics.
|
|
|
+ """
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "red.shared.add.u32 [$0], $1;",
|
|
|
+ "r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_shared_i32(addr: Int32, *, loc=None, ip=None) -> Int32:
|
|
|
+ """Volatile-style load int32 from shared memory at a 32-bit byte address."""
|
|
|
+ return Int32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Int32(addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.shared.s32 $0, [$1];",
|
|
|
+ "=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_shared_i32_relaxed(addr: Int32, *, loc=None, ip=None) -> Int32:
|
|
|
+ """Load int32 from shared memory at a 32-bit byte address."""
|
|
|
+ return Int32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Int32(addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.shared.s32 $0, [$1];",
|
|
|
+ "=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_shared_u32(addr: Int32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Load uint32 from shared memory at a 32-bit byte address."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Int32(addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.shared.u32 $0, [$1];",
|
|
|
+ "=r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_shared_i32(addr: Int32, val: Int32, *, loc=None, ip=None):
|
|
|
+ """Store int32 to shared memory at a 32-bit byte address."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.shared.s32 [$0], $1;",
|
|
|
+ "r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_shared_u32(addr: Int32, val: Uint32, *, loc=None, ip=None):
|
|
|
+ """Store uint32 to shared memory at a 32-bit byte address."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.shared.u32 [$0], $1;",
|
|
|
+ "r,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_shared_f32(addr: Int32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Load float32 from shared memory at a 32-bit byte address."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Int32(addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.shared.f32 $0, [$1];",
|
|
|
+ "=f,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_shared_f32_offset(
|
|
|
+ addr: Int32,
|
|
|
+ byte_offset: int,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Float32:
|
|
|
+ """Load shared FP32 using a compile-time displacement."""
|
|
|
+ offset = int(byte_offset)
|
|
|
+ address = f"[$1+{offset}]" if offset else "[$1]"
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Int32(addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ f"ld.shared.f32 $0, {address};",
|
|
|
+ "=f,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_shared_v4_f32(
|
|
|
+ addr: Int32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """Load 128 bits (4 x float32) from shared memory."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [Int32(addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];",
|
|
|
+ "=f,=f,=f,=f,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Float32(r0), Float32(r1), Float32(r2), Float32(r3)
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_shared_bf16_to_f32(addr: Int32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Load a BF16 from shared memory at a 32-bit byte address and convert to FP32."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Int32(addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "{.reg .b16 tmp; ld.shared.b16 tmp, [$1]; cvt.f32.bf16 $0, tmp;}",
|
|
|
+ "=f,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_shared_f32(addr: Int32, val: Float32, *, loc=None, ip=None):
|
|
|
+ """Store float32 to shared memory at a 32-bit byte address."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.shared.f32 [$0], $1;",
|
|
|
+ "r,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_shared_f32_offset(
|
|
|
+ addr: Int32,
|
|
|
+ byte_offset: int,
|
|
|
+ val: Float32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+):
|
|
|
+ """Store shared FP32 using a compile-time displacement."""
|
|
|
+ offset = int(byte_offset)
|
|
|
+ address = f"[$0+{offset}]" if offset else "[$0]"
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ f"st.shared.f32 {address}, $1;",
|
|
|
+ "r,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_shared_v4_f32(
|
|
|
+ addr: Int32,
|
|
|
+ v0: Float32,
|
|
|
+ v1: Float32,
|
|
|
+ v2: Float32,
|
|
|
+ v3: Float32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+):
|
|
|
+ """Store 128 bits (4 x float32) to shared memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.shared.v4.f32 [$0], {$1, $2, $3, $4};",
|
|
|
+ "r,f,f,f,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_shared_bf16_from_f32(addr: Int32, val: Float32, *, loc=None, ip=None):
|
|
|
+ """Convert float32 to BF16 and store one 16-bit value to shared memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "{ .reg .b16 tmp; cvt.rn.bf16.f32 tmp, $1; st.shared.b16 [$0], tmp; }",
|
|
|
+ "r,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_shared_f16_from_f32(addr: Int32, val: Float32, *, loc=None, ip=None):
|
|
|
+ """Convert float32 to FP16 and store one 16-bit value to shared memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int32(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "{ .reg .b16 tmp; cvt.rn.f16.f32 tmp, $1; st.shared.b16 [$0], tmp; }",
|
|
|
+ "r,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_global_i32(addr: Int64, val: Int32, *, loc=None, ip=None):
|
|
|
+ """Store int32 to global memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.global.s32 [$0], $1;",
|
|
|
+ "l,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# PTX Intrinsics - Global Memory Barriers
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ld_global_acquire_i32(addr: Int64, *, loc=None, ip=None) -> Int32:
|
|
|
+ """Load int32 from global memory with acquire semantics."""
|
|
|
+ return Int32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Int64(addr).ir_value(loc=loc, ip=ip)],
|
|
|
+ "ld.global.acquire.gpu.s32 $0, [$1];",
|
|
|
+ "=r,l",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def st_global_release_i32(addr: Int64, val: Int32, *, loc=None, ip=None):
|
|
|
+ """Store int32 to global memory with release semantics."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(val).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "st.global.release.gpu.s32 [$0], $1;",
|
|
|
+ "l,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def spin_wait_global_eq_i32(addr: Int64, expected: Int32, *, loc=None, ip=None):
|
|
|
+ """Spin-wait until *addr != expected (acquire semantics on load)."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(expected).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "{\n"
|
|
|
+ ".reg .pred %p0;\n"
|
|
|
+ ".reg .s32 %val;\n"
|
|
|
+ "spin_loop:\n"
|
|
|
+ " ld.global.acquire.gpu.s32 %val, [$0];\n"
|
|
|
+ " setp.eq.s32 %p0, %val, $1;\n"
|
|
|
+ " @%p0 bra spin_loop;\n"
|
|
|
+ "}",
|
|
|
+ "l,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def spin_wait_global_ge_i32(addr: Int64, target: Int32, *, loc=None, ip=None):
|
|
|
+ """Spin-wait until *addr >= target using acquire loads."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(target).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "{\n"
|
|
|
+ ".reg .pred %p0;\n"
|
|
|
+ ".reg .s32 %val;\n"
|
|
|
+ "spin_loop_ge:\n"
|
|
|
+ " ld.global.acquire.gpu.s32 %val, [$0];\n"
|
|
|
+ " setp.lt.s32 %p0, %val, $1;\n"
|
|
|
+ " @%p0 bra spin_loop_ge;\n"
|
|
|
+ "}",
|
|
|
+ "l,r",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def threadfence(*, loc=None, ip=None):
|
|
|
+ """Emit membar.gl — threadfence across global memory."""
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [],
|
|
|
+ "membar.gl;",
|
|
|
+ "",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# PTX Intrinsics - Scatter Atomics
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def scatter_add_bf16x2(addr: Int64, val0_f32, val1_f32, *, loc=None, ip=None):
|
|
|
+ """BF16x2 atomic reduction add to global memory.
|
|
|
+
|
|
|
+ Packs two f32 values into bf16x2 via cvt.rn.satfinite, then does
|
|
|
+ red.relaxed.gpu.global.add.noftz.bf16x2. Workaround for LLVM NVPTX
|
|
|
+ bug that decomposes vector<2xbf16> into .v2.bf16x2 which ptxas rejects.
|
|
|
+ """
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ val0_f32.ir_value(loc=loc, ip=ip),
|
|
|
+ val1_f32.ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "{ .reg .b32 packed; cvt.rn.satfinite.bf16x2.f32 packed, $2, $1; red.relaxed.gpu.global.add.noftz.bf16x2 [$0], packed; }",
|
|
|
+ "l,f,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def scatter_add_bf16(addr: Int64, val_f32, *, loc=None, ip=None):
|
|
|
+ """BF16 atomic reduction add to global memory.
|
|
|
+
|
|
|
+ Converts the f32 input to bf16 inside PTX and atomically accumulates it into
|
|
|
+ one bf16 output lane. This is intended for opt-in approximate reductions.
|
|
|
+ """
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ val_f32.ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "{ .reg .b16 packed; cvt.rn.satfinite.bf16.f32 packed, $1; red.relaxed.gpu.global.add.noftz.bf16 [$0], packed; }",
|
|
|
+ "l,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def scatter_add_v4_bf16x2(addr: Int64, v0, v1, v2, v3, v4, v5, v6, v7, *, loc=None, ip=None):
|
|
|
+ """Vectorized BF16x2 atomic reduction: 8 bf16 values (16 bytes) in one go.
|
|
|
+
|
|
|
+ red.global.add.noftz.v4.bf16x2 [addr], {packed0, packed1, packed2, packed3}
|
|
|
+ Each packed reg contains 2 bf16 values converted from f32 pairs.
|
|
|
+ """
|
|
|
+ llvm.inline_asm(
|
|
|
+ None,
|
|
|
+ [
|
|
|
+ Int64(addr).ir_value(loc=loc, ip=ip),
|
|
|
+ v0.ir_value(loc=loc, ip=ip), v1.ir_value(loc=loc, ip=ip),
|
|
|
+ v2.ir_value(loc=loc, ip=ip), v3.ir_value(loc=loc, ip=ip),
|
|
|
+ v4.ir_value(loc=loc, ip=ip), v5.ir_value(loc=loc, ip=ip),
|
|
|
+ v6.ir_value(loc=loc, ip=ip), v7.ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "{ .reg .b32 p0,p1,p2,p3;"
|
|
|
+ " cvt.rn.satfinite.bf16x2.f32 p0, $2, $1;"
|
|
|
+ " cvt.rn.satfinite.bf16x2.f32 p1, $4, $3;"
|
|
|
+ " cvt.rn.satfinite.bf16x2.f32 p2, $6, $5;"
|
|
|
+ " cvt.rn.satfinite.bf16x2.f32 p3, $8, $7;"
|
|
|
+ " red.global.add.noftz.v4.bf16x2 [$0], {p0, p1, p2, p3}; }",
|
|
|
+ "l,f,f,f,f,f,f,f,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# PTX Intrinsics - Math Operations
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def rcp_approx_ftz(a: Float32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Fast reciprocal using PTX rcp.approx.ftz.f32."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Float32(a).ir_value(loc=loc, ip=ip)],
|
|
|
+ "rcp.approx.ftz.f32 $0, $1;",
|
|
|
+ "=f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def fmin_f32(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Compute min of two float32 values using PTX min.f32."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)],
|
|
|
+ "min.f32 $0, $1, $2;",
|
|
|
+ "=f,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def fmax_f32(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Compute max of two float32 values using PTX max.f32."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)],
|
|
|
+ "max.f32 $0, $1, $2;",
|
|
|
+ "=f,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def fabs_f32(a: Float32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Compute absolute value of float32 using PTX abs.f32."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Float32(a).ir_value(loc=loc, ip=ip)],
|
|
|
+ "abs.f32 $0, $1;",
|
|
|
+ "=f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# Half2 SIMD Intrinsics
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def half2_mul(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Multiply two Half2 values element-wise: (a.x*b.x, a.y*b.y)."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(a).ir_value(loc=loc, ip=ip), Uint32(b).ir_value(loc=loc, ip=ip)],
|
|
|
+ "mul.f16x2 $0, $1, $2;",
|
|
|
+ "=r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def broadcast_f32_to_half2(x: Float32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Pack one float32 value into both lanes of an f16x2 register."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Float32(x).ir_value(loc=loc, ip=ip)],
|
|
|
+ "cvt.rn.f16x2.f32 $0, $1, $1;",
|
|
|
+ "=r,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def hadd2(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Add two Half2 values element-wise: (a.x+b.x, a.y+b.y)."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(a).ir_value(loc=loc, ip=ip), Uint32(b).ir_value(loc=loc, ip=ip)],
|
|
|
+ "add.f16x2 $0, $1, $2;",
|
|
|
+ "=r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def habs2(x: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Half2 absolute value - clears sign bits of both fp16 values."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(x).ir_value(loc=loc, ip=ip)],
|
|
|
+ "and.b32 $0, $1, 0x7FFF7FFF;",
|
|
|
+ "=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def hmax2(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Half2 max - element-wise max of 2 fp16 pairs."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(a).ir_value(loc=loc, ip=ip), Uint32(b).ir_value(loc=loc, ip=ip)],
|
|
|
+ "max.f16x2 $0, $1, $2;",
|
|
|
+ "=r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def hmax_to_f32(x: Uint32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Extract max of 2 fp16 values in half2 as float32."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Uint32(x).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 h0, h1;
|
|
|
+ .reg .f32 f0, f1;
|
|
|
+ mov.b32 {h0, h1}, $1;
|
|
|
+ cvt.f32.f16 f0, h0;
|
|
|
+ cvt.f32.f16 f1, h1;
|
|
|
+ max.f32 $0, f0, f1;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def half2_to_float2_scaled(
|
|
|
+ h2: Uint32, scale: Float32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Float32, Float32]:
|
|
|
+ """Convert half2 to float2 AND multiply by scale."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32()]),
|
|
|
+ [Uint32(h2).ir_value(loc=loc, ip=ip), Float32(scale).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 h0, h1;
|
|
|
+ .reg .f32 f0, f1;
|
|
|
+ mov.b32 {h0, h1}, $2;
|
|
|
+ cvt.f32.f16 f0, h0;
|
|
|
+ cvt.f32.f16 f1, h1;
|
|
|
+ mul.f32 $0, f0, $3;
|
|
|
+ mul.f32 $1, f1, $3;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,=f,r,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+ f0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ f1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+
|
|
|
+ return Float32(f0), Float32(f1)
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# BFloat2 SIMD Intrinsics
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def bfloat2_mul(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Multiply two BFloat2 values element-wise: (a.x*b.x, a.y*b.y)."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(a).ir_value(loc=loc, ip=ip), Uint32(b).ir_value(loc=loc, ip=ip)],
|
|
|
+ "mul.bf16x2 $0, $1, $2;",
|
|
|
+ "=r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def bfloat2_broadcast_lane(x: Uint32, lane: Int32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Duplicate one BF16 lane from a packed bf16x2 register into both lanes."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [
|
|
|
+ Uint32(x).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(lane).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .pred p;
|
|
|
+ .reg .b32 lo, hi, val, shifted;
|
|
|
+ and.b32 lo, $1, 0x0000ffff;
|
|
|
+ shr.u32 hi, $1, 16;
|
|
|
+ setp.eq.s32 p, $2, 0;
|
|
|
+ @p mov.b32 val, lo;
|
|
|
+ @!p mov.b32 val, hi;
|
|
|
+ shl.b32 shifted, val, 16;
|
|
|
+ or.b32 $0, val, shifted;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def broadcast_f32_to_bfloat2(x: Float32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Pack one float32 value into both lanes of a bf16x2 register."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Float32(x).ir_value(loc=loc, ip=ip)],
|
|
|
+ "cvt.rn.satfinite.bf16x2.f32 $0, $1, $1;",
|
|
|
+ "=r,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def pack_f32x2_to_bfloat2(x0: Float32, x1: Float32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Pack 2 float32 values into one bf16x2 register."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Float32(x0).ir_value(loc=loc, ip=ip), Float32(x1).ir_value(loc=loc, ip=ip)],
|
|
|
+ "cvt.rn.satfinite.bf16x2.f32 $0, $2, $1;",
|
|
|
+ "=r,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_bf16x2_to_e4m3x2(src: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Convert packed bf16x2 to packed e4m3x2 in the low 16 bits of a u32."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(src).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 out, zero;
|
|
|
+ cvt.rn.satfinite.e4m3x2.bf16x2 out, $1;
|
|
|
+ mov.u16 zero, 0;
|
|
|
+ mov.b32 $0, {out, zero};
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_bf16x2x2_to_e4m3x4(lo: Uint32, hi: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Convert two bf16x2 values and pack the e4m3x2 results into one u32."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(lo).ir_value(loc=loc, ip=ip), Uint32(hi).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 out_lo, out_hi;
|
|
|
+ cvt.rn.satfinite.e4m3x2.bf16x2 out_lo, $1;
|
|
|
+ cvt.rn.satfinite.e4m3x2.bf16x2 out_hi, $2;
|
|
|
+ mov.b32 $0, {out_lo, out_hi};
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def bfloat2_add(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Add two BFloat2 values element-wise: (a.x+b.x, a.y+b.y)."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(a).ir_value(loc=loc, ip=ip), Uint32(b).ir_value(loc=loc, ip=ip)],
|
|
|
+ "add.bf16x2 $0, $1, $2;",
|
|
|
+ "=r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def fp8x4_e4m3_to_bfloat2x2(packed: Uint32, *, loc=None, ip=None) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Widen 4 packed E4M3 bytes into 2 packed bf16x2 registers exactly."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Uint32(packed).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 low;
|
|
|
+ .reg .b32 q, out0, out1, sign0, sign1, mant0, mant1, tmp, bias, bias_f32;
|
|
|
+ prmt.b32 q, $2, 0, 0x1302;
|
|
|
+
|
|
|
+ and.b32 sign0, q, 0x80008000;
|
|
|
+ shr.u32 tmp, q, 4;
|
|
|
+ and.b32 mant0, tmp, 0x07F007F0;
|
|
|
+ or.b32 out0, mant0, sign0;
|
|
|
+
|
|
|
+ shl.b32 tmp, q, 8;
|
|
|
+ and.b32 sign1, tmp, 0x80008000;
|
|
|
+ shl.b32 tmp, q, 4;
|
|
|
+ and.b32 mant1, tmp, 0x07F007F0;
|
|
|
+ or.b32 out1, mant1, sign1;
|
|
|
+
|
|
|
+ mov.b32 bias_f32, 0x7B800000;
|
|
|
+ cvt.rn.bf16.f32 low, bias_f32;
|
|
|
+ mov.b32 bias, {low, low};
|
|
|
+
|
|
|
+ mul.bf16x2 $0, out0, bias;
|
|
|
+ mul.bf16x2 $1, out1, bias;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ lo = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ hi = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(lo), Uint32(hi)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def packed_dequant_e2m1x4_to_bfloat2x2(
|
|
|
+ packed: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """FE2M1 -> BF16 register dequant for one packed 4-value fragment."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Uint32(packed).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b32 q, out1, out2, tmp;
|
|
|
+
|
|
|
+ and.b32 out1, $2, 0x80008000;
|
|
|
+ and.b32 tmp, $2, 0x70007000;
|
|
|
+ shr.u32 tmp, tmp, 6;
|
|
|
+ or.b32 out1, out1, tmp;
|
|
|
+
|
|
|
+ shl.b32 q, $2, 4;
|
|
|
+ and.b32 out2, q, 0x80008000;
|
|
|
+ and.b32 tmp, q, 0x70007000;
|
|
|
+ shr.u32 tmp, tmp, 6;
|
|
|
+ or.b32 out2, out2, tmp;
|
|
|
+
|
|
|
+ mov.b32 $0, out2;
|
|
|
+ mov.b32 $1, out1;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ lo = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ hi = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(lo), Uint32(hi)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def packed_dequant_e2m1x4_to_half2x2(
|
|
|
+ packed: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """FE2M1 -> FP16 register dequant for one packed 4-value fragment."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Uint32(packed).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b32 q, out1, out2, tmp;
|
|
|
+
|
|
|
+ and.b32 out1, $2, 0x80008000;
|
|
|
+ and.b32 tmp, $2, 0x70007000;
|
|
|
+ shr.u32 tmp, tmp, 3;
|
|
|
+ or.b32 out1, out1, tmp;
|
|
|
+
|
|
|
+ shl.b32 q, $2, 4;
|
|
|
+ and.b32 out2, q, 0x80008000;
|
|
|
+ and.b32 tmp, q, 0x70007000;
|
|
|
+ shr.u32 tmp, tmp, 3;
|
|
|
+ or.b32 out2, out2, tmp;
|
|
|
+
|
|
|
+ mov.b32 $0, out2;
|
|
|
+ mov.b32 $1, out1;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ lo = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ hi = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(lo), Uint32(hi)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def packed_dequant_e4m3x4_to_bfloat2x2(
|
|
|
+ packed: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """FE4M3 scale dequant for one packed 4-value BF16 fragment."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Uint32(packed).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b32 q, out1, out2, tmp;
|
|
|
+
|
|
|
+ and.b32 tmp, $2, 0x80008000;
|
|
|
+ shr.u32 out1, tmp, 1;
|
|
|
+ and.b32 tmp, $2, 0x7F007F00;
|
|
|
+ shr.u32 tmp, tmp, 4;
|
|
|
+ or.b32 out1, out1, tmp;
|
|
|
+
|
|
|
+ shl.b32 q, $2, 8;
|
|
|
+ and.b32 tmp, q, 0x80008000;
|
|
|
+ shr.u32 out2, tmp, 1;
|
|
|
+ and.b32 tmp, q, 0x7F007F00;
|
|
|
+ shr.u32 tmp, tmp, 4;
|
|
|
+ or.b32 out2, out2, tmp;
|
|
|
+
|
|
|
+ mov.b32 $0, out2;
|
|
|
+ mov.b32 $1, out1;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ lo = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ hi = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(lo), Uint32(hi)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def packed_dequant_e4m3x4_to_half2x2(
|
|
|
+ packed: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """FE4M3 scale dequant for one packed 4-value FP16 fragment."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Uint32(packed).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b32 q, out1, out2;
|
|
|
+
|
|
|
+ and.b32 out1, $2, 0xFF00FF00;
|
|
|
+ shr.u32 out1, out1, 1;
|
|
|
+
|
|
|
+ shl.b32 q, $2, 8;
|
|
|
+ and.b32 out2, q, 0xFF00FF00;
|
|
|
+ shr.u32 out2, out2, 1;
|
|
|
+
|
|
|
+ mov.b32 $0, out2;
|
|
|
+ mov.b32 $1, out1;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ lo = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ hi = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(lo), Uint32(hi)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def packed_dequant_e8m0x4_to_bfloat2x2(
|
|
|
+ packed: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """E8M0 compute-scale dequant for one packed 4-value BF16 fragment.
|
|
|
+
|
|
|
+ The W4A16 FP4 unpack path represents E2M1 values scaled by 2^-126. Match
|
|
|
+ the existing NVFP4 split by materializing E8M0 scales multiplied by 2^7 in
|
|
|
+ the MMA input and applying the remaining compensation in the kernel
|
|
|
+ epilogue.
|
|
|
+ """
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Uint32(packed).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .u32 b0, b1, b2, b3;
|
|
|
+ .reg .u32 h0, h1, h2, h3;
|
|
|
+ .reg .u32 t0, t1;
|
|
|
+
|
|
|
+ and.b32 b0, $2, 0x000000ff;
|
|
|
+ shr.u32 b1, $2, 8;
|
|
|
+ and.b32 b1, b1, 0x000000ff;
|
|
|
+ shr.u32 b2, $2, 16;
|
|
|
+ and.b32 b2, b2, 0x000000ff;
|
|
|
+ shr.u32 b3, $2, 24;
|
|
|
+
|
|
|
+ add.u32 h0, b0, 7;
|
|
|
+ add.u32 h1, b1, 7;
|
|
|
+ add.u32 h2, b2, 7;
|
|
|
+ add.u32 h3, b3, 7;
|
|
|
+ shl.b32 h0, h0, 7;
|
|
|
+ shl.b32 h1, h1, 7;
|
|
|
+ shl.b32 h2, h2, 7;
|
|
|
+ shl.b32 h3, h3, 7;
|
|
|
+
|
|
|
+ shl.b32 t0, h2, 16;
|
|
|
+ or.b32 $0, h0, t0;
|
|
|
+ shl.b32 t1, h3, 16;
|
|
|
+ or.b32 $1, h1, t1;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ lo = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ hi = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(lo), Uint32(hi)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def packed_dequant_e8m0x4_to_half2x2(
|
|
|
+ packed: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """E8M0 compute-scale dequant for one packed 4-value FP16 fragment."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Uint32(packed).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .pred p0, p1, p2, p3;
|
|
|
+ .reg .u32 b0, b1, b2, b3;
|
|
|
+ .reg .s32 e0, e1, e2, e3;
|
|
|
+ .reg .f32 ef0, ef1, ef2, ef3;
|
|
|
+ .reg .f32 f0, f1, f2, f3;
|
|
|
+ .reg .b16 h0, h1, h2, h3;
|
|
|
+
|
|
|
+ and.b32 b0, $2, 0x000000ff;
|
|
|
+ shr.u32 b1, $2, 8;
|
|
|
+ and.b32 b1, b1, 0x000000ff;
|
|
|
+ shr.u32 b2, $2, 16;
|
|
|
+ and.b32 b2, b2, 0x000000ff;
|
|
|
+ shr.u32 b3, $2, 24;
|
|
|
+
|
|
|
+ setp.eq.u32 p0, b0, 0;
|
|
|
+ setp.eq.u32 p1, b1, 0;
|
|
|
+ setp.eq.u32 p2, b2, 0;
|
|
|
+ setp.eq.u32 p3, b3, 0;
|
|
|
+
|
|
|
+ cvt.s32.u32 e0, b0;
|
|
|
+ cvt.s32.u32 e1, b1;
|
|
|
+ cvt.s32.u32 e2, b2;
|
|
|
+ cvt.s32.u32 e3, b3;
|
|
|
+ sub.s32 e0, e0, 120;
|
|
|
+ sub.s32 e1, e1, 120;
|
|
|
+ sub.s32 e2, e2, 120;
|
|
|
+ sub.s32 e3, e3, 120;
|
|
|
+
|
|
|
+ cvt.rn.f32.s32 ef0, e0;
|
|
|
+ cvt.rn.f32.s32 ef1, e1;
|
|
|
+ cvt.rn.f32.s32 ef2, e2;
|
|
|
+ cvt.rn.f32.s32 ef3, e3;
|
|
|
+ ex2.approx.f32 f0, ef0;
|
|
|
+ ex2.approx.f32 f1, ef1;
|
|
|
+ ex2.approx.f32 f2, ef2;
|
|
|
+ ex2.approx.f32 f3, ef3;
|
|
|
+ selp.f32 f0, 0f00000000, f0, p0;
|
|
|
+ selp.f32 f1, 0f00000000, f1, p1;
|
|
|
+ selp.f32 f2, 0f00000000, f2, p2;
|
|
|
+ selp.f32 f3, 0f00000000, f3, p3;
|
|
|
+
|
|
|
+ cvt.rn.f16.f32 h0, f0;
|
|
|
+ cvt.rn.f16.f32 h1, f1;
|
|
|
+ cvt.rn.f16.f32 h2, f2;
|
|
|
+ cvt.rn.f16.f32 h3, f3;
|
|
|
+
|
|
|
+ setp.eq.u32 p0, b0, 255;
|
|
|
+ setp.eq.u32 p1, b1, 255;
|
|
|
+ setp.eq.u32 p2, b2, 255;
|
|
|
+ setp.eq.u32 p3, b3, 255;
|
|
|
+ selp.b16 h0, 0x7e00, h0, p0;
|
|
|
+ selp.b16 h1, 0x7e00, h1, p1;
|
|
|
+ selp.b16 h2, 0x7e00, h2, p2;
|
|
|
+ selp.b16 h3, 0x7e00, h3, p3;
|
|
|
+
|
|
|
+ mov.b32 $0, {h0, h2};
|
|
|
+ mov.b32 $1, {h1, h3};
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ lo = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
|
|
|
+ hi = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(lo), Uint32(hi)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def bf16_mma_m16n8k16_f32(
|
|
|
+ d0: Float32,
|
|
|
+ d1: Float32,
|
|
|
+ d2: Float32,
|
|
|
+ d3: Float32,
|
|
|
+ a0: Uint32,
|
|
|
+ a1: Uint32,
|
|
|
+ a2: Uint32,
|
|
|
+ a3: Uint32,
|
|
|
+ b0: Uint32,
|
|
|
+ b1: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """Warp MMA helper for `mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32`."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(a0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
|
|
|
+ {$0, $1, $2, $3},
|
|
|
+ {$4, $5, $6, $7},
|
|
|
+ {$8, $9},
|
|
|
+ {$10, $11, $12, $13};
|
|
|
+ """,
|
|
|
+ "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Float32(r0), Float32(r1), Float32(r2), Float32(r3)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def f32_to_raw_bits(x: Float32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Return the raw IEEE-754 bits of one float32 value."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Float32(x).ir_value(loc=loc, ip=ip)],
|
|
|
+ "mov.b32 $0, $1;",
|
|
|
+ "=r,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def f32_to_tf32_bits(x: Float32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Round one float32 value to TF32 format and return its 32-bit operand bits."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Float32(x).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b32 tmp;
|
|
|
+ cvt.rna.tf32.f32 tmp, $1;
|
|
|
+ mov.b32 $0, tmp;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def tf32_mma_m16n8k8_f32(
|
|
|
+ d0: Float32,
|
|
|
+ d1: Float32,
|
|
|
+ d2: Float32,
|
|
|
+ d3: Float32,
|
|
|
+ a0: Uint32,
|
|
|
+ a1: Uint32,
|
|
|
+ a2: Uint32,
|
|
|
+ a3: Uint32,
|
|
|
+ b0: Uint32,
|
|
|
+ b1: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """Warp MMA helper for `mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32`."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(a0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32
|
|
|
+ {$0, $1, $2, $3},
|
|
|
+ {$4, $5, $6, $7},
|
|
|
+ {$8, $9},
|
|
|
+ {$10, $11, $12, $13};
|
|
|
+ """,
|
|
|
+ "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Float32(r0), Float32(r1), Float32(r2), Float32(r3)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def f16_mma_m16n8k16_f32(
|
|
|
+ d0: Float32,
|
|
|
+ d1: Float32,
|
|
|
+ d2: Float32,
|
|
|
+ d3: Float32,
|
|
|
+ a0: Uint32,
|
|
|
+ a1: Uint32,
|
|
|
+ a2: Uint32,
|
|
|
+ a3: Uint32,
|
|
|
+ b0: Uint32,
|
|
|
+ b1: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """Warp MMA helper for `mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32`."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(a0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32
|
|
|
+ {$0, $1, $2, $3},
|
|
|
+ {$4, $5, $6, $7},
|
|
|
+ {$8, $9},
|
|
|
+ {$10, $11, $12, $13};
|
|
|
+ """,
|
|
|
+ "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Float32(r0), Float32(r1), Float32(r2), Float32(r3)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def bf16_mma_rhs_fragments_as_mma_a_m16n8k16_f32(
|
|
|
+ d0: Float32,
|
|
|
+ d1: Float32,
|
|
|
+ d2: Float32,
|
|
|
+ d3: Float32,
|
|
|
+ b0_0: Uint32,
|
|
|
+ b1_0: Uint32,
|
|
|
+ b0_1: Uint32,
|
|
|
+ b1_1: Uint32,
|
|
|
+ a0: Uint32,
|
|
|
+ a1: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """BF16 MMA form used by the routed m-block-size-8 path.
|
|
|
+
|
|
|
+ The dequantized RHS fragments feed the hardware A operand, while the
|
|
|
+ routed activation fragment loaded with `ldmatrix.x2` feeds hardware B.
|
|
|
+ """
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(b0_0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b1_0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b0_1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b1_1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
|
|
|
+ {$0, $1, $2, $3},
|
|
|
+ {$4, $5, $6, $7},
|
|
|
+ {$8, $9},
|
|
|
+ {$10, $11, $12, $13};
|
|
|
+ """,
|
|
|
+ "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Float32(r0), Float32(r1), Float32(r2), Float32(r3)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def f16_mma_rhs_fragments_as_mma_a_m16n8k16_f32(
|
|
|
+ d0: Float32,
|
|
|
+ d1: Float32,
|
|
|
+ d2: Float32,
|
|
|
+ d3: Float32,
|
|
|
+ b0_0: Uint32,
|
|
|
+ b1_0: Uint32,
|
|
|
+ b0_1: Uint32,
|
|
|
+ b1_1: Uint32,
|
|
|
+ a0: Uint32,
|
|
|
+ a1: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """FP16 MMA form used by the routed m-block-size-8 path."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(b0_0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b1_0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b0_1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b1_1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32
|
|
|
+ {$0, $1, $2, $3},
|
|
|
+ {$4, $5, $6, $7},
|
|
|
+ {$8, $9},
|
|
|
+ {$10, $11, $12, $13};
|
|
|
+ """,
|
|
|
+ "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Float32(r0), Float32(r1), Float32(r2), Float32(r3)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def bf16_rowsum_m16k16_f32(
|
|
|
+ d0: Float32,
|
|
|
+ d1: Float32,
|
|
|
+ a0: Uint32,
|
|
|
+ a1: Uint32,
|
|
|
+ a2: Uint32,
|
|
|
+ a3: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32]:
|
|
|
+ """Row-sum helper matching FlashInfer's m16k16 BF16 rowsum MMA."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(a0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(1065369472).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(1065369472).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d1).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
|
|
|
+ {$0, _, $1, _},
|
|
|
+ {$2, $3, $4, $5},
|
|
|
+ {$6, $7},
|
|
|
+ {$0, 0., $1, 0.};
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,=f,r,r,r,r,r,r,0,1",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Float32(r0), Float32(r1)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def bf16_mma_m16n16k16_f32(
|
|
|
+ d0: Float32,
|
|
|
+ d1: Float32,
|
|
|
+ d2: Float32,
|
|
|
+ d3: Float32,
|
|
|
+ d4: Float32,
|
|
|
+ d5: Float32,
|
|
|
+ d6: Float32,
|
|
|
+ d7: Float32,
|
|
|
+ a0: Uint32,
|
|
|
+ a1: Uint32,
|
|
|
+ a2: Uint32,
|
|
|
+ a3: Uint32,
|
|
|
+ b0: Uint32,
|
|
|
+ b1: Uint32,
|
|
|
+ b2: Uint32,
|
|
|
+ b3: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32, Float32, Float32, Float32, Float32]:
|
|
|
+ """Warp MMA helper matching FlashInfer's `m16n16k16` BF16/BF16->F32 wrapper."""
|
|
|
+ result0 = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(a0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
|
|
|
+ {$0, $1, $2, $3},
|
|
|
+ {$4, $5, $6, $7},
|
|
|
+ {$8, $9},
|
|
|
+ {$10, $11, $12, $13};
|
|
|
+ """,
|
|
|
+ "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ result1 = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(a0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b3).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d4).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d5).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d6).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d7).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
|
|
|
+ {$0, $1, $2, $3},
|
|
|
+ {$4, $5, $6, $7},
|
|
|
+ {$8, $9},
|
|
|
+ {$10, $11, $12, $13};
|
|
|
+ """,
|
|
|
+ "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
|
|
|
+ has_side_effects=True,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.f32(), result0, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.f32(), result0, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.f32(), result0, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.f32(), result0, [3], loc=loc, ip=ip)
|
|
|
+ r4 = llvm.extractvalue(T.f32(), result1, [0], loc=loc, ip=ip)
|
|
|
+ r5 = llvm.extractvalue(T.f32(), result1, [1], loc=loc, ip=ip)
|
|
|
+ r6 = llvm.extractvalue(T.f32(), result1, [2], loc=loc, ip=ip)
|
|
|
+ r7 = llvm.extractvalue(T.f32(), result1, [3], loc=loc, ip=ip)
|
|
|
+ return (
|
|
|
+ Float32(r0),
|
|
|
+ Float32(r1),
|
|
|
+ Float32(r2),
|
|
|
+ Float32(r3),
|
|
|
+ Float32(r4),
|
|
|
+ Float32(r5),
|
|
|
+ Float32(r6),
|
|
|
+ Float32(r7),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def mxfp8_mma_m16n8k32_f32_e4m3(
|
|
|
+ d0: Float32,
|
|
|
+ d1: Float32,
|
|
|
+ d2: Float32,
|
|
|
+ d3: Float32,
|
|
|
+ a0: Uint32,
|
|
|
+ a1: Uint32,
|
|
|
+ a2: Uint32,
|
|
|
+ a3: Uint32,
|
|
|
+ b0: Uint32,
|
|
|
+ b1: Uint32,
|
|
|
+ sfa: Uint32,
|
|
|
+ sfb: Uint32,
|
|
|
+ bid_a: int = 0,
|
|
|
+ tid_a: int = 0,
|
|
|
+ bid_b: int = 0,
|
|
|
+ tid_b: int = 0,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """Warp MMA helper for SM120 MXFP8 block-scaled `m16n8k32` E4M3/E4M3."""
|
|
|
+ i16_ty = cutlass._mlir.ir.IntegerType.get_signless(16)
|
|
|
+ bid_a_i16 = cutlass._mlir.ir.Operation.create(
|
|
|
+ "llvm.mlir.constant",
|
|
|
+ results=[i16_ty],
|
|
|
+ attributes={"value": cutlass._mlir.ir.IntegerAttr.get(i16_ty, int(bid_a))},
|
|
|
+ ).result
|
|
|
+ tid_a_i16 = cutlass._mlir.ir.Operation.create(
|
|
|
+ "llvm.mlir.constant",
|
|
|
+ results=[i16_ty],
|
|
|
+ attributes={"value": cutlass._mlir.ir.IntegerAttr.get(i16_ty, int(tid_a))},
|
|
|
+ ).result
|
|
|
+ bid_b_i16 = cutlass._mlir.ir.Operation.create(
|
|
|
+ "llvm.mlir.constant",
|
|
|
+ results=[i16_ty],
|
|
|
+ attributes={"value": cutlass._mlir.ir.IntegerAttr.get(i16_ty, int(bid_b))},
|
|
|
+ ).result
|
|
|
+ tid_b_i16 = cutlass._mlir.ir.Operation.create(
|
|
|
+ "llvm.mlir.constant",
|
|
|
+ results=[i16_ty],
|
|
|
+ attributes={"value": cutlass._mlir.ir.IntegerAttr.get(i16_ty, int(tid_b))},
|
|
|
+ ).result
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(a0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(sfa).ir_value(loc=loc, ip=ip),
|
|
|
+ bid_a_i16,
|
|
|
+ tid_a_i16,
|
|
|
+ Uint32(sfb).ir_value(loc=loc, ip=ip),
|
|
|
+ bid_b_i16,
|
|
|
+ tid_b_i16,
|
|
|
+ Float32(d0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ mma.sync.aligned.kind::mxf8f6f4.block_scale.scale_vec::1X.m16n8k32.row.col.f32.e4m3.e4m3.f32.ue8m0
|
|
|
+ {$0, $1, $2, $3},
|
|
|
+ {$4, $5, $6, $7},
|
|
|
+ {$8, $9},
|
|
|
+ {$0, $1, $2, $3},
|
|
|
+ {$10},
|
|
|
+ {$11, $12},
|
|
|
+ {$13},
|
|
|
+ {$14, $15};
|
|
|
+ """,
|
|
|
+ "=f,=f,=f,=f,r,r,r,r,r,r,r,h,h,r,h,h,0,1,2,3",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Float32(r0), Float32(r1), Float32(r2), Float32(r3)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def mxfp8_mma_m16n8k32_f32_e2m1(
|
|
|
+ d0: Float32,
|
|
|
+ d1: Float32,
|
|
|
+ d2: Float32,
|
|
|
+ d3: Float32,
|
|
|
+ a0: Uint32,
|
|
|
+ a1: Uint32,
|
|
|
+ a2: Uint32,
|
|
|
+ a3: Uint32,
|
|
|
+ b0: Uint32,
|
|
|
+ b1: Uint32,
|
|
|
+ sfa: Uint32,
|
|
|
+ sfb: Uint32,
|
|
|
+ bid_a: int = 0,
|
|
|
+ tid_a: int = 0,
|
|
|
+ bid_b: int = 0,
|
|
|
+ tid_b: int = 0,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """SM120 block-scaled QMMA with E4M3 A and E2M1 B containers."""
|
|
|
+ asm = f"""
|
|
|
+ mma.sync.aligned.kind::mxf8f6f4.block_scale.scale_vec::1X.m16n8k32.row.col.f32.e4m3.e2m1.f32.ue8m0
|
|
|
+ {{$0, $1, $2, $3}},
|
|
|
+ {{$4, $5, $6, $7}},
|
|
|
+ {{$8, $9}},
|
|
|
+ {{$0, $1, $2, $3}},
|
|
|
+ {{$10}},
|
|
|
+ {{{int(bid_a)}, {int(tid_a)}}},
|
|
|
+ {{$11}},
|
|
|
+ {{{int(bid_b)}, {int(tid_b)}}};
|
|
|
+ """
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(a0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(sfa).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(sfb).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ asm,
|
|
|
+ "=f,=f,=f,=f,r,r,r,r,r,r,r,r,0,1,2,3",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ return tuple(
|
|
|
+ Float32(llvm.extractvalue(T.f32(), result, [i], loc=loc, ip=ip))
|
|
|
+ for i in range(4)
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def mma_m16n8k32_f32_e4m3(
|
|
|
+ d0: Float32,
|
|
|
+ d1: Float32,
|
|
|
+ d2: Float32,
|
|
|
+ d3: Float32,
|
|
|
+ a0: Uint32,
|
|
|
+ a1: Uint32,
|
|
|
+ a2: Uint32,
|
|
|
+ a3: Uint32,
|
|
|
+ b0: Uint32,
|
|
|
+ b1: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """Plain (non-block-scaled) SM120 FP8 E4M3 warp MMA `m16n8k32`.
|
|
|
+
|
|
|
+ d{0..3} are the accumulator IN-OUT (C on input, D on output)."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(a0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(a3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(d3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32
|
|
|
+ {$0, $1, $2, $3},
|
|
|
+ {$4, $5, $6, $7},
|
|
|
+ {$8, $9},
|
|
|
+ {$0, $1, $2, $3};
|
|
|
+ """,
|
|
|
+ "=f,=f,=f,=f,r,r,r,r,r,r,0,1,2,3",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ r0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ r1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+ r2 = llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)
|
|
|
+ r3 = llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)
|
|
|
+ return Float32(r0), Float32(r1), Float32(r2), Float32(r3)
|
|
|
+
|
|
|
+
|
|
|
+# ``mma_m16n8k16_f32_bf16`` was proven during the SM120 sparse-MLA port to be
|
|
|
+# bit-identical (same signature, same emitted
|
|
|
+# ``mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32``, same constraint
|
|
|
+# string and operand grouping) to the pre-existing ``bf16_mma_m16n8k16_f32``
|
|
|
+# above. Expose it under the port's name as a thin alias so the sparse-MLA
|
|
|
+# kernels can import it without duplicating the inline asm.
|
|
|
+mma_m16n8k16_f32_bf16 = bf16_mma_m16n8k16_f32
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def byte_perm(a: Uint32, b: Uint32, selector: Int32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """PTX byte permutation helper."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [
|
|
|
+ Uint32(a).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(b).ir_value(loc=loc, ip=ip),
|
|
|
+ Int32(selector).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ "prmt.b32 $0, $1, $2, $3;",
|
|
|
+ "=r,r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def frag_layout_swizzle_16b_to_8b(x: Uint32) -> Uint32:
|
|
|
+ tmp = cute.arch.shuffle_sync_bfly(x, offset=1)
|
|
|
+ x = byte_perm(x, tmp, Int32(0x5410 if cute.arch.lane_idx() & 0x1 == 0 else 0x3276))
|
|
|
+ tmp = cute.arch.shuffle_sync_bfly(x, offset=2)
|
|
|
+ x = byte_perm(x, tmp, Int32(0x5410 if cute.arch.lane_idx() & 0x2 == 0 else 0x3276))
|
|
|
+ return x
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def frag_layout_swizzle_16b_to_8b_trans(x: Uint32) -> Uint32:
|
|
|
+ tmp = cute.arch.shuffle_sync_bfly(x, offset=4)
|
|
|
+ x = byte_perm(x, tmp, Int32(0x6420 if cute.arch.lane_idx() & 0x4 == 0 else 0x3175))
|
|
|
+ tmp = cute.arch.shuffle_sync_bfly(x, offset=8)
|
|
|
+ x = byte_perm(x, tmp, Int32(0x5410 if cute.arch.lane_idx() & 0x8 == 0 else 0x3276))
|
|
|
+ tmp = cute.arch.shuffle_sync_bfly(x, offset=16)
|
|
|
+ x = byte_perm(x, tmp, Int32(0x5410 if cute.arch.lane_idx() & 0x10 == 0 else 0x3276))
|
|
|
+ return x
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def bfloat2_habs2(x: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """BFloat16x2 absolute value - clears sign bits of both bf16 values."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(x).ir_value(loc=loc, ip=ip)],
|
|
|
+ "and.b32 $0, $1, 0x7FFF7FFF;",
|
|
|
+ "=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def bfloat2_hmax2(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """BFloat16x2 max - element-wise max of 2 bf16 pairs."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(a).ir_value(loc=loc, ip=ip), Uint32(b).ir_value(loc=loc, ip=ip)],
|
|
|
+ "max.bf16x2 $0, $1, $2;",
|
|
|
+ "=r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def bfloat2_hmax_to_f32(x: Uint32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Extract max of 2 bf16 values in bfloat2 as float32."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Uint32(x).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b32 lo, hi;
|
|
|
+ .reg .f32 f0, f1;
|
|
|
+ and.b32 lo, $1, 0xFFFF;
|
|
|
+ shr.b32 hi, $1, 16;
|
|
|
+ shl.b32 lo, lo, 16;
|
|
|
+ shl.b32 hi, hi, 16;
|
|
|
+ mov.b32 f0, lo;
|
|
|
+ mov.b32 f1, hi;
|
|
|
+ max.f32 $0, f0, f1;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def bfloat2_to_float2_scaled(
|
|
|
+ bf2: Uint32, scale: Float32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Float32, Float32]:
|
|
|
+ """Convert bfloat16x2 to float2 AND multiply by scale."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32()]),
|
|
|
+ [Uint32(bf2).ir_value(loc=loc, ip=ip), Float32(scale).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b32 lo, hi;
|
|
|
+ .reg .f32 f0, f1;
|
|
|
+ and.b32 lo, $2, 0xFFFF;
|
|
|
+ shr.b32 hi, $2, 16;
|
|
|
+ shl.b32 lo, lo, 16;
|
|
|
+ shl.b32 hi, hi, 16;
|
|
|
+ mov.b32 f0, lo;
|
|
|
+ mov.b32 f1, hi;
|
|
|
+ mul.f32 $0, f0, $3;
|
|
|
+ mul.f32 $1, f1, $3;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,=f,r,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+
|
|
|
+ f0 = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ f1 = llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)
|
|
|
+
|
|
|
+ return Float32(f0), Float32(f1)
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# FP8 E4M3 Intrinsics
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_f32_to_e4m3(a: Float32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Convert float32 to E4M3 using native cvt.rn.satfinite.e4m3x2.f32."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Float32(a).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 fp8_pair;
|
|
|
+ .reg .f32 zero;
|
|
|
+ mov.f32 zero, 0f00000000;
|
|
|
+ cvt.rn.satfinite.e4m3x2.f32 fp8_pair, zero, $1;
|
|
|
+ cvt.u32.u16 $0, fp8_pair;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_f32x4_to_e4m3x4(
|
|
|
+ v0: Float32,
|
|
|
+ v1: Float32,
|
|
|
+ v2: Float32,
|
|
|
+ v3: Float32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Uint32:
|
|
|
+ """Convert 4 float32 values to 4 packed E4M3 bytes (v0 in the low byte)."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [
|
|
|
+ Float32(v0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 lo, hi;
|
|
|
+ cvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;
|
|
|
+ cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;
|
|
|
+ mov.b32 $0, {lo, hi};
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,f,f,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def nvfp4_scale_from_amax(block_amax: Float32, global_scale: Float32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Compute block_amax * reciprocal_global_scale / 6 with CUDA tensor-scalar semantics."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [
|
|
|
+ Float32(block_amax).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(global_scale).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .f64 amax_d, gs_d, q_d, six_d;
|
|
|
+ cvt.f64.f32 amax_d, $1;
|
|
|
+ cvt.f64.f32 gs_d, $2;
|
|
|
+ mov.f64 six_d, 0d4018000000000000;
|
|
|
+ mul.rn.f64 q_d, amax_d, gs_d;
|
|
|
+ div.rn.f64 q_d, q_d, six_d;
|
|
|
+ cvt.rn.f32.f64 $0, q_d;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def fp8_e4m3_to_f32(fp8_val: Uint32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Convert FP8 E4M3 to float32."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Uint32(fp8_val).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .pred p_zero, p_neg;
|
|
|
+ .reg .u32 sign_u, exp_u, mant_u;
|
|
|
+ .reg .s32 exp_s;
|
|
|
+ .reg .f32 exp_f, mant_f, fp8_float, fp8_neg;
|
|
|
+
|
|
|
+ setp.eq.u32 p_zero, $1, 0;
|
|
|
+ and.b32 sign_u, $1, 0x80;
|
|
|
+ and.b32 mant_u, $1, 7;
|
|
|
+ shr.b32 exp_u, $1, 3;
|
|
|
+ and.b32 exp_u, exp_u, 15;
|
|
|
+ sub.s32 exp_s, exp_u, 7;
|
|
|
+ cvt.rn.f32.s32 exp_f, exp_s;
|
|
|
+ ex2.approx.f32 exp_f, exp_f;
|
|
|
+ cvt.rn.f32.u32 mant_f, mant_u;
|
|
|
+ fma.rn.f32 mant_f, mant_f, 0f3E000000, 0f3F800000;
|
|
|
+ mul.f32 fp8_float, exp_f, mant_f;
|
|
|
+ neg.f32 fp8_neg, fp8_float;
|
|
|
+ setp.ne.u32 p_neg, sign_u, 0;
|
|
|
+ selp.f32 fp8_float, fp8_neg, fp8_float, p_neg;
|
|
|
+ selp.f32 $0, 0f00000000, fp8_float, p_zero;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def fp8_e4m3_to_f32_and_rcp(fp8_val: Uint32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Convert FP8 E4M3 to float32 AND compute reciprocal."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Uint32(fp8_val).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .pred p_zero;
|
|
|
+ .reg .u32 exp_u, mant_u;
|
|
|
+ .reg .s32 exp_s;
|
|
|
+ .reg .f32 exp_f, mant_f, fp8_float, result;
|
|
|
+
|
|
|
+ setp.eq.u32 p_zero, $1, 0;
|
|
|
+ and.b32 mant_u, $1, 7;
|
|
|
+ shr.b32 exp_u, $1, 3;
|
|
|
+ and.b32 exp_u, exp_u, 15;
|
|
|
+ sub.s32 exp_s, exp_u, 7;
|
|
|
+ cvt.rn.f32.s32 exp_f, exp_s;
|
|
|
+ ex2.approx.f32 exp_f, exp_f;
|
|
|
+ cvt.rn.f32.u32 mant_f, mant_u;
|
|
|
+ fma.rn.f32 mant_f, mant_f, 0f3E000000, 0f3F800000;
|
|
|
+ mul.f32 fp8_float, exp_f, mant_f;
|
|
|
+ rcp.approx.ftz.f32 result, fp8_float;
|
|
|
+ selp.f32 $0, 0f00000000, result, p_zero;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# E4M3 -> F16/F32 Native Conversion Intrinsics
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_e4m3x2_to_f16x2_pair(
|
|
|
+ packed_u32: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Decode 4 e4m3 bytes (packed u32) into 2 x f16x2 pairs."""
|
|
|
+ res = llvm.inline_asm(
|
|
|
+ ir.Type.parse("!llvm.struct<(i32, i32)>"),
|
|
|
+ [Uint32(packed_u32).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 b16_01, b16_23;
|
|
|
+ .reg .b32 hi32;
|
|
|
+ cvt.u16.u32 b16_01, $2;
|
|
|
+ shr.b32 hi32, $2, 16;
|
|
|
+ cvt.u16.u32 b16_23, hi32;
|
|
|
+ cvt.rn.f16x2.e4m3x2 $0, b16_01;
|
|
|
+ cvt.rn.f16x2.e4m3x2 $1, b16_23;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ return (
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip)),
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), res, [1], loc=loc, ip=ip)),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def f16x2_to_f32x2(
|
|
|
+ packed_h2: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Float32, Float32]:
|
|
|
+ """Unpack f16x2 (u32) to two f32 values."""
|
|
|
+ res = llvm.inline_asm(
|
|
|
+ ir.Type.parse("!llvm.struct<(f32, f32)>"),
|
|
|
+ [Uint32(packed_h2).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 lo, hi;
|
|
|
+ mov.b32 {lo, hi}, $2;
|
|
|
+ cvt.f32.f16 $0, lo;
|
|
|
+ cvt.f32.f16 $1, hi;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,=f,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ return (
|
|
|
+ Float32(llvm.extractvalue(T.f32(), res, [0], loc=loc, ip=ip)),
|
|
|
+ Float32(llvm.extractvalue(T.f32(), res, [1], loc=loc, ip=ip)),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def cvt_e4m3x4_to_f32x4(
|
|
|
+ packed_u32: Uint32,
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """Decode 4 e4m3 bytes (packed u32) to 4 x f32 via hw-native e4m3->f16->f32."""
|
|
|
+ h01, h23 = cvt_e4m3x2_to_f16x2_pair(packed_u32)
|
|
|
+ f0, f1 = f16x2_to_f32x2(h01)
|
|
|
+ f2, f3 = f16x2_to_f32x2(h23)
|
|
|
+ return f0, f1, f2, f3
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_e4m3_to_f32_via_f16(
|
|
|
+ fp8_val: Uint32, *, loc=None, ip=None
|
|
|
+) -> Float32:
|
|
|
+ """Convert single E4M3 byte to f32 via hw-native cvt.rn.f16x2.e4m3x2."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Uint32(fp8_val).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 fp8_pair;
|
|
|
+ .reg .b32 h2;
|
|
|
+ .reg .b16 lo, hi;
|
|
|
+ cvt.u16.u32 fp8_pair, $1;
|
|
|
+ cvt.rn.f16x2.e4m3x2 h2, fp8_pair;
|
|
|
+ mov.b32 {lo, hi}, h2;
|
|
|
+ cvt.f32.f16 $0, lo;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_w4a16_packed_e4m3_scale_to_f32(
|
|
|
+ packed_byte: Uint32, *, loc=None, ip=None
|
|
|
+) -> Float32:
|
|
|
+ """Scalar mirror of packed_dequant_e4m3x4_to_bfloat2x2 for W4A16 scales."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Uint32(packed_byte).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b32 bits, tmp;
|
|
|
+ .reg .b16 bf;
|
|
|
+ and.b32 bits, $1, 0x80;
|
|
|
+ shl.b32 bits, bits, 7;
|
|
|
+ and.b32 tmp, $1, 0x7f;
|
|
|
+ shl.b32 tmp, tmp, 4;
|
|
|
+ or.b32 bits, bits, tmp;
|
|
|
+ cvt.u16.u32 bf, bits;
|
|
|
+ cvt.f32.bf16 $0, bf;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def dequant_kv_e4m3_pair_to_bf16x2(
|
|
|
+ p0: Uint16, p1: Uint16, scale_f: Float32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """FP8 -> BF16 K dequant for the BF16-QK m16n8k16 B operand.
|
|
|
+
|
|
|
+ Byte-for-byte mirror of FlashInfer's sparse_mla prefill BF16-QK inline
|
|
|
+ sequence (prefill_kernel.cuh:240-251): two ``cvt.rn.f16x2.e4m3x2`` decode
|
|
|
+ the two e4m3 byte-pairs ``p0``/``p1`` (each a u16 holding two consecutive
|
|
|
+ e4m3 K-dims), multiply the four resulting f16 lanes by the per-(token,blk)
|
|
|
+ UE8M0 ``scale_f`` in f32, then two ``cvt.rn.bf16x2.f32`` pack them back to
|
|
|
+ the bf16x2 ``b0``/``b1`` MMA operands. NON-saturating cvt (matches the
|
|
|
+ reference; the K magnitudes are bounded by the e4m3 range so satfinite is a
|
|
|
+ no-op, and the ref emits plain ``cvt.rn.bf16x2.f32``)."""
|
|
|
+ res = llvm.inline_asm(
|
|
|
+ ir.Type.parse("!llvm.struct<(i32, i32)>"),
|
|
|
+ [
|
|
|
+ Uint16(p0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint16(p1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(scale_f).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b32 h2_0, h2_1;
|
|
|
+ .reg .b16 l0, h0, l1, h1;
|
|
|
+ .reg .f32 fk0, fk1, fk2, fk3;
|
|
|
+ cvt.rn.f16x2.e4m3x2 h2_0, $2;
|
|
|
+ cvt.rn.f16x2.e4m3x2 h2_1, $3;
|
|
|
+ mov.b32 {l0, h0}, h2_0;
|
|
|
+ mov.b32 {l1, h1}, h2_1;
|
|
|
+ cvt.f32.f16 fk0, l0;
|
|
|
+ cvt.f32.f16 fk1, h0;
|
|
|
+ cvt.f32.f16 fk2, l1;
|
|
|
+ cvt.f32.f16 fk3, h1;
|
|
|
+ mul.f32 fk0, fk0, $4;
|
|
|
+ mul.f32 fk1, fk1, $4;
|
|
|
+ mul.f32 fk2, fk2, $4;
|
|
|
+ mul.f32 fk3, fk3, $4;
|
|
|
+ cvt.rn.bf16x2.f32 $0, fk1, fk0;
|
|
|
+ cvt.rn.bf16x2.f32 $1, fk3, fk2;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,h,h,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ return (
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip)),
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), res, [1], loc=loc, ip=ip)),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_e8m0_to_f32(e8m0_val: Uint32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """Convert a single E8M0 scale byte to its true f32 value 2**(byte-127).
|
|
|
+
|
|
|
+ Byte 0 maps to 0.0. Unlike packed_dequant_e8m0x4_* (which fold a 2**7 bias
|
|
|
+ into the MMA input), this returns the unbiased scale so callers can apply it
|
|
|
+ directly in an f32 accumulator.
|
|
|
+ """
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Uint32(e8m0_val).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .pred p0;
|
|
|
+ .reg .u32 b0;
|
|
|
+ .reg .s32 e0;
|
|
|
+ .reg .f32 ef0;
|
|
|
+ and.b32 b0, $1, 0x000000ff;
|
|
|
+ setp.eq.u32 p0, b0, 0;
|
|
|
+ cvt.s32.u32 e0, b0;
|
|
|
+ sub.s32 e0, e0, 127;
|
|
|
+ cvt.rn.f32.s32 ef0, e0;
|
|
|
+ ex2.approx.f32 $0, ef0;
|
|
|
+ selp.f32 $0, 0f00000000, $0, p0;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_e8m0x4_to_f32x4(
|
|
|
+ packed: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Float32, Float32, Float32, Float32]:
|
|
|
+ """Decode 4 E8M0 scale bytes (packed u32) to 4 x true f32 (2**(byte-127))."""
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]),
|
|
|
+ [Uint32(packed).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .pred p0, p1, p2, p3;
|
|
|
+ .reg .u32 b0, b1, b2, b3;
|
|
|
+ .reg .s32 e0, e1, e2, e3;
|
|
|
+ .reg .f32 ef0, ef1, ef2, ef3;
|
|
|
+ and.b32 b0, $4, 0x000000ff;
|
|
|
+ shr.u32 b1, $4, 8;
|
|
|
+ and.b32 b1, b1, 0x000000ff;
|
|
|
+ shr.u32 b2, $4, 16;
|
|
|
+ and.b32 b2, b2, 0x000000ff;
|
|
|
+ shr.u32 b3, $4, 24;
|
|
|
+ setp.eq.u32 p0, b0, 0;
|
|
|
+ setp.eq.u32 p1, b1, 0;
|
|
|
+ setp.eq.u32 p2, b2, 0;
|
|
|
+ setp.eq.u32 p3, b3, 0;
|
|
|
+ cvt.s32.u32 e0, b0;
|
|
|
+ cvt.s32.u32 e1, b1;
|
|
|
+ cvt.s32.u32 e2, b2;
|
|
|
+ cvt.s32.u32 e3, b3;
|
|
|
+ sub.s32 e0, e0, 127;
|
|
|
+ sub.s32 e1, e1, 127;
|
|
|
+ sub.s32 e2, e2, 127;
|
|
|
+ sub.s32 e3, e3, 127;
|
|
|
+ cvt.rn.f32.s32 ef0, e0;
|
|
|
+ cvt.rn.f32.s32 ef1, e1;
|
|
|
+ cvt.rn.f32.s32 ef2, e2;
|
|
|
+ cvt.rn.f32.s32 ef3, e3;
|
|
|
+ ex2.approx.f32 $0, ef0;
|
|
|
+ ex2.approx.f32 $1, ef1;
|
|
|
+ ex2.approx.f32 $2, ef2;
|
|
|
+ ex2.approx.f32 $3, ef3;
|
|
|
+ selp.f32 $0, 0f00000000, $0, p0;
|
|
|
+ selp.f32 $1, 0f00000000, $1, p1;
|
|
|
+ selp.f32 $2, 0f00000000, $2, p2;
|
|
|
+ selp.f32 $3, 0f00000000, $3, p3;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,=f,=f,=f,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ return (
|
|
|
+ Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)),
|
|
|
+ Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)),
|
|
|
+ Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)),
|
|
|
+ Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# FP4 (E2M1) Decode Intrinsics
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def fp4_decode_4bytes(
|
|
|
+ packed_u32: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32, Uint32, Uint32]:
|
|
|
+ """Decode 4 FP4 bytes (packed u32) into 4 x f16x2 pairs."""
|
|
|
+ res = llvm.inline_asm(
|
|
|
+ ir.Type.parse("!llvm.struct<(i32, i32, i32, i32)>"),
|
|
|
+ [Uint32(packed_u32).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b8 byte0, byte1, byte2, byte3;
|
|
|
+ mov.b32 {byte0, byte1, byte2, byte3}, $4;
|
|
|
+ cvt.rn.f16x2.e2m1x2 $0, byte0;
|
|
|
+ cvt.rn.f16x2.e2m1x2 $1, byte1;
|
|
|
+ cvt.rn.f16x2.e2m1x2 $2, byte2;
|
|
|
+ cvt.rn.f16x2.e2m1x2 $3, byte3;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,=r,=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ return (
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip)),
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), res, [1], loc=loc, ip=ip)),
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), res, [2], loc=loc, ip=ip)),
|
|
|
+ Uint32(llvm.extractvalue(T.i32(), res, [3], loc=loc, ip=ip)),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def fp4_decode_2(byte_val: Uint32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Decode 1 FP4 byte into 1 f16x2 pair."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Uint32(byte_val).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b8 b0;
|
|
|
+ cvt.u8.u32 b0, $1;
|
|
|
+ cvt.rn.f16x2.e2m1x2 $0, b0;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def e2m1x8_to_qmma_e2m1x8(
|
|
|
+ packed_u32: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Spread eight packed FP4 nibbles into QMMA E2M1 byte containers.
|
|
|
+
|
|
|
+ The SM120 f8f6f4 instruction expects one E2M1 value per byte with the
|
|
|
+ original nibble in bits 5:2. This six-instruction spread is cheaper than
|
|
|
+ losslessly relabeling every value as E4M3 when no residual multiplier is
|
|
|
+ required.
|
|
|
+ """
|
|
|
+ res = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Uint32(packed_u32).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b32 s, t;
|
|
|
+ shl.b32 s, $2, 2;
|
|
|
+ shr.u32 t, $2, 2;
|
|
|
+ prmt.b32 $0, s, t, 0x5140;
|
|
|
+ prmt.b32 $1, s, t, 0x7362;
|
|
|
+ and.b32 $0, $0, 0x3C3C3C3C;
|
|
|
+ and.b32 $1, $1, 0x3C3C3C3C;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ lo = llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip)
|
|
|
+ hi = llvm.extractvalue(T.i32(), res, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(lo), Uint32(hi)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def e2m1x8_to_e4m3x8(
|
|
|
+ packed_u32: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Expand 8 packed E2M1 nibbles into 8 E4M3 bytes, losslessly.
|
|
|
+
|
|
|
+ Every E2M1 value is exactly representable in E4M3, so this is a pure
|
|
|
+ relabeling done with two prmt magnitude-LUT lookups plus a sign pass.
|
|
|
+ Nibble i of the input becomes byte i of the (lo, hi) output pair.
|
|
|
+ """
|
|
|
+ res = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [Uint32(packed_u32).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b32 clo, chi, cslo, cshi, mlo, mhi, slo, shi;
|
|
|
+ .reg .b32 tbl_a, tbl_b, tbl_s, zero;
|
|
|
+ // E4M3 encodings of {0, .5, 1, 1.5, 2, 3, 4, 6}.
|
|
|
+ mov.b32 tbl_a, 0x3C383000;
|
|
|
+ mov.b32 tbl_b, 0x4C484440;
|
|
|
+ mov.b32 tbl_s, 0x00008000;
|
|
|
+ mov.b32 zero, 0;
|
|
|
+ // Magnitude codes (low 3 bits of each nibble) as prmt selectors.
|
|
|
+ and.b32 clo, $2, 0x00007777;
|
|
|
+ shr.u32 chi, $2, 16;
|
|
|
+ and.b32 chi, chi, 0x00007777;
|
|
|
+ prmt.b32 mlo, tbl_a, tbl_b, clo;
|
|
|
+ prmt.b32 mhi, tbl_a, tbl_b, chi;
|
|
|
+ // Sign bit (nibble bit 3) -> byte bit 7 via a 2-entry LUT.
|
|
|
+ shr.u32 cslo, $2, 3;
|
|
|
+ and.b32 cslo, cslo, 0x00001111;
|
|
|
+ shr.u32 cshi, $2, 19;
|
|
|
+ and.b32 cshi, cshi, 0x00001111;
|
|
|
+ prmt.b32 slo, tbl_s, zero, cslo;
|
|
|
+ prmt.b32 shi, tbl_s, zero, cshi;
|
|
|
+ or.b32 $0, mlo, slo;
|
|
|
+ or.b32 $1, mhi, shi;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ lo = llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip)
|
|
|
+ hi = llvm.extractvalue(T.i32(), res, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(lo), Uint32(hi)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def e2m1x8_mul_residual_to_e4m3x8(
|
|
|
+ packed_u32: Uint32, residual_h2: Uint32, *, loc=None, ip=None
|
|
|
+) -> Tuple[Uint32, Uint32]:
|
|
|
+ """Expand 8 E2M1 nibbles to E4M3 bytes with a residual multiplier.
|
|
|
+
|
|
|
+ Decodes to f16 pairs, multiplies by ``residual_h2`` (the same residual
|
|
|
+ broadcast to both f16 lanes), and rounds to E4M3. Used by the w4a8
|
|
|
+ NVFP4 path where the per-K/16 e4m3 scale is decomposed into a shared
|
|
|
+ UE8M0 hardware exponent and this in-register residual.
|
|
|
+ """
|
|
|
+ res = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.i32(), T.i32()]),
|
|
|
+ [
|
|
|
+ Uint32(packed_u32).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(residual_h2).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b8 b0, b1, b2, b3;
|
|
|
+ .reg .b32 h0, h1, h2, h3;
|
|
|
+ .reg .b16 e0, e1, e2, e3;
|
|
|
+ mov.b32 {b0, b1, b2, b3}, $2;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h0, b0;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h1, b1;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h2, b2;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h3, b3;
|
|
|
+ mul.f16x2 h0, h0, $3;
|
|
|
+ mul.f16x2 h1, h1, $3;
|
|
|
+ mul.f16x2 h2, h2, $3;
|
|
|
+ mul.f16x2 h3, h3, $3;
|
|
|
+ cvt.rn.satfinite.e4m3x2.f16x2 e0, h0;
|
|
|
+ cvt.rn.satfinite.e4m3x2.f16x2 e1, h1;
|
|
|
+ cvt.rn.satfinite.e4m3x2.f16x2 e2, h2;
|
|
|
+ cvt.rn.satfinite.e4m3x2.f16x2 e3, h3;
|
|
|
+ mov.b32 $0, {e0, e1};
|
|
|
+ mov.b32 $1, {e2, e3};
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,=r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ lo = llvm.extractvalue(T.i32(), res, [0], loc=loc, ip=ip)
|
|
|
+ hi = llvm.extractvalue(T.i32(), res, [1], loc=loc, ip=ip)
|
|
|
+ return Uint32(lo), Uint32(hi)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_fp32x2_to_e2m1x2(v0: Float32, v1: Float32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """Convert 2 x f32 to 1 x FP4 (E2M1) byte."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [
|
|
|
+ Float32(v0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v1).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b8 b;
|
|
|
+ cvt.rn.satfinite.e2m1x2.f32 b, $2, $1;
|
|
|
+ cvt.u32.u8 $0, b;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def quant_dequant_2(
|
|
|
+ v0: Float32, v1: Float32, sf_f32: Float32, eff_scale: Float32
|
|
|
+) -> Tuple[Float32, Float32]:
|
|
|
+ """Quantize-dequantize pair roundtrip through FP4."""
|
|
|
+ inv_scale = Float32(1.0) / eff_scale
|
|
|
+ fp4_byte = cvt_fp32x2_to_e2m1x2(v0 * inv_scale, v1 * inv_scale)
|
|
|
+ h2 = fp4_decode_2(fp4_byte)
|
|
|
+ f0, f1 = f16x2_to_f32x2(h2)
|
|
|
+ return f0 * sf_f32, f1 * sf_f32
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# Half2 FMA Dot Product Intrinsics
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def hfma2_4(
|
|
|
+ w0: Uint32, x0: Uint32,
|
|
|
+ w1: Uint32, x1: Uint32,
|
|
|
+ w2: Uint32, x2: Uint32,
|
|
|
+ w3: Uint32, x3: Uint32,
|
|
|
+ *, loc=None, ip=None
|
|
|
+) -> Tuple[Float32, Float32]:
|
|
|
+ """4-element half2 FMA dot product returning (lo, hi) f32."""
|
|
|
+ res = llvm.inline_asm(
|
|
|
+ ir.Type.parse("!llvm.struct<(f32, f32)>"),
|
|
|
+ [
|
|
|
+ Uint32(w0).ir_value(loc=loc, ip=ip), Uint32(x0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w1).ir_value(loc=loc, ip=ip), Uint32(x1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w2).ir_value(loc=loc, ip=ip), Uint32(x2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w3).ir_value(loc=loc, ip=ip), Uint32(x3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .f16x2 acc;
|
|
|
+ .reg .b16 lo, hi;
|
|
|
+ mov.b32 acc, 0;
|
|
|
+ fma.rn.f16x2 acc, $2, $3, acc;
|
|
|
+ fma.rn.f16x2 acc, $4, $5, acc;
|
|
|
+ fma.rn.f16x2 acc, $6, $7, acc;
|
|
|
+ fma.rn.f16x2 acc, $8, $9, acc;
|
|
|
+ mov.b32 {lo, hi}, acc;
|
|
|
+ cvt.f32.f16 $0, lo;
|
|
|
+ cvt.f32.f16 $1, hi;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,=f,r,r,r,r,r,r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ return (
|
|
|
+ Float32(llvm.extractvalue(T.f32(), res, [0], loc=loc, ip=ip)),
|
|
|
+ Float32(llvm.extractvalue(T.f32(), res, [1], loc=loc, ip=ip)),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def hfma2_4_sum(
|
|
|
+ w0: Uint32, x0: Uint32,
|
|
|
+ w1: Uint32, x1: Uint32,
|
|
|
+ w2: Uint32, x2: Uint32,
|
|
|
+ w3: Uint32, x3: Uint32,
|
|
|
+ *, loc=None, ip=None
|
|
|
+) -> Float32:
|
|
|
+ """4-element half2 FMA dot product returning lane sum."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [
|
|
|
+ Uint32(w0).ir_value(loc=loc, ip=ip), Uint32(x0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w1).ir_value(loc=loc, ip=ip), Uint32(x1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w2).ir_value(loc=loc, ip=ip), Uint32(x2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w3).ir_value(loc=loc, ip=ip), Uint32(x3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .f16x2 acc;
|
|
|
+ .reg .b16 lo, hi;
|
|
|
+ .reg .f32 flo, fhi;
|
|
|
+ mov.b32 acc, 0;
|
|
|
+ fma.rn.f16x2 acc, $1, $2, acc;
|
|
|
+ fma.rn.f16x2 acc, $3, $4, acc;
|
|
|
+ fma.rn.f16x2 acc, $5, $6, acc;
|
|
|
+ fma.rn.f16x2 acc, $7, $8, acc;
|
|
|
+ mov.b32 {lo, hi}, acc;
|
|
|
+ cvt.f32.f16 flo, lo;
|
|
|
+ cvt.f32.f16 fhi, hi;
|
|
|
+ add.f32 $0, flo, fhi;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r,r,r,r,r,r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def hfma2_8(
|
|
|
+ w0: Uint32, x0: Uint32,
|
|
|
+ w1: Uint32, x1: Uint32,
|
|
|
+ w2: Uint32, x2: Uint32,
|
|
|
+ w3: Uint32, x3: Uint32,
|
|
|
+ w4: Uint32, x4: Uint32,
|
|
|
+ w5: Uint32, x5: Uint32,
|
|
|
+ w6: Uint32, x6: Uint32,
|
|
|
+ w7: Uint32, x7: Uint32,
|
|
|
+ *, loc=None, ip=None
|
|
|
+) -> Tuple[Float32, Float32]:
|
|
|
+ """8-element half2 FMA dot product returning (lo, hi) f32."""
|
|
|
+ res = llvm.inline_asm(
|
|
|
+ ir.Type.parse("!llvm.struct<(f32, f32)>"),
|
|
|
+ [
|
|
|
+ Uint32(w0).ir_value(loc=loc, ip=ip), Uint32(x0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w1).ir_value(loc=loc, ip=ip), Uint32(x1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w2).ir_value(loc=loc, ip=ip), Uint32(x2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w3).ir_value(loc=loc, ip=ip), Uint32(x3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w4).ir_value(loc=loc, ip=ip), Uint32(x4).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w5).ir_value(loc=loc, ip=ip), Uint32(x5).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w6).ir_value(loc=loc, ip=ip), Uint32(x6).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w7).ir_value(loc=loc, ip=ip), Uint32(x7).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .f16x2 acc;
|
|
|
+ .reg .b16 lo, hi;
|
|
|
+ mov.b32 acc, 0;
|
|
|
+ fma.rn.f16x2 acc, $2, $3, acc;
|
|
|
+ fma.rn.f16x2 acc, $4, $5, acc;
|
|
|
+ fma.rn.f16x2 acc, $6, $7, acc;
|
|
|
+ fma.rn.f16x2 acc, $8, $9, acc;
|
|
|
+ fma.rn.f16x2 acc, $10, $11, acc;
|
|
|
+ fma.rn.f16x2 acc, $12, $13, acc;
|
|
|
+ fma.rn.f16x2 acc, $14, $15, acc;
|
|
|
+ fma.rn.f16x2 acc, $16, $17, acc;
|
|
|
+ mov.b32 {lo, hi}, acc;
|
|
|
+ cvt.f32.f16 $0, lo;
|
|
|
+ cvt.f32.f16 $1, hi;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,=f,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ return (
|
|
|
+ Float32(llvm.extractvalue(T.f32(), res, [0], loc=loc, ip=ip)),
|
|
|
+ Float32(llvm.extractvalue(T.f32(), res, [1], loc=loc, ip=ip)),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def hfma2_8_sum(
|
|
|
+ w0: Uint32, x0: Uint32,
|
|
|
+ w1: Uint32, x1: Uint32,
|
|
|
+ w2: Uint32, x2: Uint32,
|
|
|
+ w3: Uint32, x3: Uint32,
|
|
|
+ w4: Uint32, x4: Uint32,
|
|
|
+ w5: Uint32, x5: Uint32,
|
|
|
+ w6: Uint32, x6: Uint32,
|
|
|
+ w7: Uint32, x7: Uint32,
|
|
|
+ *, loc=None, ip=None
|
|
|
+) -> Float32:
|
|
|
+ """8-element half2 FMA dot product returning lane sum."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [
|
|
|
+ Uint32(w0).ir_value(loc=loc, ip=ip), Uint32(x0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w1).ir_value(loc=loc, ip=ip), Uint32(x1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w2).ir_value(loc=loc, ip=ip), Uint32(x2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w3).ir_value(loc=loc, ip=ip), Uint32(x3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w4).ir_value(loc=loc, ip=ip), Uint32(x4).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w5).ir_value(loc=loc, ip=ip), Uint32(x5).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w6).ir_value(loc=loc, ip=ip), Uint32(x6).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(w7).ir_value(loc=loc, ip=ip), Uint32(x7).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .f16x2 acc;
|
|
|
+ .reg .b16 lo, hi;
|
|
|
+ .reg .f32 flo, fhi;
|
|
|
+ mov.b32 acc, 0;
|
|
|
+ fma.rn.f16x2 acc, $1, $2, acc;
|
|
|
+ fma.rn.f16x2 acc, $3, $4, acc;
|
|
|
+ fma.rn.f16x2 acc, $5, $6, acc;
|
|
|
+ fma.rn.f16x2 acc, $7, $8, acc;
|
|
|
+ fma.rn.f16x2 acc, $9, $10, acc;
|
|
|
+ fma.rn.f16x2 acc, $11, $12, acc;
|
|
|
+ fma.rn.f16x2 acc, $13, $14, acc;
|
|
|
+ fma.rn.f16x2 acc, $15, $16, acc;
|
|
|
+ mov.b32 {lo, hi}, acc;
|
|
|
+ cvt.f32.f16 flo, lo;
|
|
|
+ cvt.f32.f16 fhi, hi;
|
|
|
+ add.f32 $0, flo, fhi;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# FP4 Dot Product Intrinsics
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def fp4_dot4_sum(
|
|
|
+ u_packed: Uint32,
|
|
|
+ x0: Uint32,
|
|
|
+ x1: Uint32,
|
|
|
+ x2: Uint32,
|
|
|
+ x3: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Float32:
|
|
|
+ """FP4 dot product: decode 4 FP4 bytes and dot with 4 x f16x2 inputs."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [
|
|
|
+ Uint32(u_packed).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x3).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b8 b0, b1, b2, b3;
|
|
|
+ .reg .b32 h0, h1, h2, h3;
|
|
|
+ .reg .f16x2 acc;
|
|
|
+ .reg .b16 lo, hi;
|
|
|
+ .reg .f32 flo, fhi;
|
|
|
+ mov.b32 {b0, b1, b2, b3}, $1;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h0, b0;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h1, b1;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h2, b2;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h3, b3;
|
|
|
+ mov.b32 acc, 0;
|
|
|
+ fma.rn.f16x2 acc, h0, $2, acc;
|
|
|
+ fma.rn.f16x2 acc, h1, $3, acc;
|
|
|
+ fma.rn.f16x2 acc, h2, $4, acc;
|
|
|
+ fma.rn.f16x2 acc, h3, $5, acc;
|
|
|
+ mov.b32 {lo, hi}, acc;
|
|
|
+ cvt.f32.f16 flo, lo;
|
|
|
+ cvt.f32.f16 fhi, hi;
|
|
|
+ add.f32 $0, flo, fhi;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r,r,r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def fp4_dot8_dual_sum(
|
|
|
+ up_a: Uint32,
|
|
|
+ up_b: Uint32,
|
|
|
+ gate_a: Uint32,
|
|
|
+ gate_b: Uint32,
|
|
|
+ x0: Uint32,
|
|
|
+ x1: Uint32,
|
|
|
+ x2: Uint32,
|
|
|
+ x3: Uint32,
|
|
|
+ x4: Uint32,
|
|
|
+ x5: Uint32,
|
|
|
+ x6: Uint32,
|
|
|
+ x7: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Float32]:
|
|
|
+ """Fused dual FP4 dot product (up + gate) over the SAME 8 shared f16x2
|
|
|
+ activation words. Decodes both weight rows and runs two INDEPENDENT
|
|
|
+ f16x2 accumulator chains in a single inline-asm block, then reduces each
|
|
|
+ to f32. Bit-identical to running fp4_dot8_sum twice (same per-chain fma
|
|
|
+ order, same lo+hi f32 reduction) but emits one asm block so ptxas can
|
|
|
+ interleave the two independent fma chains for 2-way ILP and halves the
|
|
|
+ per-pair epilogue (mov/cvt/add) op count on the issue-bound decode loop."""
|
|
|
+ res = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.f32()]),
|
|
|
+ [
|
|
|
+ Uint32(up_a).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(up_b).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(gate_a).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(gate_b).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x4).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x5).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x6).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x7).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b8 ua0, ua1, ua2, ua3, ub0, ub1, ub2, ub3;
|
|
|
+ .reg .b8 ga0, ga1, ga2, ga3, gb0, gb1, gb2, gb3;
|
|
|
+ .reg .b32 uh0, uh1, uh2, uh3, uh4, uh5, uh6, uh7;
|
|
|
+ .reg .b32 gh0, gh1, gh2, gh3, gh4, gh5, gh6, gh7;
|
|
|
+ .reg .f16x2 uacc, gacc;
|
|
|
+ .reg .b16 ulo, uhi, glo, ghi;
|
|
|
+ .reg .f32 uflo, ufhi, gflo, gfhi;
|
|
|
+ mov.b32 {ua0, ua1, ua2, ua3}, $2;
|
|
|
+ mov.b32 {ub0, ub1, ub2, ub3}, $3;
|
|
|
+ mov.b32 {ga0, ga1, ga2, ga3}, $4;
|
|
|
+ mov.b32 {gb0, gb1, gb2, gb3}, $5;
|
|
|
+ cvt.rn.f16x2.e2m1x2 uh0, ua0;
|
|
|
+ cvt.rn.f16x2.e2m1x2 uh1, ua1;
|
|
|
+ cvt.rn.f16x2.e2m1x2 uh2, ua2;
|
|
|
+ cvt.rn.f16x2.e2m1x2 uh3, ua3;
|
|
|
+ cvt.rn.f16x2.e2m1x2 uh4, ub0;
|
|
|
+ cvt.rn.f16x2.e2m1x2 uh5, ub1;
|
|
|
+ cvt.rn.f16x2.e2m1x2 uh6, ub2;
|
|
|
+ cvt.rn.f16x2.e2m1x2 uh7, ub3;
|
|
|
+ cvt.rn.f16x2.e2m1x2 gh0, ga0;
|
|
|
+ cvt.rn.f16x2.e2m1x2 gh1, ga1;
|
|
|
+ cvt.rn.f16x2.e2m1x2 gh2, ga2;
|
|
|
+ cvt.rn.f16x2.e2m1x2 gh3, ga3;
|
|
|
+ cvt.rn.f16x2.e2m1x2 gh4, gb0;
|
|
|
+ cvt.rn.f16x2.e2m1x2 gh5, gb1;
|
|
|
+ cvt.rn.f16x2.e2m1x2 gh6, gb2;
|
|
|
+ cvt.rn.f16x2.e2m1x2 gh7, gb3;
|
|
|
+ mov.b32 uacc, 0;
|
|
|
+ mov.b32 gacc, 0;
|
|
|
+ fma.rn.f16x2 uacc, uh0, $6, uacc;
|
|
|
+ fma.rn.f16x2 gacc, gh0, $6, gacc;
|
|
|
+ fma.rn.f16x2 uacc, uh1, $7, uacc;
|
|
|
+ fma.rn.f16x2 gacc, gh1, $7, gacc;
|
|
|
+ fma.rn.f16x2 uacc, uh2, $8, uacc;
|
|
|
+ fma.rn.f16x2 gacc, gh2, $8, gacc;
|
|
|
+ fma.rn.f16x2 uacc, uh3, $9, uacc;
|
|
|
+ fma.rn.f16x2 gacc, gh3, $9, gacc;
|
|
|
+ fma.rn.f16x2 uacc, uh4, $10, uacc;
|
|
|
+ fma.rn.f16x2 gacc, gh4, $10, gacc;
|
|
|
+ fma.rn.f16x2 uacc, uh5, $11, uacc;
|
|
|
+ fma.rn.f16x2 gacc, gh5, $11, gacc;
|
|
|
+ fma.rn.f16x2 uacc, uh6, $12, uacc;
|
|
|
+ fma.rn.f16x2 gacc, gh6, $12, gacc;
|
|
|
+ fma.rn.f16x2 uacc, uh7, $13, uacc;
|
|
|
+ fma.rn.f16x2 gacc, gh7, $13, gacc;
|
|
|
+ mov.b32 {ulo, uhi}, uacc;
|
|
|
+ mov.b32 {glo, ghi}, gacc;
|
|
|
+ cvt.f32.f16 uflo, ulo;
|
|
|
+ cvt.f32.f16 ufhi, uhi;
|
|
|
+ cvt.f32.f16 gflo, glo;
|
|
|
+ cvt.f32.f16 gfhi, ghi;
|
|
|
+ add.f32 $0, uflo, ufhi;
|
|
|
+ add.f32 $1, gflo, gfhi;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,=f,r,r,r,r,r,r,r,r,r,r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ up = Float32(llvm.extractvalue(T.f32(), res, [0], loc=loc, ip=ip))
|
|
|
+ gate = Float32(llvm.extractvalue(T.f32(), res, [1], loc=loc, ip=ip))
|
|
|
+ return up, gate
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def fp4_dot8_sum(
|
|
|
+ u_a: Uint32,
|
|
|
+ u_b: Uint32,
|
|
|
+ x0: Uint32,
|
|
|
+ x1: Uint32,
|
|
|
+ x2: Uint32,
|
|
|
+ x3: Uint32,
|
|
|
+ x4: Uint32,
|
|
|
+ x5: Uint32,
|
|
|
+ x6: Uint32,
|
|
|
+ x7: Uint32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Float32:
|
|
|
+ """FP4 dot product: decode 8 FP4 bytes (2 x u32) and dot with 8 x f16x2 inputs."""
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [
|
|
|
+ Uint32(u_a).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(u_b).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x0).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x1).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x2).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x3).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x4).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x5).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x6).ir_value(loc=loc, ip=ip),
|
|
|
+ Uint32(x7).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b8 a0, a1, a2, a3;
|
|
|
+ .reg .b8 b0, b1, b2, b3;
|
|
|
+ .reg .b32 h0, h1, h2, h3, h4, h5, h6, h7;
|
|
|
+ .reg .f16x2 acc;
|
|
|
+ .reg .b16 lo, hi;
|
|
|
+ .reg .f32 flo, fhi;
|
|
|
+ mov.b32 {a0, a1, a2, a3}, $1;
|
|
|
+ mov.b32 {b0, b1, b2, b3}, $2;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h0, a0;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h1, a1;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h2, a2;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h3, a3;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h4, b0;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h5, b1;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h6, b2;
|
|
|
+ cvt.rn.f16x2.e2m1x2 h7, b3;
|
|
|
+ mov.b32 acc, 0;
|
|
|
+ fma.rn.f16x2 acc, h0, $3, acc;
|
|
|
+ fma.rn.f16x2 acc, h1, $4, acc;
|
|
|
+ fma.rn.f16x2 acc, h2, $5, acc;
|
|
|
+ fma.rn.f16x2 acc, h3, $6, acc;
|
|
|
+ fma.rn.f16x2 acc, h4, $7, acc;
|
|
|
+ fma.rn.f16x2 acc, h5, $8, acc;
|
|
|
+ fma.rn.f16x2 acc, h6, $9, acc;
|
|
|
+ fma.rn.f16x2 acc, h7, $10, acc;
|
|
|
+ mov.b32 {lo, hi}, acc;
|
|
|
+ cvt.f32.f16 flo, lo;
|
|
|
+ cvt.f32.f16 fhi, hi;
|
|
|
+ add.f32 $0, flo, fhi;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r,r,r,r,r,r,r,r,r,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def _f16x2_dot_sum_f32acc(a_h2: Uint32, b_h2: Uint32) -> Float32:
|
|
|
+ a0, a1 = f16x2_to_f32x2(a_h2)
|
|
|
+ b0, b1 = f16x2_to_f32x2(b_h2)
|
|
|
+ return a0 * b0 + a1 * b1
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def fp4_dot4_sum_f32acc(
|
|
|
+ u_packed: Uint32,
|
|
|
+ x0: Uint32,
|
|
|
+ x1: Uint32,
|
|
|
+ x2: Uint32,
|
|
|
+ x3: Uint32,
|
|
|
+) -> Float32:
|
|
|
+ """FP4 dot product using f32 accumulation after FP4 decode."""
|
|
|
+ h0, h1, h2, h3 = fp4_decode_4bytes(u_packed)
|
|
|
+ return (
|
|
|
+ _f16x2_dot_sum_f32acc(h0, x0)
|
|
|
+ + _f16x2_dot_sum_f32acc(h1, x1)
|
|
|
+ + _f16x2_dot_sum_f32acc(h2, x2)
|
|
|
+ + _f16x2_dot_sum_f32acc(h3, x3)
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def fp4_dot8_sum_f32acc(
|
|
|
+ u_a: Uint32,
|
|
|
+ u_b: Uint32,
|
|
|
+ x0: Uint32,
|
|
|
+ x1: Uint32,
|
|
|
+ x2: Uint32,
|
|
|
+ x3: Uint32,
|
|
|
+ x4: Uint32,
|
|
|
+ x5: Uint32,
|
|
|
+ x6: Uint32,
|
|
|
+ x7: Uint32,
|
|
|
+) -> Float32:
|
|
|
+ """FP4 dot product using f32 accumulation after FP4 decode."""
|
|
|
+ h0, h1, h2, h3 = fp4_decode_4bytes(u_a)
|
|
|
+ h4, h5, h6, h7 = fp4_decode_4bytes(u_b)
|
|
|
+ return (
|
|
|
+ _f16x2_dot_sum_f32acc(h0, x0)
|
|
|
+ + _f16x2_dot_sum_f32acc(h1, x1)
|
|
|
+ + _f16x2_dot_sum_f32acc(h2, x2)
|
|
|
+ + _f16x2_dot_sum_f32acc(h3, x3)
|
|
|
+ + _f16x2_dot_sum_f32acc(h4, x4)
|
|
|
+ + _f16x2_dot_sum_f32acc(h5, x5)
|
|
|
+ + _f16x2_dot_sum_f32acc(h6, x6)
|
|
|
+ + _f16x2_dot_sum_f32acc(h7, x7)
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# Pack Intrinsics
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def pack_f32x2_to_f16x2(
|
|
|
+ a: Float32, b: Float32, *, loc=None, ip=None
|
|
|
+) -> Uint32:
|
|
|
+ """Pack two f32 values into one f16x2 u32."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [
|
|
|
+ Float32(a).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(b).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b16 lo, hi;
|
|
|
+ cvt.rn.f16.f32 lo, $1;
|
|
|
+ cvt.rn.f16.f32 hi, $2;
|
|
|
+ mov.b32 $0, {lo, hi};
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# UE8M0 Intrinsics (for MXFP4)
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def pow2_ceil_ue8m0(
|
|
|
+ scale: Float32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Tuple[Float32, Uint32]:
|
|
|
+ """Round a positive FP32 ``scale`` UP to a power of two, bit-exactly.
|
|
|
+
|
|
|
+ Returns ``(rounded_fp32, ue8m0_byte)`` where:
|
|
|
+ * ``rounded_fp32`` == ``__uint_as_float(power_of_2_round(__float_as_uint(scale)))``
|
|
|
+ * ``ue8m0_byte`` == ``(__float_as_uint(rounded_fp32) >> 23) & 0xFF``
|
|
|
+
|
|
|
+ Bit-exact replica of FlashInfer's fp8_quant.cuh rounding +
|
|
|
+ scale_convert.cuh fp32_to_ue8m0. Integer ops only (no lg2.approx), so unlike
|
|
|
+ ``cvt_f32_to_ue8m0`` it matches the FlashInfer reference bit-for-bit.
|
|
|
+ """
|
|
|
+ result = llvm.inline_asm(
|
|
|
+ llvm.StructType.get_literal([T.f32(), T.i32()]),
|
|
|
+ [Float32(scale).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .pred p_mant;
|
|
|
+ .reg .b32 bits, mant;
|
|
|
+
|
|
|
+ // bits = __float_as_uint(scale)
|
|
|
+ mov.b32 bits, $2;
|
|
|
+ // mant = bits & 0x007FFFFF (mantissa field)
|
|
|
+ and.b32 mant, bits, 8388607;
|
|
|
+ setp.ne.u32 p_mant, mant, 0;
|
|
|
+ // if (mant) bits = (bits + 0x00800000) & 0x7F800000
|
|
|
+ @p_mant add.u32 bits, bits, 8388608;
|
|
|
+ @p_mant and.b32 bits, bits, 2139095040;
|
|
|
+ // rounded fp32 scale = __uint_as_float(bits)
|
|
|
+ mov.b32 $0, bits;
|
|
|
+ // ue8m0 = (bits >> 23) & 0xFF
|
|
|
+ bfe.u32 $1, bits, 23, 8;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,=r,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ loc=loc,
|
|
|
+ ip=ip,
|
|
|
+ )
|
|
|
+ rounded = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)
|
|
|
+ ue8m0 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
|
|
|
+ return Float32(rounded), Uint32(ue8m0)
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_f32_to_ue8m0(max_val: Float32, *, loc=None, ip=None) -> Uint32:
|
|
|
+ """
|
|
|
+ Convert float32 max value to UE8M0 scale factor.
|
|
|
+
|
|
|
+ UE8M0 is unsigned 8-bit exponent-only format:
|
|
|
+ - value = 2^(ue8m0 - 127)
|
|
|
+ - ue8m0 = ceil(log2(max_val)) + 127
|
|
|
+
|
|
|
+ Uses lg2.approx.f32 for fast log2 approximation.
|
|
|
+ Uses cvt.rpi (round towards positive infinity, i.e., ceiling).
|
|
|
+ Returns value clamped to [0, 255].
|
|
|
+ """
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [Float32(max_val).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .pred p_zero, p_neg, p_ovf;
|
|
|
+ .reg .f32 log2_val;
|
|
|
+ .reg .s32 exp_int, result;
|
|
|
+
|
|
|
+ // Check for zero/negative
|
|
|
+ setp.le.f32 p_zero, $1, 0f00000000;
|
|
|
+
|
|
|
+ // Compute ceil(log2(max_val)) using cvt.rpi (round towards +inf)
|
|
|
+ lg2.approx.f32 log2_val, $1;
|
|
|
+ cvt.rpi.s32.f32 exp_int, log2_val;
|
|
|
+
|
|
|
+ // Add bias and clamp to [0, 255]
|
|
|
+ add.s32 result, exp_int, 127;
|
|
|
+ setp.lt.s32 p_neg, result, 0;
|
|
|
+ setp.gt.s32 p_ovf, result, 255;
|
|
|
+ selp.s32 result, 0, result, p_neg;
|
|
|
+ selp.s32 result, 255, result, p_ovf;
|
|
|
+ selp.s32 $0, 0, result, p_zero;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def ue8m0_to_output_scale(ue8m0_val: Uint32, *, loc=None, ip=None) -> Float32:
|
|
|
+ """
|
|
|
+ Convert UE8M0 to output_scale for MXFP4 quantization.
|
|
|
+
|
|
|
+ UE8M0 value = 2^(ue8m0 - 127)
|
|
|
+ Returns 1 / 2^(ue8m0 - 127) = 2^(127 - ue8m0)
|
|
|
+ """
|
|
|
+ return Float32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.f32(),
|
|
|
+ [Uint32(ue8m0_val).ir_value(loc=loc, ip=ip)],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .pred p_zero;
|
|
|
+ .reg .s32 neg_exp;
|
|
|
+ .reg .f32 neg_exp_f, result;
|
|
|
+
|
|
|
+ // Check for zero
|
|
|
+ setp.eq.u32 p_zero, $1, 0;
|
|
|
+
|
|
|
+ // Compute 2^(127 - ue8m0) = 1 / 2^(ue8m0 - 127)
|
|
|
+ sub.s32 neg_exp, 127, $1;
|
|
|
+ cvt.rn.f32.s32 neg_exp_f, neg_exp;
|
|
|
+ ex2.approx.f32 result, neg_exp_f;
|
|
|
+ selp.f32 $0, 0f00000000, result, p_zero;
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=f,r",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# E2M1 Conversion
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@dsl_user_op
|
|
|
+def cvt_e2m1x8_f32(
|
|
|
+ v0: Float32,
|
|
|
+ v1: Float32,
|
|
|
+ v2: Float32,
|
|
|
+ v3: Float32,
|
|
|
+ v4: Float32,
|
|
|
+ v5: Float32,
|
|
|
+ v6: Float32,
|
|
|
+ v7: Float32,
|
|
|
+ *,
|
|
|
+ loc=None,
|
|
|
+ ip=None,
|
|
|
+) -> Uint32:
|
|
|
+ """Convert eight float32 values to eight E2M1 (4-bit) values packed into uint32."""
|
|
|
+ return Uint32(
|
|
|
+ llvm.inline_asm(
|
|
|
+ T.i32(),
|
|
|
+ [
|
|
|
+ Float32(v0).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v1).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v2).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v3).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v4).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v5).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v6).ir_value(loc=loc, ip=ip),
|
|
|
+ Float32(v7).ir_value(loc=loc, ip=ip),
|
|
|
+ ],
|
|
|
+ """
|
|
|
+ {
|
|
|
+ .reg .b8 byte0, byte1, byte2, byte3;
|
|
|
+ cvt.rn.satfinite.e2m1x2.f32 byte0, $2, $1;
|
|
|
+ cvt.rn.satfinite.e2m1x2.f32 byte1, $4, $3;
|
|
|
+ cvt.rn.satfinite.e2m1x2.f32 byte2, $6, $5;
|
|
|
+ cvt.rn.satfinite.e2m1x2.f32 byte3, $8, $7;
|
|
|
+ mov.b32 $0, {byte0, byte1, byte2, byte3};
|
|
|
+ }
|
|
|
+ """,
|
|
|
+ "=r,f,f,f,f,f,f,f,f",
|
|
|
+ has_side_effects=False,
|
|
|
+ is_align_stack=False,
|
|
|
+ asm_dialect=llvm.AsmDialect.AD_ATT,
|
|
|
+ )
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# Warp, Block, and Cluster Reduction Utilities
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def warp_reduce(val, op, width: cutlass.Constexpr[int] = 32):
|
|
|
+ """Reduce across threads in a warp using butterfly shuffle."""
|
|
|
+ if cutlass.const_expr(isinstance(val, cute.TensorSSA)):
|
|
|
+ res = cute.make_rmem_tensor(val.shape, val.dtype)
|
|
|
+ res.store(val)
|
|
|
+ for i in cutlass.range_constexpr(cute.size(val.shape)):
|
|
|
+ res[i] = warp_reduce(res[i], op, width)
|
|
|
+ return res.load()
|
|
|
+ else:
|
|
|
+ for i in cutlass.range_constexpr(int(math.log2(width))):
|
|
|
+ val = op(val, cute.arch.shuffle_sync_bfly(val, offset=1 << i))
|
|
|
+ return val
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def block_reduce(
|
|
|
+ val: Float32,
|
|
|
+ op: Callable,
|
|
|
+ reduction_buffer: cute.Tensor,
|
|
|
+ init_val: Float32,
|
|
|
+) -> Float32:
|
|
|
+ """Block reduction across multiple warps using shared memory."""
|
|
|
+ lane_idx = cute.arch.lane_idx()
|
|
|
+ warp_idx = cute.arch.warp_idx()
|
|
|
+ warps_per_row = cute.size(reduction_buffer.shape[1])
|
|
|
+ row_idx = warp_idx // warps_per_row
|
|
|
+ col_idx = warp_idx % warps_per_row
|
|
|
+
|
|
|
+ if lane_idx == 0:
|
|
|
+ reduction_buffer[row_idx, col_idx] = val
|
|
|
+ cute.arch.barrier()
|
|
|
+
|
|
|
+ block_reduce_val = init_val
|
|
|
+ if lane_idx < warps_per_row:
|
|
|
+ block_reduce_val = reduction_buffer[row_idx, lane_idx]
|
|
|
+ return warp_reduce(block_reduce_val, op)
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def cluster_reduce(
|
|
|
+ val: Float32,
|
|
|
+ op: Callable,
|
|
|
+ reduction_buffer: cute.Tensor,
|
|
|
+ mbar_ptr: cute.Pointer,
|
|
|
+ cluster_n: cutlass.Constexpr[int],
|
|
|
+ init_val: Float32,
|
|
|
+) -> Float32:
|
|
|
+ """Cluster reduction across multiple CTAs using mbarrier."""
|
|
|
+ cta_rank_in_cluster = cute.arch.block_idx_in_cluster()
|
|
|
+ lane_idx = cute.arch.lane_idx()
|
|
|
+ warp_idx = cute.arch.warp_idx()
|
|
|
+
|
|
|
+ rows_per_block = reduction_buffer.shape[0]
|
|
|
+ warps_per_row = reduction_buffer.shape[1][0]
|
|
|
+
|
|
|
+ row_idx = warp_idx // warps_per_row
|
|
|
+ col_idx = warp_idx % warps_per_row
|
|
|
+
|
|
|
+ # Warp 0 sets up mbarrier transaction count
|
|
|
+ if warp_idx == 0:
|
|
|
+ with cute.arch.elect_one():
|
|
|
+ num_warps = rows_per_block * warps_per_row
|
|
|
+ expected_bytes = num_warps * cluster_n * 4
|
|
|
+ cute.arch.mbarrier_arrive_and_expect_tx(mbar_ptr, expected_bytes)
|
|
|
+
|
|
|
+ # Each lane < cluster_n writes to a different CTA's shared memory
|
|
|
+ if lane_idx < cluster_n:
|
|
|
+ store_shared_remote(
|
|
|
+ val,
|
|
|
+ elem_pointer(reduction_buffer, (row_idx, (col_idx, cta_rank_in_cluster))),
|
|
|
+ mbar_ptr,
|
|
|
+ peer_cta_rank_in_cluster=lane_idx,
|
|
|
+ )
|
|
|
+
|
|
|
+ # Wait for all cluster writes
|
|
|
+ cute.arch.mbarrier_wait(mbar_ptr, phase=0)
|
|
|
+
|
|
|
+ # Reduce across all values
|
|
|
+ num_total = warps_per_row * cluster_n
|
|
|
+ num_iter = cute.ceil_div(num_total, 32)
|
|
|
+
|
|
|
+ block_reduce_val = init_val
|
|
|
+ for i in cutlass.range_constexpr(num_iter):
|
|
|
+ idx = lane_idx + i * 32
|
|
|
+ if idx < num_total:
|
|
|
+ block_reduce_val = op(block_reduce_val, reduction_buffer[row_idx, idx])
|
|
|
+
|
|
|
+ return warp_reduce(block_reduce_val, op)
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def row_reduce(
|
|
|
+ x: cute.TensorSSA,
|
|
|
+ op: cute.ReductionOp,
|
|
|
+ threads_per_row: cutlass.Constexpr[int],
|
|
|
+ reduction_buffer: cute.Tensor,
|
|
|
+ mbar_ptr,
|
|
|
+ cluster_n: cutlass.Constexpr[int],
|
|
|
+ init_val: Float32,
|
|
|
+):
|
|
|
+ """Row reduction with optional cluster support."""
|
|
|
+ local_val = x.reduce(op, init_val=init_val, reduction_profile=0)
|
|
|
+
|
|
|
+ warp_op = {
|
|
|
+ cute.ReductionOp.ADD: operator.add,
|
|
|
+ cute.ReductionOp.MAX: cute.arch.fmax,
|
|
|
+ }[op]
|
|
|
+ warp_width = min(threads_per_row, 32)
|
|
|
+ warp_val = warp_reduce(local_val, warp_op, width=warp_width)
|
|
|
+
|
|
|
+ warps_per_row = max(threads_per_row // 32, 1)
|
|
|
+
|
|
|
+ if cutlass.const_expr(warps_per_row > 1 or cluster_n > 1):
|
|
|
+ if cutlass.const_expr(cluster_n == 1):
|
|
|
+ return block_reduce(warp_val, warp_op, reduction_buffer, init_val)
|
|
|
+ else:
|
|
|
+ return cluster_reduce(
|
|
|
+ warp_val, warp_op, reduction_buffer, mbar_ptr, cluster_n, init_val
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ return warp_val
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# Predicate Utility
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def predicate_k(tXcX: cute.Tensor, limit: int) -> cute.Tensor:
|
|
|
+ """Create predicate tensor for bounds checking."""
|
|
|
+ tXpX = cute.make_rmem_tensor(
|
|
|
+ cute.make_layout(
|
|
|
+ (
|
|
|
+ cute.size(tXcX, mode=[0, 1]),
|
|
|
+ cute.size(tXcX, mode=[1]),
|
|
|
+ cute.size(tXcX, mode=[2]),
|
|
|
+ ),
|
|
|
+ stride=(cute.size(tXcX, mode=[2]), 0, 1),
|
|
|
+ ),
|
|
|
+ cutlass.Boolean,
|
|
|
+ )
|
|
|
+ for rest_v in cutlass.range_constexpr(tXpX.shape[0]):
|
|
|
+ for rest_k in cutlass.range_constexpr(tXpX.shape[2]):
|
|
|
+ tXpX[rest_v, 0, rest_k] = cute.elem_less(
|
|
|
+ tXcX[(0, rest_v), 0, rest_k][1], limit
|
|
|
+ )
|
|
|
+ return tXpX
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# Helper Functions for SF Block Processing (block_size=16)
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def load_8_half2(
|
|
|
+ mX: cute.Tensor, mW: cute.Tensor, row_offset: Int32, col_offset: Int32, H: int
|
|
|
+):
|
|
|
+ """Load 16 elements (8 half2 pairs) of X and W from global memory.
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ x_h2: rmem_tensor of shape (8,) containing X as half2
|
|
|
+ w_h2: rmem_tensor of shape (8,) containing W as half2
|
|
|
+ """
|
|
|
+ x_h2 = cute.make_rmem_tensor((8,), Uint32)
|
|
|
+ w_h2 = cute.make_rmem_tensor((8,), Uint32)
|
|
|
+
|
|
|
+ # Load X (2 x 128-bit loads = 16 elements)
|
|
|
+ x_ptr0 = get_ptr_as_int64(mX, row_offset * H + col_offset)
|
|
|
+ x_ptr1 = get_ptr_as_int64(mX, row_offset * H + col_offset + Int32(8))
|
|
|
+ x_h2[0], x_h2[1], x_h2[2], x_h2[3] = ld_global_v4_u32(x_ptr0)
|
|
|
+ x_h2[4], x_h2[5], x_h2[6], x_h2[7] = ld_global_v4_u32(x_ptr1)
|
|
|
+
|
|
|
+ # Load W (2 x 128-bit loads = 16 elements)
|
|
|
+ w_ptr0 = get_ptr_as_int64(mW, col_offset)
|
|
|
+ w_ptr1 = get_ptr_as_int64(mW, col_offset + Int32(8))
|
|
|
+ w_h2[0], w_h2[1], w_h2[2], w_h2[3] = ld_global_v4_u32(w_ptr0)
|
|
|
+ w_h2[4], w_h2[5], w_h2[6], w_h2[7] = ld_global_v4_u32(w_ptr1)
|
|
|
+
|
|
|
+ return x_h2, w_h2
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def half2_mul_8(x_h2: cute.Tensor, w_h2: cute.Tensor) -> cute.Tensor:
|
|
|
+ """Multiply 8 half2 pairs element-wise."""
|
|
|
+ xw_h2 = cute.make_rmem_tensor((8,), Uint32)
|
|
|
+ for i in cutlass.range_constexpr(8):
|
|
|
+ xw_h2[i] = half2_mul(x_h2[i], w_h2[i])
|
|
|
+ return xw_h2
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def bfloat2_mul_8(x_h2: cute.Tensor, w_h2: cute.Tensor) -> cute.Tensor:
|
|
|
+ """Multiply 8 bfloat2 pairs element-wise."""
|
|
|
+ xw_h2 = cute.make_rmem_tensor((8,), Uint32)
|
|
|
+ for i in cutlass.range_constexpr(8):
|
|
|
+ xw_h2[i] = bfloat2_mul(x_h2[i], w_h2[i])
|
|
|
+ return xw_h2
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def half2_max_abs_8(xw_h2: cute.Tensor) -> Uint32:
|
|
|
+ """Compute max absolute value across 8 half2 values using tree reduction."""
|
|
|
+ # Compute abs for all 8 values
|
|
|
+ abs_h2 = cute.make_rmem_tensor((8,), Uint32)
|
|
|
+ for i in cutlass.range_constexpr(8):
|
|
|
+ abs_h2[i] = habs2(xw_h2[i])
|
|
|
+
|
|
|
+ # Tree reduction: 8 -> 4 -> 2 -> 1
|
|
|
+ max_01 = hmax2(abs_h2[0], abs_h2[1])
|
|
|
+ max_23 = hmax2(abs_h2[2], abs_h2[3])
|
|
|
+ max_45 = hmax2(abs_h2[4], abs_h2[5])
|
|
|
+ max_67 = hmax2(abs_h2[6], abs_h2[7])
|
|
|
+ max_0123 = hmax2(max_01, max_23)
|
|
|
+ max_4567 = hmax2(max_45, max_67)
|
|
|
+ return hmax2(max_0123, max_4567)
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def bfloat2_max_abs_8(xw_h2: cute.Tensor) -> Uint32:
|
|
|
+ """Compute max absolute value across 8 bfloat2 values using tree reduction."""
|
|
|
+ # Compute abs for all 8 values
|
|
|
+ abs_h2 = cute.make_rmem_tensor((8,), Uint32)
|
|
|
+ for i in cutlass.range_constexpr(8):
|
|
|
+ abs_h2[i] = bfloat2_habs2(xw_h2[i])
|
|
|
+
|
|
|
+ # Tree reduction: 8 -> 4 -> 2 -> 1
|
|
|
+ max_01 = bfloat2_hmax2(abs_h2[0], abs_h2[1])
|
|
|
+ max_23 = bfloat2_hmax2(abs_h2[2], abs_h2[3])
|
|
|
+ max_45 = bfloat2_hmax2(abs_h2[4], abs_h2[5])
|
|
|
+ max_67 = bfloat2_hmax2(abs_h2[6], abs_h2[7])
|
|
|
+ max_0123 = bfloat2_hmax2(max_01, max_23)
|
|
|
+ max_4567 = bfloat2_hmax2(max_45, max_67)
|
|
|
+ return bfloat2_hmax2(max_0123, max_4567)
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def half2_to_float16(xw_h2: cute.Tensor, scale: Float32) -> cute.Tensor:
|
|
|
+ """Convert 8 half2 to 16 float32 with scaling."""
|
|
|
+ y_f32 = cute.make_rmem_tensor((16,), Float32)
|
|
|
+ for i in cutlass.range_constexpr(8):
|
|
|
+ y_f32[i * 2], y_f32[i * 2 + 1] = half2_to_float2_scaled(xw_h2[i], scale)
|
|
|
+ return y_f32
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def bfloat2_to_float16(xw_h2: cute.Tensor, scale: Float32) -> cute.Tensor:
|
|
|
+ """Convert 8 bfloat2 to 16 float32 with scaling."""
|
|
|
+ y_f32 = cute.make_rmem_tensor((16,), Float32)
|
|
|
+ for i in cutlass.range_constexpr(8):
|
|
|
+ y_f32[i * 2], y_f32[i * 2 + 1] = bfloat2_to_float2_scaled(xw_h2[i], scale)
|
|
|
+ return y_f32
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def quantize_and_pack_16(y_f32: cute.Tensor, value_scale: Float32) -> Uint64:
|
|
|
+ """Quantize 16 float32 values to FP4 and pack into uint64."""
|
|
|
+ t025 = value_scale * Float32(0.25)
|
|
|
+ t075 = value_scale * Float32(0.75)
|
|
|
+ t125 = value_scale * Float32(1.25)
|
|
|
+ t175 = value_scale * Float32(1.75)
|
|
|
+ t250 = value_scale * Float32(2.5)
|
|
|
+ t350 = value_scale * Float32(3.5)
|
|
|
+ t500 = value_scale * Float32(5.0)
|
|
|
+ packed = Uint64(0)
|
|
|
+ for i in cutlass.range_constexpr(16):
|
|
|
+ q = y_f32[i]
|
|
|
+ mag = fabs_f32(q)
|
|
|
+ nibble = Uint8(0)
|
|
|
+ if mag > t025 and mag < t075:
|
|
|
+ nibble = Uint8(1)
|
|
|
+ elif mag >= t075 and mag <= t125:
|
|
|
+ nibble = Uint8(2)
|
|
|
+ elif mag > t125 and mag < t175:
|
|
|
+ nibble = Uint8(3)
|
|
|
+ elif mag >= t175 and mag <= t250:
|
|
|
+ nibble = Uint8(4)
|
|
|
+ elif mag > t250 and mag < t350:
|
|
|
+ nibble = Uint8(5)
|
|
|
+ elif mag >= t350 and mag <= t500:
|
|
|
+ nibble = Uint8(6)
|
|
|
+ elif mag > t500:
|
|
|
+ nibble = Uint8(7)
|
|
|
+ if nibble != Uint8(0) and q < Float32(0.0):
|
|
|
+ nibble = nibble | Uint8(0x8)
|
|
|
+ packed = packed | (Uint64(nibble) << Uint64(i * 4))
|
|
|
+ return packed
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def quantize_and_pack_16_fast(y_f32: cute.Tensor, inv_scale: Float32) -> Uint64:
|
|
|
+ """Fast approximate FP4 quantize/pack for 16 float32 values."""
|
|
|
+ q = cute.make_rmem_tensor((16,), Float32)
|
|
|
+ for i in cutlass.range_constexpr(16):
|
|
|
+ q[i] = y_f32[i] * inv_scale
|
|
|
+
|
|
|
+ packed_lo = cvt_e2m1x8_f32(q[0], q[1], q[2], q[3], q[4], q[5], q[6], q[7])
|
|
|
+ packed_hi = cvt_e2m1x8_f32(q[8], q[9], q[10], q[11], q[12], q[13], q[14], q[15])
|
|
|
+ return (Uint64(packed_hi) << Uint64(32)) | Uint64(packed_lo)
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def quantize_block_fp4(
|
|
|
+ values: cute.Tensor, max_abs: Float32, global_scale_val: Float32,
|
|
|
+) -> Tuple[Uint64, cutlass.Uint8]:
|
|
|
+ """Quantize 16 float32 values to packed FP4 + e4m3 scale byte.
|
|
|
+
|
|
|
+ Given 16 values and their pre-computed max_abs, derives the NVFP4 block
|
|
|
+ scale, quantizes to FP4, and packs into a uint64. Returns
|
|
|
+ (packed_fp4_u64, scale_byte). The caller handles storage writes.
|
|
|
+ """
|
|
|
+ scale_float = max_abs * global_scale_val / Float32(FLOAT4_E2M1_MAX)
|
|
|
+ scale_float = fmin_f32(scale_float, Float32(FLOAT8_E4M3_MAX))
|
|
|
+ scale_u32 = cvt_f32_to_e4m3(scale_float)
|
|
|
+ scale_byte = cutlass.Uint8(scale_u32 & Uint32(0xFF))
|
|
|
+ quantized_scale = fp8_e4m3_to_f32(scale_u32)
|
|
|
+ packed64 = Uint64(0)
|
|
|
+ if quantized_scale != Float32(0.0) and global_scale_val != Float32(0.0):
|
|
|
+ packed64 = quantize_and_pack_16(values, quantized_scale / global_scale_val)
|
|
|
+ return packed64, scale_byte
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def quantize_block_fp4_fast(
|
|
|
+ values: cute.Tensor, max_abs: Float32, global_scale_val: Float32,
|
|
|
+) -> Tuple[Uint64, cutlass.Uint8]:
|
|
|
+ """Fast approximate FP4 block quantization using reciprocal/vector path."""
|
|
|
+ scale_u32 = Uint32(0)
|
|
|
+ scale_byte = cutlass.Uint8(0)
|
|
|
+ packed64 = Uint64(0)
|
|
|
+ if global_scale_val != Float32(0.0):
|
|
|
+ fp4_max_rcp = rcp_approx_ftz(Float32(FLOAT4_E2M1_MAX))
|
|
|
+ scale_float = global_scale_val * (max_abs * fp4_max_rcp)
|
|
|
+ scale_float = fmin_f32(scale_float, Float32(FLOAT8_E4M3_MAX))
|
|
|
+ scale_u32 = cvt_f32_to_e4m3(scale_float)
|
|
|
+ scale_byte = cutlass.Uint8(scale_u32 & Uint32(0xFF))
|
|
|
+ inv_quantized_scale = fp8_e4m3_to_f32_and_rcp(scale_u32)
|
|
|
+ if inv_quantized_scale != Float32(0.0):
|
|
|
+ packed64 = quantize_and_pack_16_fast(values, inv_quantized_scale * global_scale_val)
|
|
|
+ return packed64, scale_byte
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def max_abs_16(values: cute.Tensor) -> Float32:
|
|
|
+ """Compute the maximum absolute value of 16 float32 values."""
|
|
|
+ result = fabs_f32(values[0])
|
|
|
+ for i in cutlass.range_constexpr(1, 16):
|
|
|
+ result = fmax_f32(result, fabs_f32(values[i]))
|
|
|
+ return result
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def silu_mul_16(
|
|
|
+ gate: cute.Tensor, up: cute.Tensor,
|
|
|
+) -> cute.Tensor:
|
|
|
+ """Fused SiLU(gate) * up for 16 float32 element pairs.
|
|
|
+
|
|
|
+ Used in the CompactMoEKernel epilogue to fuse the activation function
|
|
|
+ between GEMM1 and GEMM2, avoiding the gmem round-trip.
|
|
|
+ """
|
|
|
+ out = cute.make_rmem_tensor((16,), Float32)
|
|
|
+ for i in cutlass.range_constexpr(16):
|
|
|
+ g = gate[i]
|
|
|
+ sigmoid_g = cute.arch.rcp_approx(
|
|
|
+ Float32(1.0) + cute.math.exp(-g, fastmath=False)
|
|
|
+ )
|
|
|
+ out[i] = g * sigmoid_g * up[i]
|
|
|
+ return out
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def silu_mul_quantize_block_fp4(
|
|
|
+ gate: cute.Tensor, up: cute.Tensor, global_scale_val: Float32,
|
|
|
+) -> Tuple[Uint64, cutlass.Uint8]:
|
|
|
+ """Fused SiLU(gate)*up + FP4 quantize for 16 element pairs.
|
|
|
+
|
|
|
+ Combines silu_mul_16, max_abs_16, and quantize_block_fp4 into a single
|
|
|
+ helper for the CompactMoEKernel epilogue.
|
|
|
+ """
|
|
|
+ activated = silu_mul_16(gate, up)
|
|
|
+ block_max = max_abs_16(activated)
|
|
|
+ return quantize_block_fp4(activated, block_max, global_scale_val)
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def relu2_16(x: cute.Tensor) -> cute.Tensor:
|
|
|
+ """Compute ReLU^2 for 16 float32 values."""
|
|
|
+ out = cute.make_rmem_tensor((16,), Float32)
|
|
|
+ for i in cutlass.range_constexpr(16):
|
|
|
+ v = fmax_f32(x[i], Float32(0.0))
|
|
|
+ out[i] = v * v
|
|
|
+ return out
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def relu2_quantize_block_fp4(
|
|
|
+ x: cute.Tensor, global_scale_val: Float32,
|
|
|
+) -> Tuple[Uint64, cutlass.Uint8]:
|
|
|
+ """Fuse ReLU^2 and FP4 quantization for 16 float32 values."""
|
|
|
+ activated = relu2_16(x)
|
|
|
+ block_max = max_abs_16(activated)
|
|
|
+ return quantize_block_fp4(activated, block_max, global_scale_val)
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# MXFP8 (w4a8) Block Quantization: 32-wide UE8M0 + E4M3
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def max_abs_32(values: cute.Tensor) -> Float32:
|
|
|
+ """Compute the maximum absolute value of 32 float32 values."""
|
|
|
+ result = fabs_f32(values[0])
|
|
|
+ for i in cutlass.range_constexpr(1, 32):
|
|
|
+ result = fmax_f32(result, fabs_f32(values[i]))
|
|
|
+ return result
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def quantize_block_fp8_mx(
|
|
|
+ values: cute.Tensor, max_abs: Float32,
|
|
|
+) -> Tuple[cute.Tensor, Uint32]:
|
|
|
+ """Quantize 32 float32 values to E4M3 payload bytes + UE8M0 scale byte.
|
|
|
+
|
|
|
+ The scale is ``pow2_ceil(max_abs / 448)`` so the scaled payload never
|
|
|
+ exceeds the E4M3 range; the payload carries its own exponent, so the
|
|
|
+ power-of-two block scale handles range only. No global scale is
|
|
|
+ involved (the MX scheme is self-ranging). Returns
|
|
|
+ ``(payload_u32x8, ue8m0_byte)`` with value ``4*i + j`` in byte ``j`` of
|
|
|
+ ``payload[i]``. An all-zero block yields scale byte 0 (decoded as a
|
|
|
+ zero output scale) and a zero payload.
|
|
|
+ """
|
|
|
+ scale_f = max_abs * Float32(1.0 / FLOAT8_E4M3_MAX)
|
|
|
+ _, scale_byte = pow2_ceil_ue8m0(scale_f)
|
|
|
+ inv_scale = ue8m0_to_output_scale(scale_byte)
|
|
|
+ payload = cute.make_rmem_tensor((8,), Uint32)
|
|
|
+ for i in cutlass.range_constexpr(8):
|
|
|
+ payload[i] = cvt_f32x4_to_e4m3x4(
|
|
|
+ values[i * 4 + 0] * inv_scale,
|
|
|
+ values[i * 4 + 1] * inv_scale,
|
|
|
+ values[i * 4 + 2] * inv_scale,
|
|
|
+ values[i * 4 + 3] * inv_scale,
|
|
|
+ )
|
|
|
+ return payload, scale_byte
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def silu_mul_32(
|
|
|
+ gate: cute.Tensor, up: cute.Tensor,
|
|
|
+) -> cute.Tensor:
|
|
|
+ """Fused SiLU(gate) * up for 32 float32 element pairs."""
|
|
|
+ out = cute.make_rmem_tensor((32,), Float32)
|
|
|
+ for i in cutlass.range_constexpr(32):
|
|
|
+ g = gate[i]
|
|
|
+ sigmoid_g = cute.arch.rcp_approx(
|
|
|
+ Float32(1.0) + cute.math.exp(-g, fastmath=False)
|
|
|
+ )
|
|
|
+ out[i] = g * sigmoid_g * up[i]
|
|
|
+ return out
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def silu_mul_quantize_block_fp8_mx(
|
|
|
+ gate: cute.Tensor, up: cute.Tensor,
|
|
|
+) -> Tuple[cute.Tensor, Uint32]:
|
|
|
+ """Fused SiLU(gate)*up + MXFP8 quantize for 32 element pairs."""
|
|
|
+ activated = silu_mul_32(gate, up)
|
|
|
+ block_max = max_abs_32(activated)
|
|
|
+ return quantize_block_fp8_mx(activated, block_max)
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def relu2_32(x: cute.Tensor) -> cute.Tensor:
|
|
|
+ """Compute ReLU^2 for 32 float32 values."""
|
|
|
+ out = cute.make_rmem_tensor((32,), Float32)
|
|
|
+ for i in cutlass.range_constexpr(32):
|
|
|
+ v = fmax_f32(x[i], Float32(0.0))
|
|
|
+ out[i] = v * v
|
|
|
+ return out
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def relu2_quantize_block_fp8_mx(
|
|
|
+ x: cute.Tensor,
|
|
|
+) -> Tuple[cute.Tensor, Uint32]:
|
|
|
+ """Fuse ReLU^2 and MXFP8 quantization for 32 float32 values."""
|
|
|
+ activated = relu2_32(x)
|
|
|
+ block_max = max_abs_32(activated)
|
|
|
+ return quantize_block_fp8_mx(activated, block_max)
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def mx_scale_from_amax32(amax: Float32) -> Tuple[Float32, Float32]:
|
|
|
+ """Per-32 MXFP8 block scale from a precomputed amax.
|
|
|
+
|
|
|
+ Returns ``(scale, inv_scale)`` with ``scale = pow2_ceil(amax/448)`` —
|
|
|
+ identical numerics to ``quantize_block_fp8_mx`` (zero amax yields
|
|
|
+ (0, 0), silencing the block).
|
|
|
+ """
|
|
|
+ rounded, byte = pow2_ceil_ue8m0(amax * Float32(1.0 / FLOAT8_E4M3_MAX))
|
|
|
+ inv_scale = ue8m0_to_output_scale(byte)
|
|
|
+ return rounded, inv_scale
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def quant_dequant_e4m3_2(
|
|
|
+ v0: Float32, v1: Float32, inv_scale: Float32, scale: Float32,
|
|
|
+) -> Tuple[Float32, Float32]:
|
|
|
+ """Quantize-dequantize a pair through E4M3 with a power-of-two block scale.
|
|
|
+
|
|
|
+ Bit-matches the w4a8 prefill activation quantizer
|
|
|
+ (``quantize_block_fp8_mx``): RN-saturating E4M3 of ``v*inv_scale``,
|
|
|
+ decoded back and rescaled. Used by the micro decode kernel's a8_mx mode
|
|
|
+ so decode numerics track the prefill recipe.
|
|
|
+ """
|
|
|
+ q0 = fp8_e4m3_to_f32(cvt_f32_to_e4m3(v0 * inv_scale)) * scale
|
|
|
+ q1 = fp8_e4m3_to_f32(cvt_f32_to_e4m3(v1 * inv_scale)) * scale
|
|
|
+ return q0, q1
|
|
|
+
|
|
|
+
|
|
|
+# =============================================================================
|
|
|
+# Helper Functions for Float32 SF Block Processing
|
|
|
+# =============================================================================
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def load_f32_16_from_smem(
|
|
|
+ sH: cute.Tensor, row_idx: Int32, col_offset: Int32
|
|
|
+) -> cute.Tensor:
|
|
|
+ """Load 16 Float32 values from shared memory."""
|
|
|
+ h_f32 = cute.make_rmem_tensor((16,), Float32)
|
|
|
+ for i in cutlass.range_constexpr(16):
|
|
|
+ h_f32[i] = Float32(sH[row_idx, col_offset + i])
|
|
|
+ return h_f32
|
|
|
+
|
|
|
+
|
|
|
+@cute.jit
|
|
|
+def compute_y_and_max_abs_f32(
|
|
|
+ h_f32: cute.Tensor, w_f32: cute.Tensor, rstd: Float32
|
|
|
+) -> Tuple[cute.Tensor, Float32]:
|
|
|
+ """Compute y = h * rstd * w and max_abs for 16 Float32 values."""
|
|
|
+ y_f32 = cute.make_rmem_tensor((16,), Float32)
|
|
|
+
|
|
|
+ # Compute y and track max_abs
|
|
|
+ y_f32[0] = h_f32[0] * rstd * w_f32[0]
|
|
|
+ max_abs = fabs_f32(y_f32[0])
|
|
|
+
|
|
|
+ for i in cutlass.range_constexpr(1, 16):
|
|
|
+ y_f32[i] = h_f32[i] * rstd * w_f32[i]
|
|
|
+ max_abs = fmax_f32(max_abs, fabs_f32(y_f32[i]))
|
|
|
+
|
|
|
+ return y_f32, max_abs
|
