diff --git a/aie_kernels/generic/dequant_i32_bf16.cc b/aie_kernels/generic/dequant_i32_bf16.cc new file mode 100644 index 0000000000..7b630cea82 --- /dev/null +++ b/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 +#include + +template 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 sf_f32 = aie::broadcast((float)sf); + + for (int k = 0; k < blocks_per_group; k++) { + aie::vector i32_vec = aie::load_v(pI); + pI += block_size; + + // Convert int32 -> float, scale in float32, then convert to bfloat16 + aie::vector f32_vec = aie::to_float(i32_vec, 0); + aie::vector scaled_f32 = aie::mul(f32_vec, sf_f32); + aie::accum acc; + acc.from_vector(scaled_f32, 0); + aie::store_v(pO, acc.to_vector()); + 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(a_in, c_out); +} + +} // extern "C" diff --git a/iron/operators/dequant_i32_bf16/design.py b/iron/operators/dequant_i32_bf16/design.py new file mode 100644 index 0000000000..8a1f7ed429 --- /dev/null +++ b/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()) diff --git a/iron/operators/dequant_i32_bf16/op.py b/iron/operators/dequant_i32_bf16/op.py new file mode 100644 index 0000000000..ca268a3bfe --- /dev/null +++ b/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), + ] diff --git a/iron/operators/dequant_i32_bf16/reference.py b/iron/operators/dequant_i32_bf16/reference.py new file mode 100644 index 0000000000..8c07f78cb3 --- /dev/null +++ b/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, + } diff --git a/iron/operators/dequant_i32_bf16/test.py b/iron/operators/dequant_i32_bf16/test.py new file mode 100644 index 0000000000..5e1f7ec6dc --- /dev/null +++ b/iron/operators/dequant_i32_bf16/test.py @@ -0,0 +1,92 @@ +#!/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.dequant_i32_bf16.op import DequantI32 +from iron.operators.dequant_i32_bf16.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 8192 + if tile_size > 8192: + tile_size = 8192 + + # 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_dequant_i32( + 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 = DequantI32( + 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()} + + errors, latency_us, bandwidth_gbps = run_test( + operator, input_buffers, output_buffers, rel_tol=0.01, abs_tol=1e-6 + ) + + 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}"