From 25af8cd4d2a4df3bbd8298cecf276c262e5c0220 Mon Sep 17 00:00:00 2001 From: Guru Deshpande <66385690+guru-desh@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:49:37 -0700 Subject: [PATCH 1/3] refactor(coreai_utils): use raw CoreAI ops in the sparsity pass --- src/coreai_opt/coreai_utils/_coreai_imports.py | 6 ------ .../coreai_utils/passes/weight_sparsification.py | 9 ++------- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/src/coreai_opt/coreai_utils/_coreai_imports.py b/src/coreai_opt/coreai_utils/_coreai_imports.py index aa84c638..dcb6a4a9 100644 --- a/src/coreai_opt/coreai_utils/_coreai_imports.py +++ b/src/coreai_opt/coreai_utils/_coreai_imports.py @@ -137,7 +137,6 @@ def _import_coreai() -> tuple: DebugInfo, TensorType, blockwise_shift_scale, - build_sparse_with_bitmask, cast as authoring_cast, float4_e2m1fn, float8_e4m3fn, @@ -146,7 +145,6 @@ def _import_coreai() -> tuple: float16, float32, lut_to_dense, - sparse_with_bitmask_to_dense, ) from coreai.authoring.constant import constant as authoring_constant # noqa: PLC0415 # isort: on @@ -194,7 +192,6 @@ def _import_coreai() -> tuple: authoring_cast, authoring_constant, blockwise_shift_scale, - build_sparse_with_bitmask, float4_e2m1fn, float8_e4m3fn, float8_e5m2, @@ -202,7 +199,6 @@ def _import_coreai() -> tuple: float16, float32, lut_to_dense, - sparse_with_bitmask_to_dense, TensorType, DebugInfo, ) @@ -230,7 +226,6 @@ def _import_coreai() -> tuple: authoring_cast, authoring_constant, blockwise_shift_scale, - build_sparse_with_bitmask, float4_e2m1fn, float8_e4m3fn, float8_e5m2, @@ -238,7 +233,6 @@ def _import_coreai() -> tuple: float16, float32, lut_to_dense, - sparse_with_bitmask_to_dense, TensorType, DebugInfo, ) = lazy_import_coreai_torch(_import_coreai) diff --git a/src/coreai_opt/coreai_utils/passes/weight_sparsification.py b/src/coreai_opt/coreai_utils/passes/weight_sparsification.py index 0246cb45..9aa100f0 100644 --- a/src/coreai_opt/coreai_utils/passes/weight_sparsification.py +++ b/src/coreai_opt/coreai_utils/passes/weight_sparsification.py @@ -23,13 +23,11 @@ _get_constant_value_as_np_array, authoring_constant as _authoring_constant, blockwise_shift_scale as _blockwise_shift_scale, - build_sparse_with_bitmask as _build_sparse_with_bitmask, compression_types as _compression_types, coreai as _coreai, float8_e4m3fn as _float8_e4m3fn, float8_e5m2 as _float8_e5m2, lut_to_dense as _lut_to_dense, - sparse_with_bitmask_to_dense as _sparse_with_bitmask_to_dense, ) from coreai_opt.coreai_utils._utils.graph_utils import ( _apply_compression_transform, @@ -383,11 +381,8 @@ def replace_weight_with_compression_op(op: Any) -> Any: sparse_params.mask, _IntegerType.get_unsigned(1), ) - sparse_tensor = _build_sparse_with_bitmask( - values=nonzero_data_const, - bitmask=mask, - )._to_mlir() - compressed_weight = _sparse_with_bitmask_to_dense(sparse_tensor)._to_mlir() + sparse_tensor = _coreai.build_sparse_with_bitmask(nonzero_data_const, mask) + compressed_weight = _coreai.sparse_with_bitmask_to_dense(sparse_tensor) if require_dequantization: compressed_weight = _blockwise_shift_scale( From f4aceddcc973aa4a635487063d37817cfd1bded2 Mon Sep 17 00:00:00 2001 From: Guru Deshpande <66385690+guru-desh@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:01:12 -0700 Subject: [PATCH 2/3] refactor(coreai_utils): import the raw sparse ops by name --- src/coreai_opt/coreai_utils/_coreai_imports.py | 11 +++++++++++ .../coreai_utils/passes/weight_sparsification.py | 9 +++++++-- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/coreai_opt/coreai_utils/_coreai_imports.py b/src/coreai_opt/coreai_utils/_coreai_imports.py index dcb6a4a9..4effa972 100644 --- a/src/coreai_opt/coreai_utils/_coreai_imports.py +++ b/src/coreai_opt/coreai_utils/_coreai_imports.py @@ -150,6 +150,13 @@ def _import_coreai() -> tuple: # isort: on from coreai._compiler.dialects import coreai as coreai_dialect # noqa: PLC0415 + + # Raw dialect functions: unlike the coreai.authoring wrappers of the same name, they take and + # return raw MLIR values. + from coreai._compiler.dialects.coreai import ( # noqa: PLC0415 + build_sparse_with_bitmask, + sparse_with_bitmask_to_dense, + ) from coreai._compiler.dialects.coreai.constant import create_elements_attr # noqa: PLC0415 from coreai._compiler.ir import ( # noqa: PLC0415 BoolAttr, @@ -192,6 +199,7 @@ def _import_coreai() -> tuple: authoring_cast, authoring_constant, blockwise_shift_scale, + build_sparse_with_bitmask, float4_e2m1fn, float8_e4m3fn, float8_e5m2, @@ -199,6 +207,7 @@ def _import_coreai() -> tuple: float16, float32, lut_to_dense, + sparse_with_bitmask_to_dense, TensorType, DebugInfo, ) @@ -226,6 +235,7 @@ def _import_coreai() -> tuple: authoring_cast, authoring_constant, blockwise_shift_scale, + build_sparse_with_bitmask, float4_e2m1fn, float8_e4m3fn, float8_e5m2, @@ -233,6 +243,7 @@ def _import_coreai() -> tuple: float16, float32, lut_to_dense, + sparse_with_bitmask_to_dense, TensorType, DebugInfo, ) = lazy_import_coreai_torch(_import_coreai) diff --git a/src/coreai_opt/coreai_utils/passes/weight_sparsification.py b/src/coreai_opt/coreai_utils/passes/weight_sparsification.py index 9aa100f0..628d2734 100644 --- a/src/coreai_opt/coreai_utils/passes/weight_sparsification.py +++ b/src/coreai_opt/coreai_utils/passes/weight_sparsification.py @@ -23,11 +23,13 @@ _get_constant_value_as_np_array, authoring_constant as _authoring_constant, blockwise_shift_scale as _blockwise_shift_scale, + build_sparse_with_bitmask as _build_sparse_with_bitmask, compression_types as _compression_types, coreai as _coreai, float8_e4m3fn as _float8_e4m3fn, float8_e5m2 as _float8_e5m2, lut_to_dense as _lut_to_dense, + sparse_with_bitmask_to_dense as _sparse_with_bitmask_to_dense, ) from coreai_opt.coreai_utils._utils.graph_utils import ( _apply_compression_transform, @@ -381,8 +383,11 @@ def replace_weight_with_compression_op(op: Any) -> Any: sparse_params.mask, _IntegerType.get_unsigned(1), ) - sparse_tensor = _coreai.build_sparse_with_bitmask(nonzero_data_const, mask) - compressed_weight = _coreai.sparse_with_bitmask_to_dense(sparse_tensor) + sparse_tensor = _build_sparse_with_bitmask( + values=nonzero_data_const, + bitmask=mask, + ) + compressed_weight = _sparse_with_bitmask_to_dense(sparse_tensor) if require_dequantization: compressed_weight = _blockwise_shift_scale( From ace2ffb2ccb8151b958f6fd8fd6467eef3ab6ea9 Mon Sep 17 00:00:00 2001 From: Guru Deshpande <66385690+guru-desh@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:33:33 -0700 Subject: [PATCH 3/3] refactor(coreai_utils): drop the comment on the raw sparse imports --- src/coreai_opt/coreai_utils/_coreai_imports.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/coreai_opt/coreai_utils/_coreai_imports.py b/src/coreai_opt/coreai_utils/_coreai_imports.py index 4effa972..c19c78f4 100644 --- a/src/coreai_opt/coreai_utils/_coreai_imports.py +++ b/src/coreai_opt/coreai_utils/_coreai_imports.py @@ -150,9 +150,6 @@ def _import_coreai() -> tuple: # isort: on from coreai._compiler.dialects import coreai as coreai_dialect # noqa: PLC0415 - - # Raw dialect functions: unlike the coreai.authoring wrappers of the same name, they take and - # return raw MLIR values. from coreai._compiler.dialects.coreai import ( # noqa: PLC0415 build_sparse_with_bitmask, sparse_with_bitmask_to_dense,