Add sparsity-aware modelling pipeline

Dockerfile

@@ -46,9 +46,10 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
             && rm -rf /var/lib/apt/lists/*
 
 # Update certificates (needed for downloading)
-RUN apt-get upgrade -y ca-certificates && \
-    update-ca-certificates && \
-    rm -rf /var/lib/apt/lists/*
+RUN apt-get upgrade -y ca-certificates \
+    && update-ca-certificates \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/* /var/cache/apt/archives/*
 
 WORKDIR /home/build
 COPY Makefile ./

Makefile

@@ -1,11 +1,12 @@
 DOCKER_EXE ?= docker
 DOCKER_NAME ?= accelforge
-DOCKER_BUILD ?= ${DOCKER_EXE} buildx build --load
+DOCKER_BUILD ?= ${DOCKER_EXE} buildx build --load --pull
 
-VERSION := 0.1.3
+VERSION := 0.1.4
 
 USER    := timeloopaccelergy
 REPO    := accelforge
+INFRA_REPO := accelforge-extra
 
 NAME    := ${USER}/${REPO}
 TAG     := $$(git log -1 --pretty=%h)
@@ -14,6 +15,10 @@ IMG     := ${NAME}:${TAG}
 ALTTAG  := latest
 ALTIMG  := ${NAME}:${ALTTAG}
 
+INFRA_NAME    := ${USER}/${INFRA_REPO}
+INFRA_IMG     := ${INFRA_NAME}:${TAG}
+INFRA_ALTIMG  := ${INFRA_NAME}:${ALTTAG}
+
 # Install hwcomponents packages from PyPI for Docker builds.
 .PHONY: install-hwcomponents
 install-hwcomponents:
@@ -32,6 +37,12 @@ build-amd64:
           -t ${IMG}-amd64 \
           -t ${ALTIMG}-amd64 .
 
+build-extra-amd64:
+	${DOCKER_BUILD} ${BUILD_FLAGS} --platform linux/amd64 \
+          -f infrastructure/Dockerfile \
+          -t ${INFRA_IMG}-amd64 \
+          -t ${INFRA_ALTIMG}-amd64 .
+
 build-arm64:
 	${DOCKER_BUILD} ${BUILD_FLAGS} --platform linux/arm64 \
           --build-arg BUILD_DATE=`date -u +"%Y-%m-%dT%H:%M:%SZ"` \
@@ -40,29 +51,72 @@ build-arm64:
           -t ${IMG}-arm64 \
           -t ${ALTIMG}-arm64 .
 
+build-extra-arm64:
+	${DOCKER_BUILD} ${BUILD_FLAGS} --platform linux/arm64 \
+          -f infrastructure/Dockerfile \
+          -t ${INFRA_IMG}-arm64 \
+          -t ${INFRA_ALTIMG}-arm64 .
+
 # Push docker image
 push-amd64:
 	@echo "Pushing ${NAME}:${ALTTAG}-amd64"
 	#Push Amd64 version
 	"${DOCKER_EXE}" push ${NAME}:${ALTTAG}-amd64
 	#Combine Amd64 version into multi-architecture docker image.
+	"${DOCKER_EXE}" manifest rm ${NAME}:${ALTTAG} || true
 	"${DOCKER_EXE}" manifest create \
 		${NAME}:${ALTTAG} \
-		--amend ${NAME}:${ALTTAG}-amd64 \
-	  --amend ${NAME}:${ALTTAG}-arm64
+		${NAME}:${ALTTAG}-amd64 \
+		${NAME}:${ALTTAG}-arm64
 	"${DOCKER_EXE}" manifest push ${NAME}:${ALTTAG}
+	@echo "Pushing ${INFRA_NAME}:${ALTTAG}-amd64"
+
+
+push-extra-amd64:
+	@echo "Pushing ${INFRA_NAME}:${ALTTAG}-amd64"
+	"${DOCKER_EXE}" push ${INFRA_NAME}:${ALTTAG}-amd64
+	#Combine Amd64 infrastructure version into multi-architecture docker image.
+	"${DOCKER_EXE}" manifest rm ${INFRA_NAME}:${ALTTAG} || true
+	"${DOCKER_EXE}" manifest create \
+		${INFRA_NAME}:${ALTTAG} \
+		${INFRA_NAME}:${ALTTAG}-amd64 \
+		${INFRA_NAME}:${ALTTAG}-arm64
+	"${DOCKER_EXE}" manifest push ${INFRA_NAME}:${ALTTAG}
 
 push-arm64:
 	@echo "Pushing ${NAME}:${ALTTAG}-arm64"
 	#Push Arm64 version
 	"${DOCKER_EXE}" push ${NAME}:${ALTTAG}-arm64
 	#Combine Arm64 version into multi-architecture docker image.
+	"${DOCKER_EXE}" manifest rm ${NAME}:${ALTTAG} || true
 	"${DOCKER_EXE}" manifest create \
 		${NAME}:${ALTTAG} \
-		--amend ${NAME}:${ALTTAG}-amd64 \
-	  --amend ${NAME}:${ALTTAG}-arm64
+		${NAME}:${ALTTAG}-amd64 \
+		${NAME}:${ALTTAG}-arm64
 	"${DOCKER_EXE}" manifest push ${NAME}:${ALTTAG}
 
+push-extra-arm64:
+	@echo "Pushing ${INFRA_NAME}:${ALTTAG}-arm64"
+	#Push Arm64 infrastructure version
+	"${DOCKER_EXE}" push ${INFRA_NAME}:${ALTTAG}-arm64
+	#Combine Arm64 infrastructure version into multi-architecture docker image.
+	"${DOCKER_EXE}" manifest rm ${INFRA_NAME}:${ALTTAG} || true
+	"${DOCKER_EXE}" manifest create \
+		${INFRA_NAME}:${ALTTAG} \
+		${INFRA_NAME}:${ALTTAG}-amd64 \
+		${INFRA_NAME}:${ALTTAG}-arm64
+	"${DOCKER_EXE}" manifest push ${INFRA_NAME}:${ALTTAG}
+
+all-infra:
+	make build-arm64
+	make build-amd64
+	make push-arm64
+	make push-amd64
+	make build-extra-arm64
+	make build-extra-amd64
+	make push-extra-arm64
+	make push-extra-amd64
+
 run-docker:
 	docker-compose up
 

accelforge/frontend/arch/components.py

@@ -29,6 +29,11 @@
 from accelforge.util._eval_expressions import eval_expression
 from accelforge.util._setexpressions import InvertibleSet, eval_set_expression
 from accelforge.frontend.renames import TensorName
+from accelforge.frontend.sparse import (
+    RepresentationFormat,
+    ActionOptimization,
+    ComputeOptimization,
+)
 from accelforge.frontend.arch.constraints import Comparison
 from accelforge.frontend.arch.structure import ArchNode, Branch, Leaf
 from accelforge.frontend.arch.spatialable import Spatial, Spatialable
@@ -893,6 +898,14 @@ class TensorHolder(Component, Leaf):
     value for the bits_per_action of all actions of this component.
     """
 
+    representation_format: EvalableList[RepresentationFormat] = EvalableList()
+    """Compressed representation formats for tensors at this storage level.
+    Inline alternative to specifying in a separate sparse_optimizations file."""
+
+    action_optimization: EvalableList[ActionOptimization] = EvalableList()
+    """Storage action optimizations (gating/skipping) at this level.
+    Inline alternative to specifying in a separate sparse_optimizations file."""
+
     def model_post_init(self, __context__=None) -> None:
         self._update_actions(MEMORY_ACTIONS)
 
@@ -999,6 +1012,10 @@ class Compute(Component, Leaf):
     actions: EvalableList[Action] = COMPUTE_ACTIONS
     """ The actions that this `Compute` can perform. """
 
+    compute_optimization: EvalableList[ComputeOptimization] = EvalableList()
+    """Compute-level optimizations (gating/skipping at the MAC).
+    Inline alternative to specifying in a separate sparse_optimizations file."""
+
     def model_post_init(self, __context__=None) -> None:
         self._update_actions(COMPUTE_ACTIONS)
 

accelforge/frontend/sparse.py

@@ -0,0 +1,191 @@
+"""Sparse optimization specification for AccelForge."""
+
+from typing import Literal, Optional
+
+from pydantic import Field
+
+from accelforge.util._basetypes import EvalableModel, EvalableList
+
+
+class RankFormat(EvalableModel):
+    """Per-rank format specification for explicit (expert) format definitions."""
+
+    format: str
+    """ Format primitive name: UOP, CP, B, or RLE. """
+
+    metadata_word_bits: Optional[int] = None
+    """ Bits per metadata word. None = auto-derived from format primitive. """
+
+    payload_word_bits: Optional[int] = None
+    """ Bits per payload word. None = auto-derived from dimension size. """
+
+    flattened_rank_ids: Optional[list[list[str]]] = None
+    """ Dimension names flattened into this rank, e.g. [["C", "R"]]. """
+
+    def model_post_init(self, __context__=None) -> None:
+        if self.format.upper() not in ("UOP", "CP", "B", "RLE"):
+            raise ValueError(
+                f"Unknown format primitive {self.format!r}. "
+                f"Expected one of: UOP, CP, B, RLE"
+            )
+
+
+class RepresentationFormat(EvalableModel):
+    """Per-tensor compressed format at a storage level.
+
+    Specify ``format`` as one of: csr, coo, bitmask, rle.
+    """
+
+    _VALID_FORMATS = {"csr", "coo", "bitmask", "b", "rle"}
+
+    name: str
+    """ Tensor name (must match a tensor in the workload). """
+
+    format: Optional[str] = None
+    """ User-friendly format name (csr, coo, bitmask, rle), auto-expanded to per-rank primitives. """
+
+    ranks: Optional[EvalableList[RankFormat]] = Field(None, exclude=True)
+    """ Explicit per-rank format specification (internal), outer-to-inner. """
+
+    metadata_word_bits: Optional[int] = None
+    """ Default bits per metadata word for auto-expanded ranks. None = auto-derived per rank. """
+
+    metadata_storage_width: Optional[int] = None
+    """ Physical SRAM width in bits for metadata packing. None = fall back to arch. """
+
+    uop_payload_word_bits: Optional[int] = None
+    """ Override payload_word_bits for auto-expanded UOP ranks. None = auto-derived. """
+
+    def has_explicit_ranks(self) -> bool:
+        """True if explicit per-rank formats were provided (internal)."""
+        return self.ranks is not None
+
+    def model_post_init(self, __context__=None) -> None:
+        if self.format is not None and self.format.lower() not in self._VALID_FORMATS:
+            raise ValueError(
+                f"Unknown format {self.format!r}. "
+                f"Expected one of: csr, coo, bitmask, rle"
+            )
+
+    def get_rank_formats(self, num_ranks: Optional[int] = None) -> list[RankFormat]:
+        """Return per-rank formats, auto-expanding from ``format`` if needed."""
+        if self.ranks is not None:
+            return list(self.ranks)
+        if self.format is None:
+            return []
+        if num_ranks is None:
+            raise ValueError(
+                f"num_ranks required to auto-expand format {self.format!r} "
+                f"for tensor {self.name}"
+            )
+        from accelforge.model.sparse_formats import expand_format
+
+        primitives = expand_format(self.format, num_ranks)
+        result = []
+        for p in primitives:
+            if p.upper() == "UOP" and self.uop_payload_word_bits is not None:
+                result.append(RankFormat(format=p, payload_word_bits=self.uop_payload_word_bits))
+            else:
+                result.append(RankFormat(format=p))
+        return result
+
+
+class ActionOptimization(EvalableModel):
+    """Storage action optimization at a memory level."""
+
+    kind: Literal["gating", "skipping", "position_skipping"]
+    """ Optimization type: gating (filter after access), skipping (skip access), or position_skipping (self-conditioned skip). """
+
+    target: str
+    """ Tensor whose read accesses are reduced. """
+
+    condition_on: list[str]
+    """ Tensors whose sparsity determines the filtering probability. Empty for position_skipping. """
+
+    def model_post_init(self, __context__=None) -> None:
+        if self.kind == "position_skipping" and self.condition_on:
+            raise ValueError(
+                f"position_skipping requires condition_on=[], "
+                f"got {self.condition_on!r}"
+            )
+
+    @property
+    def is_self_conditioned(self) -> bool:
+        """True when the optimization is position-skipping (self-conditioned)."""
+        return self.kind == "position_skipping" and not self.condition_on
+
+
+class ComputeOptimization(EvalableModel):
+    """Compute-level optimization (gating or skipping at the MAC)."""
+
+    kind: Literal["gating", "skipping"]
+    """ Optimization type: gating (discard result) or skipping (skip entirely). """
+
+    target: str
+    """ Target tensor or operation name (e.g., Z, GEMM). """
+
+    condition_on: list[str]
+    """ Operand tensors for compute classification. """
+
+
+class SparseTarget(EvalableModel):
+    """Sparse optimization configuration for one hardware component."""
+
+    target: str
+    """ Component name from arch YAML (e.g., DRAM, Buffer, Reg, MAC). """
+
+    representation_format: EvalableList[RepresentationFormat] = EvalableList()
+    """ Compressed formats for tensors at this level. """
+
+    action_optimization: EvalableList[ActionOptimization] = EvalableList()
+    """ Storage action filtering optimizations at this level. """
+
+    compute_optimization: EvalableList[ComputeOptimization] = EvalableList()
+    """ Compute-level optimizations (only meaningful on Compute nodes). """
+
+
+class SparseOptimizations(EvalableModel):
+    """Top-level sparse optimizations specification."""
+
+    targets: EvalableList[SparseTarget] = EvalableList()
+    """ Per-component sparse optimization configurations. """
+
+    def get_targets_for(self, component_name: str) -> list[SparseTarget]:
+        """Return all SparseTarget entries matching a component name."""
+        return [t for t in self.targets if t.target == component_name]
+
+    def get_formats_for(
+        self, component_name: str, tensor_name: str
+    ) -> list[RepresentationFormat]:
+        """Return all RepresentationFormat entries for a (component, tensor) pair."""
+        results = []
+        for t in self.get_targets_for(component_name):
+            for rf in t.representation_format:
+                if rf.name == tensor_name:
+                    results.append(rf)
+        return results
+
+    def get_action_optimizations_for(
+        self, component_name: str
+    ) -> list[ActionOptimization]:
+        """Return all ActionOptimization entries for a component."""
+        results = []
+        for t in self.get_targets_for(component_name):
+            results.extend(t.action_optimization)
+        return results
+
+    def get_compute_optimizations_for(
+        self, component_name: str
+    ) -> list[ComputeOptimization]:
+        """Return all ComputeOptimization entries for a component."""
+        results = []
+        for t in self.get_targets_for(component_name):
+            results.extend(t.compute_optimization)
+        return results
+
+    def has_format(self, component_name: str, tensor_name: str) -> bool:
+        """True if the tensor has a compressed format at the component."""
+        return any(
+            rf.format is not None or rf.has_explicit_ranks()
+            for rf in self.get_formats_for(component_name, tensor_name)
+        )