Add initial support for joint sparsity pipeline

src/coreai_opt/palettization/kmeans/_prepare_for_export.py

@@ -6,6 +6,7 @@
 from dataclasses import dataclass
 from os import PathLike
 from pathlib import Path
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -23,6 +24,7 @@
 from coreai_opt.palettization.spec.fake_palettize import (
     _FakePalettizeImplBase,
 )
+from coreai_opt.palettization.spec.granularity import PerTensorGranularity
 
 _DEFAULT_VECTOR_AXIS = 0
 
@@ -44,6 +46,34 @@ class PalettizationInfo:
     lut_quantization: LUTQuantizationInfo | None = None
 
 
+class _SparsePalettizeReconstruction(nn.Module):
+    """Parametrization module inserted to reconstruct a sparse-palettized weight.
+
+    Traces ``coreai.lut_to_dense`` and ``coreai.sparse_to_dense`` in that order.
+    """
+
+    def __init__(
+        self,
+        nonzero_indices: torch.Tensor,
+        lut: torch.Tensor,
+        mask: torch.Tensor,
+        vector_axis: int | None,
+    ) -> None:
+        super().__init__()
+        self.register_buffer("nonzero_indices", nonzero_indices)
+        # nonzero_indices is rank 1 (flattened by masking), so lut must be
+        # reshaped to rank 3 (lut_to_dense requires lut.rank == indices.rank + 2).
+        self.register_buffer("lut", lut.reshape(1, lut.shape[-2], lut.shape[-1]))
+        self.register_buffer("mask", mask)
+        self.vector_axis = 0 if vector_axis is None else vector_axis
+
+    def forward(self, _: Any) -> torch.Tensor:
+        nonzero_values = torch.ops.coreai.lut_to_dense(
+            self.nonzero_indices, self.lut, self.vector_axis
+        )
+        return torch.ops.coreai.sparse_to_dense(nonzero_values, self.mask)
+
+
 def _expand_rank(
     tensor: torch.Tensor,
     target_rank: int,
@@ -116,6 +146,7 @@ def _register_mil_compression_metadata(
     module: nn.Module,
     param_name: str,
     palett_info: PalettizationInfo,
+    fake_palett_mod: _FakePalettizeImplBase,
 ) -> None:
     """
     Remove the fake palettization parametrization from the module
@@ -132,12 +163,20 @@ def _register_mil_compression_metadata(
         leave_parametrized=True,
     )
 
-    # Determine compression type(s)
+    # PRUNING must be listed first for coremltools to chain the sparse LUT op.
     lut_quant = palett_info.lut_quantization
     if lut_quant is not None:
         compression_type = [CompressionType.PALETTIZATION, CompressionType.QUANTIZATION]
     else:
-        compression_type = CompressionType.PALETTIZATION
+        compression_type = [CompressionType.PALETTIZATION]
+    if fake_palett_mod.sparsity is not None:
+        # Vector palettization & per-channel scales aren't supported jointly with sparsity.
+        if fake_palett_mod.cluster_dim != 1 or fake_palett_mod.enable_per_channel_scale:
+            raise ValueError(
+                "cluster_dim != 1 (vector palettization) and enable_per_channel_scale "
+                "are not supported for joint sparsity + palettization."
+            )
+        compression_type = [CompressionType.PRUNING, *compression_type]
 
     metadata = MILCompressionMetadata(
         param_name=param_name,
@@ -203,11 +242,29 @@ def _resolve_mlir_lut_and_scale(
     return lut, scale, offset
 
 
+def _validate_sparsity_for_export(fake_palett_mod: _FakePalettizeImplBase) -> None:
+    """Reject sparsity combined with anything that isn't a single, position-independent LUT."""
+    if not isinstance(fake_palett_mod.granularity, PerTensorGranularity):
+        raise ValueError(
+            f"granularity={fake_palett_mod.granularity} not supported for joint sparsity."
+        )
+    if fake_palett_mod.cluster_dim != 1:
+        raise ValueError(
+            f"cluster_dim={fake_palett_mod.cluster_dim} (vector palettization) "
+            "not supported for joint sparsity."
+        )
+    if fake_palett_mod.lut_qspec is not None:
+        raise ValueError("lut_qspec not supported for joint sparsity.")
+    if fake_palett_mod.enable_per_channel_scale:
+        raise ValueError("enable_per_channel_scale not supported for joint sparsity.")
+
+
 def _insert_mlir_custom_op(
     module: nn.Module,
     module_name: str,
     param_name: str,
     palett_info: PalettizationInfo,
+    fake_palett_mod: _FakePalettizeImplBase,
     fake_palett_idx: int,
     mmap_dir: str | PathLike[str] | None,
 ) -> None:
@@ -225,6 +282,10 @@ def _insert_mlir_custom_op(
     4. Both: lut_to_dense(int LUT) + constexpr_blockwise_shift_scale(fused_scale)
        where fused_scale = lut_scale * per_channel_scale
 
+    When ``fake_palett_mod.sparsity`` is set, the LUT lookup runs on the
+    nonzero-only indices and the result is packed via ``coreai::sparse_to_dense``
+    instead of installing a plain Palettize/ScaledPalettize parametrization.
+
     When ``mmap_dir`` is provided, the new MLIR module is serialized to a
     safetensors file under that directory and reloaded via mmap before being
     swapped in.
@@ -259,7 +320,18 @@ def _import_coreai_torch_modules():
 
     vector_axis = _DEFAULT_VECTOR_AXIS if palett_info.cluster_dim > 1 else None
 
-    if needs_scale:
+    if fake_palett_mod.sparsity is not None:
+        _validate_sparsity_for_export(fake_palett_mod)
+        # Reuses the mask from prepare()'s forward pass.
+        mask = fake_palett_mod._sparsity_mask.to(torch.bool)
+        nonzero_indices = palett_info.indices[mask]
+        mlir_palett_mod = _SparsePalettizeReconstruction(
+            nonzero_indices=nonzero_indices,
+            lut=palett_info.lut,
+            mask=mask,
+            vector_axis=vector_axis,
+        )
+    elif needs_scale:
         lut, scale, zero_point = _resolve_mlir_lut_and_scale(palett_info)
         mlir_palett_mod = ScaledPalettizeParametrization(
             indices=palett_info.indices,
@@ -331,10 +403,10 @@ def _process_palettized_parameter(
     )
 
     if backend == ExportBackend.CoreML:
-        _register_mil_compression_metadata(module, param_name, palett_info)
+        _register_mil_compression_metadata(module, param_name, palett_info, fake_palett_mod)
     elif backend == ExportBackend.CoreAI:
         _insert_mlir_custom_op(
-            module, module_name, param_name, palett_info, fake_palett_idx, mmap_dir
+            module, module_name, param_name, palett_info, fake_palett_mod, fake_palett_idx, mmap_dir
         )
 
 

src/coreai_opt/palettization/kmeans/kmeans_fake_palettize.py

@@ -100,13 +100,15 @@ def __init__(
         rounding_precision: int = 4,
         op_to_optimize: Callable | None = None,
         training_strategy_spec: TrainingStrategySpec | None = None,
+        sparsity: float | None = None,
     ):
         super().__init__(
             n_bits=n_bits,
             lut_qspec=lut_qspec,
             granularity=granularity,
             cluster_dim=cluster_dim,
             enable_per_channel_scale=enable_per_channel_scale,
+            sparsity=sparsity,
         )
 
         self.enable_fast_kmeans_mode = enable_fast_kmeans_mode

src/coreai_opt/palettization/kmeans/palettizer.py

@@ -493,6 +493,7 @@ def _spec_to_partial(
         # Serialize the spec, then layer in the owning module's compressor-specific
         # settings (e.g. enable_fast_kmeans_mode, rounding_precision).
         args = spec.model_dump_preserve_objects()
+        args["sparsity"] = spec._sparsity
         args.update(module_config._get_fake_module_kwargs())
         return _KMeansFakePalettize.with_args(**args)
 

src/coreai_opt/palettization/spec/fake_palettize.py

@@ -18,6 +18,7 @@
 from coreai_opt.palettization.spec import (
     PalettizationGranularity,
 )
+from coreai_opt.pruning.spec import PruneImplBase, Unstructured
 from coreai_opt.quantization.spec import QuantizationSpec
 
 
@@ -38,6 +39,7 @@ def __init__(
         granularity: PalettizationGranularity,
         cluster_dim: int,
         enable_per_channel_scale: bool,
+        sparsity: float | None = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -46,6 +48,7 @@ def __init__(
         self.granularity = granularity
         self.cluster_dim = cluster_dim
         self.enable_per_channel_scale = enable_per_channel_scale
+        self.sparsity = sparsity
 
         self.register_buffer("fake_palett_enabled", torch.tensor([1], dtype=torch.uint8))
         # Non-persistent (kept out of new checkpoints); when set to 1 (at runtime or
@@ -54,6 +57,7 @@ def __init__(
             "observer_enabled", torch.tensor([0], dtype=torch.uint8), persistent=False
         )
         self._disabled = False
+        self.register_buffer("_sparsity_mask", None, persistent=False)
 
         self.register_buffer("indices", None)
         self.register_buffer("per_channel_scale", None)
@@ -71,6 +75,12 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         if self._disabled:
             return tensor
 
+        if self.sparsity is not None:
+            self._sparsity_mask = PruneImplBase.resolve("default").compute_mask(
+                tensor, self.sparsity, Unstructured()
+            )
+            tensor = tensor * self._sparsity_mask
+
         self.ensure_initialized(tensor)
 
         # Check for self._disabled again in case ensure_initialized disabled the palettizer.

src/coreai_opt/palettization/spec/spec.py

@@ -108,6 +108,17 @@ class PalettizationSpec(CompressionSpec):
     # Private attribute for compression type
     _compression_type: CompressionType = PrivateAttr(default=CompressionType.PALETTIZATION)
 
+    # Sparsity level, in [0, 1]. Set via the `_sparsity` constructor/dict key.
+    _sparsity: float | None = PrivateAttr(default=None)
+
+    def __init__(self, **data: Any) -> None:
+        sparsity = data.pop("_sparsity", None)
+        super().__init__(**data)
+        if sparsity is not None:
+            if not (0.0 <= sparsity <= 1.0):
+                raise ValueError(f"_sparsity must be in [0, 1], got {sparsity}")
+            self._sparsity = sparsity
+
     @model_validator(mode="after")
     def validate_lut_qspec(self) -> "PalettizationSpec":
         """Validate that lut_qspec only uses supported configurations."""

src/coreai_opt/quantization/_graph/_prepare_for_export.py

@@ -232,6 +232,16 @@ def _register_quantization_buffers(
     return buffer_names
 
 
+def _validate_sparsity_for_export(
+    fake_quant_mod: FakeQuantizeImplBase, zero_point: torch.Tensor | None
+) -> None:
+    """Reject sparsity combined with anything raw-0 padding can't represent."""
+    if is_float4_dtype(fake_quant_mod.dtype):
+        raise ValueError("FP4 dtype not supported for joint sparsity.")
+    if zero_point is not None and not torch.all(zero_point == 0):
+        raise ValueError("nonzero zero_point not supported for joint sparsity.")
+
+
 def _process_mlir_weight_quantization(
     model: torch.fx.GraphModule,
     node: Node,
@@ -269,8 +279,14 @@ def _import_coreai_custom_ops():
     if minval is not None:
         minval = minval.to(dtype=_compute_dtype_for_export)
 
-    # Construct quantized weights
+    # Construct quantized weights, reusing the mask computed during
+    # prepare()'s forward pass if sparsity is set.
     dense_weight = resolve_attr(model, input_node.target).data
+    mask: torch.Tensor | None = None
+    if fake_quant_mod.sparsity is not None:
+        _validate_sparsity_for_export(fake_quant_mod, zero_point)
+        mask = fake_quant_mod._sparsity_mask.to(torch.bool)
+        dense_weight = dense_weight * mask
     quantized_data = fake_quant_mod.quantize(dense_weight, scale, zero_point, minval)
 
     # Drop one of the offsets so that the export
@@ -286,13 +302,28 @@ def _import_coreai_custom_ops():
 
     # Register buffers and get buffer names
     param_name = str(input_node.target).replace(".", "_")
-    buffer_names = _register_quantization_buffers(
-        model, param_name, scale, zero_point, quantized_data, minval
-    )
+    if mask is not None:
+        nonzero_data = quantized_data[mask]
+        buffer_names = _register_quantization_buffers(
+            model, param_name, scale, zero_point, minval=minval
+        )
+        model.register_buffer(f"{param_name}_nonzero", nonzero_data)
+        model.register_buffer(f"{param_name}_mask", mask)
+    else:
+        buffer_names = _register_quantization_buffers(
+            model, param_name, scale, zero_point, quantized_data, minval
+        )
 
     # Create graph nodes and replace fake quantization
     with model.graph.inserting_before(node):
-        quantized_data_node = model.graph.get_attr(buffer_names["quantized_data"])
+        if mask is not None:
+            nonzero_node = model.graph.get_attr(f"{param_name}_nonzero")
+            mask_node = model.graph.get_attr(f"{param_name}_mask")
+            quantized_data_node = model.graph.call_function(
+                coreai.sparse_to_dense, (nonzero_node, mask_node)
+            )
+        else:
+            quantized_data_node = model.graph.get_attr(buffer_names["quantized_data"])
         scale_node = model.graph.get_attr(buffer_names["scale"])
 
         if zero_point is not None:
@@ -439,10 +470,22 @@ def _process_mil_weight_quantization(
     # Get the module that owns the weight parameter
     weight_module: torch.nn.Module = _get_weight_module(modules, module_name)
 
+    # coremltools' own torch-frontend converter auto-detects sparsity from raw
+    # zeros in the registered weight value when compression_type lists PRUNING
+    # first, then chains QUANTIZATION onto its constexpr_sparse_to_dense output
+    # (see coremltools.converters.mil.frontend.torch.converter._construct_compression_op).
+    # Unlike CoreAI's sparse_to_dense, that chain dequantizes the compact
+    # nonzero_data before scattering it into the padded dense tensor, so the
+    # padding is a real float 0.0, not a raw int later reinterpreted through
+    # zero_point -- safe for any zero_point, so there's nothing to gate here.
+    compression_type = [CompressionType.QUANTIZATION]
+    if fake_quant_mod.sparsity is not None:
+        compression_type = [CompressionType.PRUNING, CompressionType.QUANTIZATION]
+
     # Create and register metadata
     metadata = MILCompressionMetadata(
         param_name=param_name,
-        compression_type=CompressionType.QUANTIZATION,
+        compression_type=compression_type,
         quantization_n_bits=fake_quant_mod.n_bits,
         quantization_scale=scale,
         zero_point=zero_point,

src/coreai_opt/quantization/_graph/_provisional_qspec_generation.py

@@ -64,6 +64,7 @@
     FieldName.QPARAM_CALCULATOR_CLS: "qparam_calculator_cls",
     FieldName.RANGE_CALCULATOR_CLS: "range_calculator_cls",
     FieldName.SCALE_DTYPE: "scale_dtype",
+    FieldName.SPARSITY: "_sparsity",
 }
 
 

src/coreai_opt/quantization/_graph/_qspec_constraints.py

@@ -320,6 +320,7 @@ def _policy_float_range_union(proposals: Sequence[FieldValue]) -> FieldValue:
     FieldName.QPARAM_CALCULATOR_CLS: _policy_priority_wins,
     FieldName.RANGE_CALCULATOR_CLS: _policy_priority_wins,
     FieldName.SCALE_DTYPE: _policy_priority_wins,
+    FieldName.SPARSITY: _policy_priority_wins,
     # Covering every member's values is a correctness constraint, not a
     # preference, so this one unions instead of deferring to priority.
     FieldName.FLOAT_RANGE: _policy_float_range_union,

src/coreai_opt/quantization/_graph/_qspec_resolution.py

@@ -178,6 +178,7 @@ def _shared_spec_pointing_at(anchor: NodeSlot) -> _SharedQuantizationSpec:
     FieldName.QPARAM_CALCULATOR_CLS: "qparam_calculator_cls",
     FieldName.RANGE_CALCULATOR_CLS: "range_calculator_cls",
     FieldName.SCALE_DTYPE: "scale_dtype",
+    FieldName.SPARSITY: "_sparsity",
 }
 
 

src/coreai_opt/quantization/_graph/_qspec_types.py

@@ -61,6 +61,7 @@ class FieldName(enum.Enum):
     QPARAM_CALCULATOR_CLS = enum.auto()
     RANGE_CALCULATOR_CLS = enum.auto()
     SCALE_DTYPE = enum.auto()
+    SPARSITY = enum.auto()  # QuantizationSpec's private, settable `_sparsity` input
     # Weight or activation, set by which config dict the spec came from. Not a
     # QuantizationSpec attribute, but an input to construct_partial alongside
     # them, so resolution needs it to rebuild the observer.

src/coreai_opt/quantization/spec/factory.py

@@ -209,6 +209,7 @@ def create_fake_quantizer(
             "quant_max": spec.quant_max,
             "qparams_calculator": qparams_calculator,
             "n_bits": spec.n_bits,
+            "sparsity": spec._sparsity,
         }
 
         # Automatically detect and include any extra arguments
@@ -289,6 +290,7 @@ def create_fake_quantizer_partial(
             "quant_min": spec.quant_min,
             "quant_max": spec.quant_max,
             "n_bits": spec.n_bits,
+            "sparsity": spec._sparsity,
         }
 
         # Automatically detect and include any extra arguments