From abf0178c787d0fd7e986a7771ebc93c3f56709ee Mon Sep 17 00:00:00 2001 From: albiol2004 Date: Thu, 9 Apr 2026 14:22:59 +0200 Subject: [PATCH 1/2] =?UTF-8?q?Add=20quantize=20operator=20(bfloat16=20?= =?UTF-8?q?=E2=86=92=20int8)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a per-group symmetric quantization operator that converts bfloat16 tensors to int8 with per-group scale factors. Completes the W8A8 quantized inference pipeline alongside INT8 GEMM (#93) and dequant i32→bf16 (#95). - C++ kernel: two-pass (find max abs, then scale+round+saturate) - Python operator following existing dequant/ pattern - 32 parametrized tests (1-8 columns, 1-2 channels, 1K-8K elements) --- aie_kernels/generic/quantize_bf16_i8.cc | 97 ++++++++++++++++ iron/operators/quantize/design.py | 148 ++++++++++++++++++++++++ iron/operators/quantize/op.py | 96 +++++++++++++++ iron/operators/quantize/reference.py | 78 +++++++++++++ iron/operators/quantize/test.py | 94 +++++++++++++++ 5 files changed, 513 insertions(+) create mode 100644 aie_kernels/generic/quantize_bf16_i8.cc create mode 100644 iron/operators/quantize/design.py create mode 100644 iron/operators/quantize/op.py create mode 100644 iron/operators/quantize/reference.py create mode 100644 iron/operators/quantize/test.py diff --git a/aie_kernels/generic/quantize_bf16_i8.cc b/aie_kernels/generic/quantize_bf16_i8.cc new file mode 100644 index 0000000000..10d6fb7953 --- /dev/null +++ b/aie_kernels/generic/quantize_bf16_i8.cc @@ -0,0 +1,97 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "../aie_kernel_utils.h" + +#include +#include + +template void quantize_bf16_to_i8_impl(bfloat16 *in, int8_t *out) +{ + constexpr int block_size = 16; + constexpr int blocks_per_group = G / block_size; + constexpr int groups_per_tile = N / G; + static_assert((G % block_size) == 0, "GROUP_SIZE must be a multiple of 16"); + + bfloat16 *__restrict pI = in; + int8_t *__restrict pO = out; + // Scale factors stored after int8 data. N is always a multiple of group_size (>= 16), + // so out + N is guaranteed to be 2-byte aligned for bfloat16 access. + bfloat16 *__restrict pSF = (bfloat16 *)(out + N); + + event0(); + for (int i = 0; i < groups_per_tile; i++) + chess_prepare_for_pipelining chess_loop_range(groups_per_tile, ) + { + // Pass 1: Find max absolute value in this group + float max_abs = 0.0f; + bfloat16 *pScan = pI; + for (int k = 0; k < blocks_per_group; k++) { + aie::vector bf16_vec = aie::load_v(pScan); + pScan += block_size; + aie::vector abs_vec = aie::abs(bf16_vec); + float block_max = (float)aie::reduce_max(abs_vec); + if (block_max > max_abs) + max_abs = block_max; + } + + // Compute scale factor and inverse scale + float inv_scale; + if (max_abs > 0.0f) { + inv_scale = ::aie::div(127.0f, max_abs); + *pSF = (bfloat16)::aie::div(max_abs, 127.0f); + } else { + inv_scale = 0.0f; + *pSF = (bfloat16)0.0f; + } + pSF += 1; + + // Pass 2: Quantize values + aie::vector inv_scale_vec = aie::broadcast(inv_scale); + for (int k = 0; k < blocks_per_group; k++) { + aie::vector bf16_vec = aie::load_v(pI); + pI += block_size; + + // Convert bf16 -> float, scale, round, saturate to int8 + aie::accum acc; + acc.from_vector(bf16_vec, 0); + aie::vector f32_vec = acc.to_vector(); + aie::vector scaled = aie::mul(f32_vec, inv_scale_vec); + + // Round to nearest (half-away-from-zero) and saturate to [-128, 127] + // TODO: vectorize with aie::to_fixed when available for int8 targets + for (int e = 0; e < block_size; e++) { + float val = scaled[e]; + int32_t rounded; + if (val >= 0.0f) + rounded = (int32_t)(val + 0.5f); + else + rounded = (int32_t)(val - 0.5f); + if (rounded > 127) + rounded = 127; + else if (rounded < -128) + rounded = -128; + pO[e] = (int8_t)rounded; + } + pO += block_size; + } + } + event1(); +} + +extern "C" { + +#ifndef GROUP_SIZE +#define GROUP_SIZE 32 +#endif + +#ifndef TILE_SIZE +#define TILE_SIZE 1024 +#endif + +void quantize_bfloat16_to_i8(bfloat16 *a_in, int8_t *c_out) +{ + quantize_bf16_to_i8_impl(a_in, c_out); +} + +} // extern "C" diff --git a/iron/operators/quantize/design.py b/iron/operators/quantize/design.py new file mode 100644 index 0000000000..840ce7ca31 --- /dev/null +++ b/iron/operators/quantize/design.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from ml_dtypes import bfloat16 +import numpy as np + +from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker +from aie.iron.placers import SequentialPlacer +from aie.helpers.taplib.tap import TensorAccessPattern +from aie.iron.controlflow import range_ + + +def my_quantize_kernel( + dev, + num_elements, + num_columns, + num_channels, + trace_size, + tile_size, + group_size, +): + # Cap tile size to fit in L1 memory + # bf16 input: tile_size*2 bytes + # int8 output + scales: tile_size + (tile_size/group_size)*2 bytes + # Total ~ tile_size*3 bytes; 16384 elements ~ 49KB fits in 64KB L1 + per_tile_elements = 16384 if tile_size > 16384 else tile_size + total_cores = num_columns * num_channels + per_core_elements = num_elements // total_cores + if num_elements % total_cores != 0: + raise ValueError( + f"Number of elements ({num_elements}) must be a multiple of {total_cores}." + ) + N_div_n = per_core_elements // per_tile_elements + chunk = num_elements // num_columns // num_channels + + in_dtype = bfloat16 + out_dtype = np.uint8 + + # Input: bf16 values + input_tensor_size = num_elements + input_tile_size = per_tile_elements + + # Output: int8 data + bf16 scale factors (packed as uint8) + output_tensor_size = num_elements + (num_elements // group_size) * 2 + output_tile_size = per_tile_elements + (per_tile_elements // group_size) * 2 + + # Define tensor types + in_tensor_ty = np.ndarray[(input_tensor_size,), np.dtype[in_dtype]] + out_tensor_ty = np.ndarray[(output_tensor_size,), np.dtype[out_dtype]] + in_tile_ty = np.ndarray[(input_tile_size,), np.dtype[in_dtype]] + out_tile_ty = np.ndarray[(output_tile_size,), np.dtype[out_dtype]] + + fifodepth = 1 if tile_size > 8192 else 2 + enable_trace = trace_size > 0 + + # ObjectFIFOs for data movement + of_ins = [ + ObjectFifo(in_tile_ty, name=f"in_{i}_{j}", depth=fifodepth) + for i in range(num_columns) + for j in range(num_channels) + ] + of_outs = [ + ObjectFifo(out_tile_ty, name=f"out_{i}_{j}", depth=fifodepth) + for i in range(num_columns) + for j in range(num_channels) + ] + + # Kernel declaration + quantize_kernel = Kernel( + "quantize_bfloat16_to_i8", + f"quantize_bf16_i8_{tile_size}.o", + [in_tile_ty, out_tile_ty], + ) + + # Core task + def core_body(of_in, of_out, quantize_kernel): + for _ in range_(N_div_n): + elem_in = of_in.acquire(1) + elem_out = of_out.acquire(1) + quantize_kernel(elem_in, elem_out) + of_in.release(1) + of_out.release(1) + + # Workers + my_workers = [ + Worker( + core_body, + [ + of_ins[i * num_channels + j].cons(), + of_outs[i * num_channels + j].prod(), + quantize_kernel, + ], + ) + for i in range(num_columns) + for j in range(num_channels) + ] + + # TensorAccessPatterns + out_chunk = chunk + (chunk // group_size) * 2 # bytes per core in output + taps_in = [ + TensorAccessPattern( + (1, input_tensor_size), + chunk * (i * num_channels + j), + [1, 1, 1, chunk], + [0, 0, 0, 1], + ) + for i in range(num_columns) + for j in range(num_channels) + ] + taps_out = [ + TensorAccessPattern( + (1, output_tensor_size), + out_chunk * (i * num_channels + j), + [1, 1, 1, out_chunk], + [0, 0, 0, 1], + ) + for i in range(num_columns) + for j in range(num_channels) + ] + + # Runtime + rt = Runtime() + with rt.sequence(in_tensor_ty, out_tensor_ty) as (A, C): + if enable_trace: + rt.enable_trace(trace_size) + rt.start(*my_workers) + + tg = rt.task_group() + for i in range(num_columns): + for j in range(num_channels): + rt.fill( + of_ins[i * num_channels + j].prod(), + A, + taps_in[i * num_channels + j], + task_group=tg, + ) + for i in range(num_columns): + for j in range(num_channels): + rt.drain( + of_outs[i * num_channels + j].cons(), + C, + taps_out[i * num_channels + j], + wait=True, + task_group=tg, + ) + rt.finish_task_group(tg) + + return Program(dev, rt).resolve_program(SequentialPlacer()) diff --git a/iron/operators/quantize/op.py b/iron/operators/quantize/op.py new file mode 100644 index 0000000000..245e6ef8a4 --- /dev/null +++ b/iron/operators/quantize/op.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +import numpy as np +from ml_dtypes import bfloat16 + +from iron.common import ( + MLIROperator, + AIERuntimeArgSpec, + KernelObjectArtifact, + SourceArtifact, + PythonGeneratedMLIRArtifact, + DesignGenerator, +) +import aie.utils as aie_utils + + +@dataclass +class Quantize(MLIROperator): + """AIE-accelerated quantization operator (bfloat16 -> int8)""" + + size: int + num_aie_columns: int + num_channels: int + tile_size: int + group_size: int = field(default=32, repr=False) + context: object = field(default=None, repr=False) + + def __post_init__(self): + # Output: int8 data (1 byte each) + bf16 scale factors (2 bytes each, one per group) + self.input_size = self.size + self.output_size = self.size + (self.size // self.group_size) * 2 + + total_cores = self.num_aie_columns * self.num_channels + if self.size % self.tile_size != 0: + raise ValueError( + f"size ({self.size}) must be divisible by tile_size ({self.tile_size})" + ) + if self.size % total_cores != 0: + raise ValueError( + f"size ({self.size}) must be divisible by total cores ({total_cores})" + ) + if total_cores > 16: + raise ValueError(f"total cores ({total_cores}) must be <= 16") + if self.tile_size % self.group_size != 0: + raise ValueError( + f"tile_size ({self.tile_size}) must be divisible by group_size ({self.group_size})" + ) + if self.group_size % 16 != 0: + raise ValueError(f"group_size ({self.group_size}) must be a multiple of 16") + MLIROperator.__init__(self, context=self.context) + + def get_mlir_artifact(self): + return PythonGeneratedMLIRArtifact( + f"{self.name}.mlir", + DesignGenerator( + self.operator_dir / "design.py", + "my_quantize_kernel", + ( + aie_utils.get_current_device(), + self.size, + self.num_aie_columns, + self.num_channels, + 0, # trace_size + self.tile_size, + self.group_size, + ), + ), + ) + + def get_kernel_artifacts(self): + return [ + KernelObjectArtifact( + f"quantize_bf16_i8_{self.tile_size}.o", + dependencies=[ + SourceArtifact( + self.context.base_dir + / "aie_kernels" + / "generic" + / "quantize_bf16_i8.cc" + ) + ], + extra_flags=[ + f"-DTILE_SIZE={self.tile_size}", + f"-DGROUP_SIZE={self.group_size}", + ], + ) + ] + + def get_arg_spec(self): + return [ + AIERuntimeArgSpec("in", (self.input_size,), dtype=bfloat16), + AIERuntimeArgSpec("out", (self.output_size,), dtype=np.uint8), + ] diff --git a/iron/operators/quantize/reference.py b/iron/operators/quantize/reference.py new file mode 100644 index 0000000000..e76fb4e3de --- /dev/null +++ b/iron/operators/quantize/reference.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +import numpy as np + + +def generate_golden_reference(input_length, tile_size, group_size): + """Generate test data and reference output for bf16 -> int8 quantization. + + Generates random bfloat16 values, computes per-group symmetric quantization + to int8, and packs the output per-tile as [int8_data | bf16_scale_factors]. + """ + torch.manual_seed(42) + + if input_length % tile_size != 0: + raise ValueError("Input length must be a multiple of tile size.") + if tile_size % group_size != 0: + raise ValueError("Tile size must be a multiple of group size.") + + num_groups = input_length // group_size + + # Generate random bf16 input (range similar to model activations) + bf16_input = torch.randn(input_length, dtype=torch.bfloat16) + + # Per-group symmetric quantization + i8_output = torch.zeros(input_length, dtype=torch.int8) + scales = torch.zeros(num_groups, dtype=torch.bfloat16) + + for g in range(num_groups): + start = g * group_size + end = start + group_size + group = bf16_input[start:end] + + max_abs = group.float().abs().max() + if max_abs > 0: + scale = (max_abs / 127.0).to(torch.bfloat16) + inv_scale = 127.0 / max_abs.float() + # Use round-half-away-from-zero to match AIE kernel rounding + scaled = group.float().mul(inv_scale) + quantized = torch.sign(scaled) * torch.floor(torch.abs(scaled) + 0.5) + quantized = quantized.clamp(-128, 127) + else: + scale = torch.tensor(0.0, dtype=torch.bfloat16) + quantized = torch.zeros(group_size, dtype=torch.float32) + + i8_output[start:end] = quantized.to(torch.int8) + scales[g] = scale + + # Pack output per-tile: [i8_data | bf16_scale_factors] as uint8 + num_tiles = input_length // tile_size + scales_per_tile = tile_size // group_size + tile_output_bytes = tile_size + scales_per_tile * 2 + + packed_output = np.zeros(num_tiles * tile_output_bytes, dtype=np.uint8) + for t in range(num_tiles): + tile_start = t * tile_size + tile_end = tile_start + tile_size + byte_offset = t * tile_output_bytes + + # Pack int8 data as bytes + i8_bytes = i8_output[tile_start:tile_end].numpy().view(np.uint8) + packed_output[byte_offset : byte_offset + tile_size] = i8_bytes + + # Pack bf16 scale factors as bytes + scale_start = t * scales_per_tile + scale_end = scale_start + scales_per_tile + scale_bytes = ( + scales[scale_start:scale_end].view(torch.uint16).numpy().view(np.uint8) + ) + packed_output[byte_offset + tile_size : byte_offset + tile_output_bytes] = ( + scale_bytes + ) + + return { + "input": bf16_input, + "output": torch.from_numpy(packed_output), + } diff --git a/iron/operators/quantize/test.py b/iron/operators/quantize/test.py new file mode 100644 index 0000000000..28c549a259 --- /dev/null +++ b/iron/operators/quantize/test.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import aie.utils as aie_utils + +from iron.operators.quantize.op import Quantize +from iron.operators.quantize.reference import generate_golden_reference +from iron.common.test_utils import run_test + + +def get_params(): + max_aie_columns = aie_utils.get_current_device().cols + + input_lengths = [1024, 2048, 4096, 8192] + group_size = 32 + + params = [] + for input_length in input_lengths: + for num_columns in range(1, max_aie_columns + 1): + for num_channels in range(1, 3): # 1 or 2 channels + total_cores = num_columns * num_channels + tile_size = input_length // total_cores + + # Cap tile_size at 16384 + if tile_size > 16384: + tile_size = 16384 + + # Only proceed if exact division + if tile_size * total_cores != input_length: + continue + + # Ensure tile_size is divisible by group_size + if tile_size % group_size != 0: + continue + + is_regular = input_length == 2048 + marks = [] if is_regular else [pytest.mark.extensive] + + params.append( + pytest.param( + input_length, + num_columns, + num_channels, + tile_size, + group_size, + marks=marks, + ) + ) + return params + + +@pytest.mark.metrics( + Latency=r"Latency \(us\): (?P[\d\.]+)", + Bandwidth=r"Effective Bandwidth: (?P[\d\.e\+-]+) GB/s", +) +@pytest.mark.parametrize( + "input_length,num_aie_columns,num_channels,tile_size,group_size", + get_params(), +) +def test_quantize( + input_length, num_aie_columns, num_channels, tile_size, group_size, aie_context +): + golden_ref = generate_golden_reference( + input_length=input_length, + tile_size=tile_size, + group_size=group_size, + ) + + operator = Quantize( + size=input_length, + num_aie_columns=num_aie_columns, + num_channels=num_channels, + tile_size=tile_size, + group_size=group_size, + context=aie_context, + ) + + input_buffers = { + "input": golden_ref["input"].flatten(), + } + output_buffers = {"output": golden_ref["output"].flatten()} + + # Allow ±1 tolerance for int8 rounding differences between AIE and x86 + # abs_tol=1.01 because verify_buffer uses >= (so 1.0 fails on exact ±1 diff) + errors, latency_us, bandwidth_gbps = run_test( + operator, input_buffers, output_buffers, rel_tol=0.01, abs_tol=1.01 + ) + + print(f"\nLatency (us): {latency_us:.1f}") + print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") + + assert not errors, f"Test failed with errors: {errors}" From b80e4ed344c0965ed0c8cdd07e4297186764cfd8 Mon Sep 17 00:00:00 2001 From: albiol2004 Date: Thu, 9 Apr 2026 15:54:33 +0200 Subject: [PATCH 2/2] Rename quantize operator to quantize_bf16_i8 for naming consistency Align operator directory name with the kernel file convention (quantize_bf16_i8.cc). This establishes a consistent full input_output type naming pattern for quantization operators. --- iron/operators/{quantize => quantize_bf16_i8}/design.py | 0 iron/operators/{quantize => quantize_bf16_i8}/op.py | 0 iron/operators/{quantize => quantize_bf16_i8}/reference.py | 0 iron/operators/{quantize => quantize_bf16_i8}/test.py | 4 ++-- 4 files changed, 2 insertions(+), 2 deletions(-) rename iron/operators/{quantize => quantize_bf16_i8}/design.py (100%) rename iron/operators/{quantize => quantize_bf16_i8}/op.py (100%) rename iron/operators/{quantize => quantize_bf16_i8}/reference.py (100%) rename iron/operators/{quantize => quantize_bf16_i8}/test.py (95%) diff --git a/iron/operators/quantize/design.py b/iron/operators/quantize_bf16_i8/design.py similarity index 100% rename from iron/operators/quantize/design.py rename to iron/operators/quantize_bf16_i8/design.py diff --git a/iron/operators/quantize/op.py b/iron/operators/quantize_bf16_i8/op.py similarity index 100% rename from iron/operators/quantize/op.py rename to iron/operators/quantize_bf16_i8/op.py diff --git a/iron/operators/quantize/reference.py b/iron/operators/quantize_bf16_i8/reference.py similarity index 100% rename from iron/operators/quantize/reference.py rename to iron/operators/quantize_bf16_i8/reference.py diff --git a/iron/operators/quantize/test.py b/iron/operators/quantize_bf16_i8/test.py similarity index 95% rename from iron/operators/quantize/test.py rename to iron/operators/quantize_bf16_i8/test.py index 28c549a259..b9efd46a87 100644 --- a/iron/operators/quantize/test.py +++ b/iron/operators/quantize_bf16_i8/test.py @@ -5,8 +5,8 @@ import pytest import aie.utils as aie_utils -from iron.operators.quantize.op import Quantize -from iron.operators.quantize.reference import generate_golden_reference +from iron.operators.quantize_bf16_i8.op import Quantize +from iron.operators.quantize_bf16_i8.reference import generate_golden_reference from iron.common.test_utils import run_test