diff --git a/helion/_compiler/pallas/sc_base.py b/helion/_compiler/pallas/sc_base.py index abe27bfd6..295312dd1 100644 --- a/helion/_compiler/pallas/sc_base.py +++ b/helion/_compiler/pallas/sc_base.py @@ -104,7 +104,7 @@ 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 444c70d55..0b24c6bbc 100644 --- a/helion/_compiler/pallas/sparsecore_compute.py +++ b/helion/_compiler/pallas/sparsecore_compute.py @@ -16,12 +16,15 @@ 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 CachedLoadPlan from .sparsecore_plan import IndirectLoadPlan if TYPE_CHECKING: + from collections.abc import Iterable + from ..generate_ast import GenerateAST from .sparsecore_program import SparseCoreProgram @@ -60,22 +63,117 @@ 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 not isinstance(plan, CachedLoadPlan) + 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 + + +def classify_cached_loads( + program: SparseCoreProgram, +) -> tuple[frozenset[torch.fx.Node], frozenset[torch.fx.Node]]: + """Cached loads used by packed and ordinary values.""" + packed: set[torch.fx.Node] = set() + ordinary: set[torch.fx.Node] = set() + memo: dict[torch.fx.Node, tuple[int, ...]] = {} + + for plan in program.stores: + root = plan.access.value_node + if root is None: + continue + result = packed if packed_bf16_segments(program, root, memo) else ordinary + pending = [root] + seen: set[torch.fx.Node] = set() + while pending: + node = pending.pop() + if node in seen: + continue + seen.add(node) + if isinstance(program.plan_by_node.get(node), CachedLoadPlan): + result.add(node) + pending.extend(node.all_input_nodes) + return frozenset(packed), frozenset(ordinary) + + @dataclass(frozen=True) class _Reduction: node: torch.fx.Node @@ -143,6 +241,7 @@ def __init__( self, owner: SparseCoreCompute, chunk: LaneChunk, + part: str, reduction_values: dict[torch.fx.Node, ast.AST], *, entry: int | None = None, @@ -150,6 +249,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) @@ -198,7 +298,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() @@ -218,11 +323,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 in (memory_ops.store, self.owner.atomic_add_target): return None if target is _mask_to: @@ -264,6 +374,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( @@ -343,6 +454,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})]" @@ -350,6 +485,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] @@ -360,7 +496,7 @@ def load_expr( value_size = _value_size(value) if isinstance(plan, CachedLoadPlan): - return self._cached_expr(plan.access.tensor, buffer, chunk) + return self._cached_expr(plan.access.tensor, buffer, chunk, part) entries = plan.stream.elements_per_item if plan.stream is not None else 1 if isinstance(plan, IndirectLoadPlan): @@ -380,13 +516,22 @@ 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 _cached_expr( self, tensor: torch.Tensor, buffer: str, chunk: LaneChunk, + part: str, ) -> ast.AST: value_size = int(tensor.shape[-1]) if tensor.ndim else 1 start = chunk.start % value_size @@ -397,12 +542,16 @@ def _cached_expr( prefix = f"{group}, " else: raise NotImplementedError("SparseCore cached inputs support rank 1 or 2") + if part in ("even", "odd") and chunk.size == 32: + half = value_size // 2 + start = start // 2 + (half if part == "odd" else 0) return expr_from_string(f"{buffer}[{prefix}pl.ds({start}, {SC_LANES})]") def vector_reduction_expr( self, node: torch.fx.Node, chunk: LaneChunk, + part: str, reduction_values: dict[torch.fx.Node, object], ) -> ast.AST: info = next( @@ -420,7 +569,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) + ")") @@ -429,6 +578,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], @@ -437,6 +592,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, " @@ -444,15 +606,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( @@ -464,10 +632,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( @@ -485,26 +654,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/helion/_compiler/pallas/sparsecore_launcher.py b/helion/_compiler/pallas/sparsecore_launcher.py index aac355514..d108616c9 100644 --- a/helion/_compiler/pallas/sparsecore_launcher.py +++ b/helion/_compiler/pallas/sparsecore_launcher.py @@ -13,6 +13,7 @@ from .sc_base import SC_LANES from .sc_base import _reject from .sc_base import shared_acc_items +from .sparsecore_compute import classify_cached_loads from .sparsecore_plan import AtomicAddPlan from .sparsecore_plan import CachedLoadPlan from .sparsecore_plan import DirectLoadPlan @@ -103,6 +104,7 @@ def build_sparsecore_launcher_spec( index_padding = _index_padding(program) pad_inputs: list[list[object]] = [] stream_inputs: list[list[object]] = [] + packed_inputs: list[list[object]] = [] seen_inputs: dict[int, tuple[object, ...]] = {} emitted_inputs: set[int] = set() @@ -123,6 +125,7 @@ def record_input( emitted_inputs.add(position) return True + packed_loads, ordinary_loads = classify_cached_loads(program) for plan in program.loads: position = arg_position(plan.access.tensor) if position is None: @@ -169,7 +172,21 @@ def record_input( ] ) elif isinstance(plan, CachedLoadPlan): - record_input(position, ("raw",), plan) + if plan.access.node in packed_loads and plan.access.node in ordinary_loads: + _reject( + "launcher", + "one cached load is used by packed BF16 and ordinary values", + node=plan.access.node, + ) + if plan.access.node in packed_loads: + transform = ( + "packed_cached", + int(plan.access.tensor.shape[-1]), + ) + if record_input(position, transform, plan): + packed_inputs.append([position, int(plan.access.tensor.shape[-1])]) + else: + record_input(position, ("raw",), plan) else: record_input(position, ("raw",), plan) @@ -216,6 +233,7 @@ def record_input( return { "pad_inputs": pad_inputs, "stream_inputs": stream_inputs, + "packed_inputs": packed_inputs, "pad_outputs": pad_outputs, "reshape_outputs": reshape_outputs, "scalar_outputs": scalar_outputs, diff --git a/helion/runtime/pallas/launcher.py b/helion/runtime/pallas/launcher.py index bf795dba4..a7e73b054 100644 --- a/helion/runtime/pallas/launcher.py +++ b/helion/runtime/pallas/launcher.py @@ -1968,6 +1968,13 @@ def _pallas_compile_sc_jit_fn( "list[list[int]]", _sc_launcher_spec.get("stream_inputs") or [] ) ] + packed_inputs = [ + (int(pos), int(value_size)) + for pos, value_size in cast( + "list[list[int]]", + _sc_launcher_spec.get("packed_inputs") or [], + ) + ] reshape_outputs = { int(pos): tuple(int(s) for s in shape) # type: ignore[union-attr] for pos, shape in cast( @@ -2023,6 +2030,14 @@ def jit_fn(*jax_args: object) -> object: if stored_size > value_size: x = jnp.pad(x, ((0, 0), (0, stored_size - value_size))) xs[tpos] = x + for orig_pos, value_size in packed_inputs: + tpos = arg_to_tensor_pos[orig_pos] + x = xs[tpos] + values = jnp.reshape(x, (-1, value_size)) # type: ignore[arg-type] + xs[tpos] = jnp.reshape( + jnp.concatenate([values[:, 0::2], values[:, 1::2]], axis=1), + x.shape, # type: ignore[union-attr] + ) result = kern(*xs) outs: list[Any] = ( list(result) if isinstance(result, (tuple, list)) else [result] diff --git a/test/test_pallas_sparsecore.py b/test/test_pallas_sparsecore.py index f231c8324..b9947282f 100644 --- a/test/test_pallas_sparsecore.py +++ b/test/test_pallas_sparsecore.py @@ -198,6 +198,25 @@ def _grouped_scale( return out +@helion.kernel(backend="pallas", static_shapes=True) +def _bf16_gather_with_side_output( + table: torch.Tensor, + index: torch.Tensor, + scale: torch.Tensor, + side: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + out = torch.empty( + (index.size(0), table.size(1)), dtype=table.dtype, device=table.device + ) + side_out = torch.empty( + (index.size(0), side.size(0)), dtype=side.dtype, device=side.device + ) + for tile in hl.tile(index.size(0)): + out[tile, :] = table[index[tile], :] * scale[None, :] + side_out[tile, :] = side[None, :] + 1 + return out, side_out + + def _code(kernel: object, args: tuple[object, ...], block: int = 256) -> str: with patch.dict("os.environ", {"HELION_SC_ASSUME_MESH": "2x16"}): bound = kernel.bind(args) # type: ignore[attr-defined] @@ -251,6 +270,21 @@ def test_grouped_layer_norm_composes_access_and_compute_lowerings(self) -> None: block=256, ) + def test_packed_bf16_rearranges_only_its_cached_inputs(self) -> None: + code = _code( + _bf16_gather_with_side_output, + ( + torch.randn(512, 64, dtype=torch.bfloat16), + torch.randint(0, 512, (513,), dtype=torch.int32), + torch.randn(64), + torch.randn(16), + ), + block=512, + ) + + self.assertIn("'packed_inputs': [[0, 64]]", code) + self.assertNotIn("[1, 16]", code.split("'packed_inputs':", 1)[1]) + def test_atomic_add_uses_shared_accumulator_phases(self) -> None: source = torch.randn(1024, 64) index = torch.randint(0, 128, (1024,), dtype=torch.int32) @@ -271,6 +305,25 @@ def test_atomic_add_requires_zero_initialized_output(self) -> None: with self.assertRaisesRegex(exc.InvalidConfig, "code: atomic_init"): _code(_scatter_add_onto_ones, (source, index, 128), block=512) + def test_bf16_chunks_compose_across_stack_boundaries(self) -> None: + index = torch.randint(0, 512, (2, 513), dtype=torch.int32) + code = _code( + _grouped_layer_norm, + ( + torch.randn(1024, 48, dtype=torch.bfloat16), + index, + torch.randn(2, 48), + torch.randn(2, 48), + ), + block=512, + ) + + self.assertIn("plsc.unpack", code) + self.assertIn("plsc.pack", code) + self.assertIn("sc_output[_sc_item, pl.ds(48, 32)]", code) + self.assertNotIn("jnp.concatenate", code) + self.assertNotIn("one_hot", code) + def test_indirect_store_uses_layout_derived_tail_row(self) -> None: source = torch.randn(513, 64) index = torch.randperm(513, dtype=torch.int32) @@ -471,11 +524,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, @@ -483,8 +537,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 @@ -660,3 +723,37 @@ def test_scatter_add_computed_values(self) -> None: expected = torch.zeros(output_rows, width) expected.index_add_(0, index.cpu().long(), source.cpu() * 1.25) torch.testing.assert_close(result.cpu(), expected, rtol=1e-5, atol=1e-4) + + def test_grouped_layer_norm_bf16(self) -> None: + groups, table_rows, width, rows = 2, 512, 48, 513 + tables = torch.randn( + groups * table_rows, width, dtype=torch.bfloat16, device=DEVICE + ) + index = torch.randint( + 0, table_rows, (groups, rows), dtype=torch.int32, device=DEVICE + ) + weight = torch.randn(groups, width, device=DEVICE) + bias = torch.randn(groups, width, device=DEVICE) + + _, result = code_and_output( + _grouped_layer_norm, + (tables, index, weight, bias), + block_sizes=[512], + core_type="sparsecore", + ) + + gathered = torch.stack( + [ + tables.cpu().float()[index.cpu()[group].long() + group * table_rows] + for group in range(groups) + ], + dim=1, + ) + flat = gathered.reshape(rows, groups * width) + mean = flat.mean(dim=1, keepdim=True) + variance = ((flat - mean) ** 2).mean(dim=1, keepdim=True) + expected = ((flat - mean) * torch.rsqrt(variance + 1e-5)).reshape( + rows, groups, width + ) + expected = expected * weight.cpu()[None] + bias.cpu()[None] + self.assertLess((result.cpu().float() - expected).abs().max().item(), 0.11) 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 547b0eb69..da45570e6 100644 --- a/test/test_pallas_sparsecore_plan.py +++ b/test/test_pallas_sparsecore_plan.py @@ -53,9 +53,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))))