Add dequant operator (int32 → bfloat16)

aie_kernels/generic/dequant_i32_bf16.cc

@@ -0,0 +1,59 @@
+// 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 dequant_i32_to_bf16_impl(int32_t *in, bfloat16 *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");
+
+    int32_t *__restrict pI = in;
+    bfloat16 *__restrict pSF = (bfloat16 *)(in + N); // Scale factors after int32 data
+    bfloat16 *__restrict pO = out;
+
+    event0();
+    for (int i = 0; i < groups_per_tile; i++)
+        chess_prepare_for_pipelining chess_loop_range(groups_per_tile, )
+        {
+            bfloat16 sf = *pSF;
+            pSF += 1;
+            aie::vector<float, block_size> sf_f32 = aie::broadcast<float, block_size>((float)sf);
+
+            for (int k = 0; k < blocks_per_group; k++) {
+                aie::vector<int32_t, block_size> i32_vec = aie::load_v<block_size>(pI);
+                pI += block_size;
+
+                // Convert int32 -> float, scale in float32, then convert to bfloat16
+                aie::vector<float, block_size> f32_vec = aie::to_float<float>(i32_vec, 0);
+                aie::vector<float, block_size> scaled_f32 = aie::mul(f32_vec, sf_f32);
+                aie::accum<accfloat, block_size> acc;
+                acc.from_vector(scaled_f32, 0);
+                aie::store_v(pO, acc.to_vector<bfloat16>());
+                pO += block_size;
+            }
+        }
+    event1();
+}
+
+extern "C" {
+
+#ifndef GROUP_SIZE
+#define GROUP_SIZE 32
+#endif
+
+#ifndef TILE_SIZE
+#define TILE_SIZE 1024
+#endif
+
+void dequant_i32_to_bfloat16(int32_t *a_in, bfloat16 *c_out)
+{
+    dequant_i32_to_bf16_impl<TILE_SIZE, GROUP_SIZE>(a_in, c_out);
+}
+
+} // extern "C"

iron/operators/dequant_i32_bf16/design.py

@@ -0,0 +1,144 @@
+# 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_dequant_i32_kernel(
+    dev,
+    num_elements,
+    num_columns,
+    num_channels,
+    trace_size,
+    tile_size,
+    group_size,
+):
+    # Cap tile size to fit in L1 memory
+    # int32 input: tile_size*4 bytes + scale factors: (tile_size/group_size)*2 bytes
+    # bf16 output: tile_size*2 bytes
+    # Total per tile ~ tile_size*6 bytes; 8192 elements ~ 49KB fits in 64KB L1
+    per_tile_elements = 8192 if tile_size > 8192 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 = np.uint8
+    out_dtype = bfloat16
+
+    # Input: int32 data (4 bytes each) + bf16 scale factors (2 bytes each, one per group)
+    input_tensor_size = num_elements * 4 + (num_elements // group_size) * 2
+    input_tile_size = per_tile_elements * 4 + (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[(num_elements,), np.dtype[out_dtype]]
+    in_tile_ty = np.ndarray[(input_tile_size,), np.dtype[in_dtype]]
+    out_tile_ty = np.ndarray[(per_tile_elements,), np.dtype[out_dtype]]
+
+    fifodepth = 1 if tile_size > 4096 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
+    dequant_kernel = Kernel(
+        "dequant_i32_to_bfloat16",
+        f"dequant_i32_bf16_{tile_size}.o",
+        [in_tile_ty, out_tile_ty],
+    )
+
+    # Core task
+    def core_body(of_in, of_out, dequant_kernel):
+        for _ in range_(N_div_n):
+            elem_in = of_in.acquire(1)
+            elem_out = of_out.acquire(1)
+            dequant_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(),
+                dequant_kernel,
+            ],
+        )
+        for i in range(num_columns)
+        for j in range(num_channels)
+    ]
+
+    # TensorAccessPatterns
+    in_chunk = chunk * 4 + (chunk // group_size) * 2  # bytes per core in input
+    taps_in = [
+        TensorAccessPattern(
+            (1, input_tensor_size),
+            in_chunk * (i * num_channels + j),
+            [1, 1, 1, in_chunk],
+            [0, 0, 0, 1],
+        )
+        for i in range(num_columns)
+        for j in range(num_channels)
+    ]
+    taps_out = [
+        TensorAccessPattern(
+            (1, num_elements),
+            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)
+    ]
+
+    # 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/dequant_i32_bf16/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 DequantI32(MLIROperator):
+    """AIE-accelerated dequantization operator (int32 -> bfloat16)"""
+
+    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):
+        # Input: int32 data (4 bytes each) + bf16 scale factors (2 bytes each, one per group)
+        self.input_size = self.size * 4 + (self.size // self.group_size) * 2
+        self.output_size = self.size
+
+        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_dequant_i32_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"dequant_i32_bf16_{self.tile_size}.o",
+                dependencies=[
+                    SourceArtifact(
+                        self.context.base_dir
+                        / "aie_kernels"
+                        / "generic"
+                        / "dequant_i32_bf16.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=np.uint8),
+            AIERuntimeArgSpec("out", (self.output_size,), dtype=bfloat16),
+        ]

iron/operators/dequant_i32_bf16/reference.py

@@ -0,0 +1,70 @@
+# 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 i32 -> bf16 dequantization.
+
+    Simulates the output of an INT8 GEMM: generates random int32 values
+    in a realistic accumulator range, with per-group bf16 scale factors.
+    The reference output is: bf16(int32_val * scale).
+    """
+    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 int32 values in realistic INT8 GEMM accumulator range
+    i32_values = torch.randint(-100000, 100001, (input_length,), dtype=torch.int32)
+
+    # Generate per-group scale factors in a realistic range
+    # Combined scale = scale_activations * scale_weights
+    scales = (0.0001 + 0.01 * torch.rand(num_groups, dtype=torch.float32)).to(
+        torch.bfloat16
+    )
+
+    # Compute reference output: float(i32) * scale -> bf16
+    output = torch.zeros(input_length, dtype=torch.bfloat16)
+    for g in range(num_groups):
+        start = g * group_size
+        end = start + group_size
+        output[start:end] = (
+            i32_values[start:end].to(torch.float32) * scales[g].to(torch.float32)
+        ).to(torch.bfloat16)
+
+    # Pack input buffer as uint8: per-tile layout of int32 data followed by scale factors
+    num_tiles = input_length // tile_size
+    scales_per_tile = tile_size // group_size
+    tile_input_bytes = tile_size * 4 + scales_per_tile * 2
+
+    packed_input = np.zeros(num_tiles * tile_input_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_input_bytes
+
+        # Pack int32 data as bytes
+        i32_bytes = i32_values[tile_start:tile_end].numpy().view(np.uint8)
+        packed_input[byte_offset : byte_offset + tile_size * 4] = i32_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_input[byte_offset + tile_size * 4 : byte_offset + tile_input_bytes] = (
+            scale_bytes
+        )
+
+    return {
+        "input": torch.from_numpy(packed_input),
+        "output": output,
+    }