Add quantize operator (bfloat16 → int8)

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 <aie_api/aie.hpp>
+#include <stdint.h>
+
+template <const int N, const int G> 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<bfloat16, block_size> bf16_vec = aie::load_v<block_size>(pScan);
+                pScan += block_size;
+                aie::vector<bfloat16, block_size> 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<float, block_size> inv_scale_vec = aie::broadcast<float, block_size>(inv_scale);
+            for (int k = 0; k < blocks_per_group; k++) {
+                aie::vector<bfloat16, block_size> bf16_vec = aie::load_v<block_size>(pI);
+                pI += block_size;
+
+                // Convert bf16 -> float, scale, round, saturate to int8
+                aie::accum<accfloat, block_size> acc;
+                acc.from_vector(bf16_vec, 0);
+                aie::vector<float, block_size> f32_vec = acc.to_vector<float>();
+                aie::vector<float, block_size> 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<TILE_SIZE, GROUP_SIZE>(a_in, c_out);
+}
+
+} // extern "C"

iron/operators/quantize_bf16_i8/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())

iron/operators/quantize_bf16_i8/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),
+        ]