From 3e5c6594f41749aba8f40230dbceb9f549239f17 Mon Sep 17 00:00:00 2001 From: Theotime Combes Date: Mon, 3 Aug 2026 12:23:04 +0000 Subject: [PATCH] [Pallas] Add SparseCore BF16 lane packing stack-info: PR: https://github.com/pytorch/helion/pull/3262, branch: gh/thcmbs/14/head --- helion/_compiler/pallas/sc_base.py | 2 +- helion/_compiler/pallas/sparsecore_compute.py | 201 +++++++++++++++--- test/test_pallas_sparsecore.py | 22 +- test/test_pallas_sparsecore_compute.py | 10 +- test/test_pallas_sparsecore_plan.py | 4 +- 5 files changed, 208 insertions(+), 31 deletions(-) diff --git a/helion/_compiler/pallas/sc_base.py b/helion/_compiler/pallas/sc_base.py index f56e2865f..e3a03c8aa 100644 --- a/helion/_compiler/pallas/sc_base.py +++ b/helion/_compiler/pallas/sc_base.py @@ -96,5 +96,5 @@ def _reject( ) -_INDIRECT_DTYPES = (torch.float32, torch.int32) +_INDIRECT_DTYPES = (torch.float32, torch.int32, torch.bfloat16) _CAST_STORE_DTYPES = (torch.int8, torch.int32, torch.bool) diff --git a/helion/_compiler/pallas/sparsecore_compute.py b/helion/_compiler/pallas/sparsecore_compute.py index a07b2dbc2..f90ad2fc3 100644 --- a/helion/_compiler/pallas/sparsecore_compute.py +++ b/helion/_compiler/pallas/sparsecore_compute.py @@ -16,11 +16,14 @@ from ..helper_function import CodegenInterface from ..inductor_lowering import GraphInterpreter from ..inductor_lowering import ReductionLowering +from .memory_access import MemoryAccessKind from .sc_base import _CAST_STORE_DTYPES from .sc_base import SC_LANES from .sparsecore_plan import IndirectLoadPlan if TYPE_CHECKING: + from collections.abc import Iterable + from ..generate_ast import GenerateAST from .sparsecore_program import SparseCoreProgram @@ -59,22 +62,90 @@ def __getattr__(self, name: str) -> object: class LaneChunk: start: int size: int + unique_start: int -def chunk_schedule(value_size: int) -> list[LaneChunk]: +def chunk_schedule(value_size: int, dtype: torch.dtype) -> list[LaneChunk]: if value_size < 1: raise NotImplementedError("SparseCore value is empty") - if value_size % SC_LANES: + if dtype is not torch.bfloat16: + if value_size % SC_LANES: + raise NotImplementedError( + f"SparseCore value size {value_size} must be a multiple of {SC_LANES}" + ) + return [ + LaneChunk(start, SC_LANES, start) + for start in range(0, value_size, SC_LANES) + ] + if value_size < 2 * SC_LANES or value_size % SC_LANES: raise NotImplementedError( - f"SparseCore value size {value_size} must be a multiple of {SC_LANES}" + f"SparseCore bf16 value size {value_size} must be a multiple of " + f"{SC_LANES} and at least {2 * SC_LANES}" ) - return [LaneChunk(start, SC_LANES) for start in range(0, value_size, SC_LANES)] + starts = list(range(0, max(value_size - 31, 1), 32)) + if starts[-1] != value_size - 32: + starts.append(value_size - 32) + chunks: list[LaneChunk] = [] + covered = 0 + for start in starts: + chunks.append(LaneChunk(start, 32, max(covered, start))) + covered = max(covered, start + 32) + return chunks def _value_size(value: torch.Tensor) -> int: return math.prod(int(dim) for dim in value.shape[1:]) if value.ndim > 1 else 1 +def packed_bf16_segments( + program: SparseCoreProgram, + node: torch.fx.Node, + memo: dict[torch.fx.Node, tuple[int, ...]] | None = None, +) -> tuple[int, ...]: + """Find packed BF16 boundaries inherited by a value.""" + if memo is None: + memo = {} + if node in memo: + return memo[node] + plan = program.plan_by_node.get(node) + if ( + plan is not None + and plan.access.kind is MemoryAccessKind.LOAD + and plan.layout.logical_dtype is torch.bfloat16 + ): + result = (plan.layout.value_size,) + else: + value = node.meta.get("val") + if isinstance(value, torch.Tensor) and _value_size(value) == 1: + result = () + elif node.target is torch.ops.aten.stack.default: + values = node.args[0] + if not isinstance(values, (list, tuple)): + raise NotImplementedError("SparseCore stack has no value sequence") + result = tuple( + segment + for child in values + if isinstance(child, torch.fx.Node) + for segment in packed_bf16_segments(program, child, memo) + ) + else: + candidates = [ + segments + for parent in node.all_input_nodes + if (segments := packed_bf16_segments(program, parent, memo)) + ] + if not candidates: + result = () + elif all(candidate == candidates[0] for candidate in candidates[1:]): + result = candidates[0] + else: + raise NotImplementedError( + "SparseCore pointwise operands have incompatible packed BF16 layouts" + ) + memo[node] = result + return result + + @dataclass(frozen=True) class _Reduction: node: torch.fx.Node @@ -142,6 +213,7 @@ def __init__( self, owner: SparseCoreCompute, chunk: LaneChunk, + part: str, reduction_values: dict[torch.fx.Node, ast.AST], *, entry: int | None = None, @@ -149,6 +221,7 @@ def __init__( super().__init__(owner.program.graph, owner.local_codegen) self.owner = owner self.chunk = chunk + self.part = part self.entry = entry self.env.update(reduction_values) @@ -197,7 +270,12 @@ def _stack(self, node: torch.fx.Node) -> ast.AST: assert isinstance(selected, torch.fx.Node) child = _ChunkInterpreter( self.owner, - LaneChunk(local, self.chunk.size), + LaneChunk( + local, + self.chunk.size, + max(0, self.chunk.unique_start - group * value_size), + ), + self.part, { key: value for key, value in self.env.items() @@ -217,11 +295,16 @@ def run_node(self, n: torch.fx.Node) -> object: target = n.target lowering = n.meta.get("lowering") if isinstance(lowering, ReductionLowering): - return self.owner.vector_reduction_expr(n, self.chunk, self.env) + return self.owner.vector_reduction_expr( + n, + self.chunk, + self.part, + self.env, + ) if target in (_host_tensor, _get_symnode): return expr_from_string("0") if target is memory_ops.load: - return self.owner.load_expr(n, self.chunk, self.entry) + return self.owner.load_expr(n, self.chunk, self.part, self.entry) if target is memory_ops.store: return None if target is _mask_to: @@ -263,6 +346,7 @@ def run_node(self, n: torch.fx.Node) -> object: if isinstance(value, torch.Tensor) and value.dtype in ( torch.int8, torch.bool, + torch.bfloat16, ): if any(user.target is not memory_ops.store for user in n.users): raise NotImplementedError( @@ -339,6 +423,30 @@ def new_var(self, prefix: str) -> str: self.counter += 1 return f"_sc_{prefix}{self.counter}" + def value_chunks( + self, node: torch.fx.Node, value_size: int, dtype: torch.dtype + ) -> tuple[list[LaneChunk], bool]: + segments = packed_bf16_segments(self.program, node) + if not segments: + return chunk_schedule(value_size, dtype), False + if sum(segments) != value_size: + raise NotImplementedError( + "SparseCore packed BF16 layout does not cover the logical value" + ) + result: list[LaneChunk] = [] + offset = 0 + for segment in segments: + result.extend( + LaneChunk( + offset + chunk.start, + chunk.size, + offset + chunk.unique_start, + ) + for chunk in chunk_schedule(segment, torch.bfloat16) + ) + offset += segment + return result, True + def _buffer_ref(self, buffer: str, item: str, start: int, size: int) -> str: return f"{buffer}[_sc_q, {item}, pl.ds({start}, {size})]" @@ -346,6 +454,7 @@ def load_expr( self, node: torch.fx.Node, chunk: LaneChunk, + part: str, entry: int | None, ) -> ast.AST: plan = self.program.plan_by_node[node] @@ -373,12 +482,21 @@ def load_expr( start = chunk.start % value_size ref = self._buffer_ref(buffer, "_sc_item", start, chunk.size) - return expr_from_string(ref) + if value.dtype is not torch.bfloat16: + return expr_from_string(ref) + even = self.new_var("e") + odd = self.new_var("o") + self.lines.append( + f"{self.indent}{even}, {odd} = plsc.unpack({ref}, " + "format=plsc.PackFormat.INTERLEAVED)" + ) + return expr_from_string(even if part == "even" else odd) def vector_reduction_expr( self, node: torch.fx.Node, chunk: LaneChunk, + part: str, reduction_values: dict[torch.fx.Node, object], ) -> ast.AST: info = next( @@ -396,7 +514,7 @@ def vector_reduction_expr( if isinstance(value, ast.AST) } for entry in range(info.input_count): - interpreter = _ChunkInterpreter(self, chunk, inherited, entry=entry) + interpreter = _ChunkInterpreter(self, chunk, part, inherited, entry=entry) expressions.append(ast.unparse(interpreter.run_until(info.source))) if info.kind == "sum": return expr_from_string("(" + " + ".join(expressions) + ")") @@ -405,6 +523,12 @@ def vector_reduction_expr( result = f"jnp.maximum({result}, {expression})" return expr_from_string(result) + def _part_mask(self, chunk: LaneChunk, part: str) -> str | None: + if chunk.size != 32 or chunk.unique_start <= chunk.start: + return None + parity = 0 if part == "even" else 1 + return f"({chunk.start + parity} + 2 * _sc_lane) >= {chunk.unique_start}" + def emit_scalar_reductions( self, reduction_values: dict[torch.fx.Node, ast.AST], @@ -413,6 +537,13 @@ def emit_scalar_reductions( for info in self.reductions: if not info.scalar or info.node not in active_nodes: continue + source_value = info.source.meta.get("val") + if not isinstance(source_value, torch.Tensor): + raise NotImplementedError("SparseCore reduction source is not a tensor") + chunks, packed_bf16 = self.value_chunks( + info.source, info.input_size, source_value.dtype + ) + neutral = "0.0" if info.kind == "sum" else "-jnp.inf" acc = self.new_var("acc") init = "jnp.zeros" if info.kind == "sum" else "jnp.full" init_args = "" if info.kind == "sum" else "-jnp.inf, " @@ -420,15 +551,21 @@ def emit_scalar_reductions( f"{self.indent}{acc} = {init}(({SC_LANES},), " f"{init_args}dtype=jnp.float32)" ) - for chunk in chunk_schedule(info.input_size): - interpreter = _ChunkInterpreter(self, chunk, reduction_values) - expression = ast.unparse(interpreter.run_until(info.source)) - if info.kind == "sum": - self.lines.append(f"{self.indent}{acc} = {acc} + {expression}") - else: - self.lines.append( - f"{self.indent}{acc} = jnp.maximum({acc}, {expression})" - ) + for chunk in chunks: + parts: Iterable[str] = ("even", "odd") if packed_bf16 else ("even",) + for part in parts: + interpreter = _ChunkInterpreter(self, chunk, part, reduction_values) + expression = ast.unparse(interpreter.run_until(info.source)) + mask = self._part_mask(chunk, part) + if mask is not None: + expression = f"jnp.where({mask}, {expression}, {neutral})" + combine = "+" if info.kind == "sum" else "jnp.maximum" + if combine == "+": + self.lines.append(f"{self.indent}{acc} = {acc} + {expression}") + else: + self.lines.append( + f"{self.indent}{acc} = jnp.maximum({acc}, {expression})" + ) result = self.new_var("red") aggregate = "jnp.sum" if info.kind == "sum" else "jnp.max" self.lines.append( @@ -440,10 +577,11 @@ def _value( self, node: torch.fx.Node, chunk: LaneChunk, + part: str, reduction_values: dict[torch.fx.Node, ast.AST], ) -> str: return ast.unparse( - _ChunkInterpreter(self, chunk, reduction_values).run_until(node) + _ChunkInterpreter(self, chunk, part, reduction_values).run_until(node) ) def emit_store( @@ -461,26 +599,37 @@ def emit_store( raise NotImplementedError("SparseCore store value is not a tensor") value_size = _value_size(value) if value_size == 1: - chunk = LaneChunk(0, SC_LANES) - expression = self._value(value_node, chunk, reduction_values) + chunk = LaneChunk(0, SC_LANES, 0) + expression = self._value(value_node, chunk, "even", reduction_values) if plan.layout.logical_dtype in _CAST_STORE_DTYPES: expression = f"({expression}).astype(jnp.int32)" self.lines.append(f"{self.indent}{buffer}[_sc_item] = {expression}") return cast_output = plan.layout.logical_dtype in _CAST_STORE_DTYPES - for chunk in chunk_schedule(value_size): + chunks, _ = self.value_chunks(value_node, value_size, value.dtype) + for chunk in chunks: dst = f"{buffer}[_sc_item, pl.ds({chunk.start}, {chunk.size})]" - expression = self._value(value_node, chunk, reduction_values) - if cast_output: + even = self._value(value_node, chunk, "even", reduction_values) + if value.dtype is torch.bfloat16: + odd = self._value(value_node, chunk, "odd", reduction_values) self.lines.append( - f"{self.indent}{dst} = ({expression}).astype(jnp.int32)" + f"{self.indent}{dst} = plsc.pack({even}, {odd}, " + "format=plsc.PackFormat.INTERLEAVED, " + "preferred_element_type=jnp.bfloat16)" ) + elif cast_output: + self.lines.append(f"{self.indent}{dst} = ({even}).astype(jnp.int32)") else: - self.lines.append(f"{self.indent}{dst} = {expression}") + self.lines.append(f"{self.indent}{dst} = {even}") def emit_body(self, indent: str, *, store_nodes: set[torch.fx.Node]) -> list[str]: self.lines = [] self.indent = indent + if any( + plan.layout.logical_dtype is torch.bfloat16 + for plan in self.program.memory_plans + ): + self.lines.append(f"{indent}_sc_lane = lax.iota(jnp.int32, {SC_LANES})") reduction_values: dict[torch.fx.Node, ast.AST] = {} active_nodes: set[torch.fx.Node] = set() diff --git a/test/test_pallas_sparsecore.py b/test/test_pallas_sparsecore.py index 06e77b3d2..c33952f56 100644 --- a/test/test_pallas_sparsecore.py +++ b/test/test_pallas_sparsecore.py @@ -174,6 +174,16 @@ def test_tile_invariant_load_rejects_sparsecore_config(self) -> None: with self.assertRaisesRegex(exc.InvalidConfig, "code: access_pattern"): _code(_scaled_rows, (source, scale)) + def test_bf16_gather_composes_interleaved_pack_unpack(self) -> None: + index = torch.randint(0, 512, (513,), dtype=torch.int32) + table = torch.randn(512, 64, dtype=torch.bfloat16) + + code = _code(_embedding, (index, table), block=512) + + self.assertIn("plsc.unpack", code) + self.assertIn("plsc.pack", code) + self.assertNotIn("jnp.concatenate", code) + def test_indirect_store_uses_layout_derived_tail_row(self) -> None: source = torch.randn(513, 64) index = torch.randperm(513, dtype=torch.int32) @@ -374,11 +384,12 @@ def test_boolean_scalar_output(self) -> None: torch.testing.assert_close(result.cpu(), source.cpu().amax(dim=1) > 0.5) - def test_stock_embedding(self) -> None: + def test_stock_embedding_and_bf16_gather(self) -> None: from examples.embedding import embedding index = torch.randint(0, 2048, (4, 129), dtype=torch.int32, device=DEVICE) table = torch.randn(2048, 64, device=DEVICE) + bf16_table = table.to(torch.bfloat16) _, stock = code_and_output( embedding, @@ -386,8 +397,17 @@ def test_stock_embedding(self) -> None: block_sizes=[256, 64], core_type="sparsecore", ) + _, bf16 = code_and_output( + _embedding, + (index.reshape(-1), bf16_table), + block_sizes=[256], + core_type="sparsecore", + ) torch.testing.assert_close(stock.cpu(), table.cpu()[index.cpu().long()]) + torch.testing.assert_close( + bf16.cpu(), bf16_table.cpu()[index.cpu().reshape(-1).long()] + ) def test_weighted_and_masked_gather_reductions(self) -> None: from examples.sparsecore_ops import masked_gather_sum diff --git a/test/test_pallas_sparsecore_compute.py b/test/test_pallas_sparsecore_compute.py index b55aec650..004f43d57 100644 --- a/test/test_pallas_sparsecore_compute.py +++ b/test/test_pallas_sparsecore_compute.py @@ -1,11 +1,19 @@ from __future__ import annotations +import torch + from helion._compiler.pallas.sparsecore_compute import chunk_schedule def test_lane_chunks() -> None: - assert [(chunk.start, chunk.size) for chunk in chunk_schedule(48)] == [ + assert [ + (chunk.start, chunk.size) for chunk in chunk_schedule(48, torch.float32) + ] == [ (0, 16), (16, 16), (32, 16), ] + assert [ + (chunk.start, chunk.size, chunk.unique_start) + for chunk in chunk_schedule(48, torch.bfloat16) + ] == [(0, 32, 0), (16, 32, 32)] diff --git a/test/test_pallas_sparsecore_plan.py b/test/test_pallas_sparsecore_plan.py index f7026db15..745bd184c 100644 --- a/test/test_pallas_sparsecore_plan.py +++ b/test/test_pallas_sparsecore_plan.py @@ -54,9 +54,9 @@ def test_direct_stream_is_lowered_without_graph_topology() -> None: def test_indirect_load_retains_its_local_dependency() -> None: graph = torch.fx.Graph() - table = torch.empty(4096, 64) + table = torch.empty(4096, 64, dtype=torch.bfloat16) index = torch.empty(32, 4, dtype=torch.int32) - result = torch.empty(32, 4, 64) + result = torch.empty(32, 4, 64, dtype=torch.bfloat16) table_node = _placeholder(graph, "table", table) index_node = _placeholder(graph, "index", index) load = graph.call_function(memory_ops.load, (table_node, (index_node, slice(None))))