[Mlir-commits] [mlir] [mlir][arith] Add LUT-based lowering of arith.extf for FP8 to F32. (PR #196321)
Javed Absar
llvmlistbot at llvm.org
Sun May 10 08:17:38 PDT 2026
https://github.com/javedabsar1 updated https://github.com/llvm/llvm-project/pull/196321
>From ad1a316142c48c0320f7f1c0144eb16e6cfcb43f Mon Sep 17 00:00:00 2001
From: Javed Absar <javed.absar at gmail.com>
Date: Sat, 25 Apr 2026 13:36:17 -0400
Subject: [PATCH 1/3] [mlir][arith] Add LUT-based lowering of arith.extf for
FP8 to F32.
FP8 formats are for efficient storage and for the actual computation
the data format is often converted to FP32. MLIR's existing arith.extf
lowers FP8 to a sequence of arithmetic operations.
This PR adds an alternative pass, --convert-arith-fp8-extf-to-lut,
that replaces arith.extf %v : f8X to f32 with a table lookup
into a 256-entry memref<256xf32> global constant.
One table is emitted per distinct FP8 format; values are precomputed at
compile time via APFloat to match exact IEEE semantics.
Supported formats:
f8E4M3FN, f8E5M2, f8E4M3FNUZ, f8E5M2FNUZ, f8E4M3B11FNUZ, f8E3M4, f8E4M3.
Includes FileCheck tests, an integration test (JIT), and unit test that
verifies all 256 bit patterns for each format against APFloat reference values.
Signed-off-by: Javed Absar <javed.absar at gmail.com>
---
.../mlir/Conversion/ArithToLUT/ArithToLUT.h | 23 ++
mlir/include/mlir/Conversion/Passes.h | 1 +
mlir/include/mlir/Conversion/Passes.td | 26 +++
mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp | 197 ++++++++++++++++++
mlir/lib/Conversion/ArithToLUT/CMakeLists.txt | 18 ++
mlir/lib/Conversion/CMakeLists.txt | 1 +
.../Conversion/ArithToLUT/extf-to-lut.mlir | 36 ++++
.../ArithToLUT/linalg-generic-fp8-addf.mlir | 33 +++
.../Dialect/Arith/CPU/test-lut-extf-f8.mlir | 59 ++++++
.../Conversion/ArithToLUT/ArithToLUTTest.cpp | 177 ++++++++++++++++
.../Conversion/ArithToLUT/CMakeLists.txt | 27 +++
mlir/unittests/Conversion/CMakeLists.txt | 1 +
12 files changed, 599 insertions(+)
create mode 100644 mlir/include/mlir/Conversion/ArithToLUT/ArithToLUT.h
create mode 100644 mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp
create mode 100644 mlir/lib/Conversion/ArithToLUT/CMakeLists.txt
create mode 100644 mlir/test/Conversion/ArithToLUT/extf-to-lut.mlir
create mode 100644 mlir/test/Conversion/ArithToLUT/linalg-generic-fp8-addf.mlir
create mode 100644 mlir/test/Integration/Dialect/Arith/CPU/test-lut-extf-f8.mlir
create mode 100644 mlir/unittests/Conversion/ArithToLUT/ArithToLUTTest.cpp
create mode 100644 mlir/unittests/Conversion/ArithToLUT/CMakeLists.txt
diff --git a/mlir/include/mlir/Conversion/ArithToLUT/ArithToLUT.h b/mlir/include/mlir/Conversion/ArithToLUT/ArithToLUT.h
new file mode 100644
index 0000000000000..09545ebf3f7ac
--- /dev/null
+++ b/mlir/include/mlir/Conversion/ArithToLUT/ArithToLUT.h
@@ -0,0 +1,23 @@
+//===- ArithToLUT.h - Arith FP8 extf to LUT conversion -------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_ARITHTOLUT_ARITHTOLUT_H
+#define MLIR_CONVERSION_ARITHTOLUT_ARITHTOLUT_H
+
+#include "mlir/Pass/Pass.h"
+#include <memory>
+
+namespace mlir {
+class Pass;
+
+#define GEN_PASS_DECL_CONVERTARITHFP8EXTFTOLUT
+#include "mlir/Conversion/Passes.h.inc"
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_ARITHTOLUT_ARITHTOLUT_H
diff --git a/mlir/include/mlir/Conversion/Passes.h b/mlir/include/mlir/Conversion/Passes.h
index a54b98004c3b6..51965e5df8f84 100644
--- a/mlir/include/mlir/Conversion/Passes.h
+++ b/mlir/include/mlir/Conversion/Passes.h
@@ -17,6 +17,7 @@
#include "mlir/Conversion/ArithToArmSME/ArithToArmSME.h"
#include "mlir/Conversion/ArithToEmitC/ArithToEmitCPass.h"
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
+#include "mlir/Conversion/ArithToLUT/ArithToLUT.h"
#include "mlir/Conversion/ArithToSPIRV/ArithToSPIRV.h"
#include "mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h"
#include "mlir/Conversion/ArmSMEToLLVM/ArmSMEToLLVM.h"
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index d401b56c7602d..2d3214f670a08 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -186,6 +186,32 @@ def ArithToLLVMConversionPass : Pass<"convert-arith-to-llvm"> {
];
}
+//===----------------------------------------------------------------------===//
+// ArithToLUT
+//===----------------------------------------------------------------------===//
+
+def ConvertArithFP8ExtFToLUT : Pass<"convert-arith-fp8-extf-to-lut", "ModuleOp"> {
+ let summary = "Replace arith.extf from f8 types to f32 with a memref LUT lookup";
+ let description = [{
+ Replaces each `arith.extf %v : f8X to f32` with a load from a 256-entry
+ f32 global constant table indexed by the f8 bit pattern (0–255). One table
+ is emitted per distinct f8 source format; the values are precomputed at
+ compile time via APFloat so they match the format's exact IEEE semantics.
+ This LUT approach is alternative to direct arith operations to convert to f32.
+
+ The expansion sequence is:
+ %tbl = memref.get_global @__extf_lut_<fmt> : memref<256xf32>
+ %i8 = arith.bitcast %v : f8X -> i8
+ %ui32 = arith.extui %i8 : i8 -> i32
+ %idx = arith.index_cast %ui32 : i32 -> index
+ %res = memref.load %tbl[%idx] : memref<256xf32>
+
+ Supported source types: f8E4M3FN, f8E5M2, f8E4M3FNUZ, f8E5M2FNUZ,
+ f8E4M3B11FNUZ, f8E3M4, f8E4M3.
+ }];
+ let dependentDialects = ["arith::ArithDialect", "memref::MemRefDialect"];
+}
+
//===----------------------------------------------------------------------===//
// ArithToAPFloat
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp b/mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp
new file mode 100644
index 0000000000000..ad03c70149fdb
--- /dev/null
+++ b/mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp
@@ -0,0 +1,197 @@
+//===- ArithToLUT.cpp - Replace arith.extf f8→f32 with LUT lookup ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// For each arith.extf %v : f8X to f32, this pass emits a 256-entry f32 global
+// constant (one per distinct f8 format) and replaces the op with LUT:
+// e.g.
+// ```
+// ...
+// %f32_x = arith.extf %f8_x : f8E4M3FN to f32
+// ...
+// ```
+// results in this sequence:
+// ```
+// memref.global "private" constant @__extf_lut_f8E4M3FN : memref<256xf32>
+// = dense<"0x000000000000003B0000803B...">
+// ...
+// func.func @foo (...) {
+// ...
+// %tbl = memref.get_global @__extf_lut_f8E4M3FN : memref<256xf32>
+// %i8 = arith.bitcast %v : f8X -> i8
+// %ui32 = arith.extui %i8 : i8 -> i32
+// %idx = arith.index_cast %ui32 : i32 -> index
+// %res = memref.load %tbl[%idx] : memref<256xf32>
+// ```
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/ArithToLUT/ArithToLUT.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_CONVERTARITHFP8EXTFTOLUT
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Helpers
+//===----------------------------------------------------------------------===//
+
+// Returns true for the f8 float types that need LUT-based extf lowering.
+static bool isSupportedF8Type(Type t) {
+ return isa<Float8E4M3FNType, Float8E5M2Type, Float8E4M3FNUZType,
+ Float8E5M2FNUZType, Float8E4M3B11FNUZType, Float8E3M4Type,
+ Float8E4M3Type>(t);
+}
+
+// Returns a stable, symbol-safe name for the global LUT of the given f8 type.
+
+// Returns a name for the global LUT by appending the MLIR textual
+// // representation of the given f8 type to a fixed prefix.
+static std::string lutSymbolName(FloatType srcType) {
+ std::string name = "__extf_lut_";
+ llvm::raw_string_ostream os(name);
+ srcType.print(os);
+ return name;
+}
+
+// Precomputes 256 f32 values by enumerating every 8-bit pattern for srcType.
+static SmallVector<float, 256> buildExtFLUT(FloatType srcType) {
+ const llvm::fltSemantics &sem = srcType.getFloatSemantics();
+ SmallVector<float, 256> table;
+ table.reserve(256);
+ for (unsigned i = 0; i < 256; ++i) {
+ APFloat val(sem, APInt(8, i));
+ bool losesInfo = false;
+ val.convert(APFloat::IEEEsingle(), APFloat::rmNearestTiesToEven,
+ &losesInfo);
+ table.push_back(val.convertToFloat());
+ }
+ return table;
+}
+
+// Inserts (or returns existing) memref.global constant for the given f8 type.
+static memref::GlobalOp
+getOrCreateLUT(ModuleOp module, FloatType srcType,
+ llvm::DenseMap<Type, memref::GlobalOp> &cache) {
+ auto it = cache.find(srcType);
+ if (it != cache.end())
+ return it->second;
+
+ std::string symName = lutSymbolName(srcType);
+ if (auto existing = module.lookupSymbol<memref::GlobalOp>(symName))
+ return existing;
+
+ OpBuilder builder(module.getContext());
+ builder.setInsertionPointToStart(module.getBody());
+ auto f32Ty = builder.getF32Type();
+ auto memrefTy = MemRefType::get({256}, f32Ty);
+ auto tensorTy = RankedTensorType::get({256}, f32Ty);
+
+ SmallVector<float, 256> values = buildExtFLUT(srcType);
+ auto denseAttr = DenseElementsAttr::get(tensorTy, ArrayRef<float>(values));
+
+ auto global = memref::GlobalOp::create(
+ builder, module.getLoc(),
+ /*sym_name=*/symName,
+ /*sym_visibility=*/builder.getStringAttr("private"),
+ /*type=*/memrefTy,
+ /*initial_value=*/denseAttr,
+ /*constant=*/true,
+ /*alignment=*/builder.getI64IntegerAttr(64));
+ cache[srcType] = global;
+ return global;
+}
+
+//===----------------------------------------------------------------------===//
+// Rewrite pattern
+//===----------------------------------------------------------------------===//
+
+struct ExtFToLUTPattern : public OpRewritePattern<arith::ExtFOp> {
+ ExtFToLUTPattern(MLIRContext *ctx,
+ llvm::DenseMap<Type, memref::GlobalOp> &lutCache)
+ : OpRewritePattern(ctx), lutCache(lutCache) {}
+
+ LogicalResult matchAndRewrite(arith::ExtFOp op,
+ PatternRewriter &rewriter) const override {
+ Type srcTy = op.getIn().getType();
+ Type dstTy = op.getType();
+
+ if (!isSupportedF8Type(srcTy) || !isa<Float32Type>(dstTy))
+ return failure();
+
+ auto srcFloatTy = cast<FloatType>(srcTy);
+ auto module = op->getParentOfType<ModuleOp>();
+ memref::GlobalOp global = getOrCreateLUT(module, srcFloatTy, lutCache);
+
+ Location loc = op.getLoc();
+ auto memrefTy = cast<MemRefType>(global.getType());
+
+ // %tbl = memref.get_global @__extf_lut_<fmt>
+ Value tbl = memref::GetGlobalOp::create(rewriter, loc, memrefTy,
+ global.getSymName());
+ // %i8 = arith.bitcast %in : f8X -> i8
+ Value i8val = arith::BitcastOp::create(rewriter, loc,
+ rewriter.getIntegerType(8),
+ op.getIn());
+
+ // %ui32 = arith.extui %i8 : i8 -> i32
+ Value ui32val =
+ arith::ExtUIOp::create(rewriter, loc, rewriter.getI32Type(), i8val);
+
+ // %idx = arith.index_cast %ui32 : i32 -> index
+ Value idx = arith::IndexCastOp::create(rewriter, loc,
+ rewriter.getIndexType(), ui32val);
+
+ // %res = memref.load %tbl[%idx]
+ Value result = memref::LoadOp::create(rewriter, loc, tbl, idx);
+
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+
+private:
+ llvm::DenseMap<Type, memref::GlobalOp> &lutCache;
+};
+
+//===----------------------------------------------------------------------===//
+// Pass
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct ConvertArithFP8ExtFToLUTPass
+ : public impl::ConvertArithFP8ExtFToLUTBase<ConvertArithFP8ExtFToLUTPass> {
+
+ void runOnOperation() override {
+ ModuleOp module = getOperation();
+ MLIRContext *ctx = &getContext();
+
+ // Cache so each format gets exactly one global, inserted once.
+ llvm::DenseMap<Type, memref::GlobalOp> lutCache;
+
+ RewritePatternSet patterns(ctx);
+ patterns.add<ExtFToLUTPattern>(ctx, lutCache);
+
+ if (failed(applyPatternsGreedily(module, std::move(patterns))))
+ signalPassFailure();
+ }
+};
+} // namespace
diff --git a/mlir/lib/Conversion/ArithToLUT/CMakeLists.txt b/mlir/lib/Conversion/ArithToLUT/CMakeLists.txt
new file mode 100644
index 0000000000000..e94e4a1835e11
--- /dev/null
+++ b/mlir/lib/Conversion/ArithToLUT/CMakeLists.txt
@@ -0,0 +1,18 @@
+add_mlir_conversion_library(MLIRArithToLUT
+ ArithToLUT.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/ArithToLUT
+
+ DEPENDS
+ MLIRConversionPassIncGen
+
+ LINK_COMPONENTS
+ Core
+
+ LINK_LIBS PUBLIC
+ MLIRArithDialect
+ MLIRMemRefDialect
+ MLIRPass
+ MLIRTransforms
+)
diff --git a/mlir/lib/Conversion/CMakeLists.txt b/mlir/lib/Conversion/CMakeLists.txt
index e17988b12cade..b5f298a777d5d 100644
--- a/mlir/lib/Conversion/CMakeLists.txt
+++ b/mlir/lib/Conversion/CMakeLists.txt
@@ -6,6 +6,7 @@ add_subdirectory(ArithAndMathToAPFloat)
add_subdirectory(ArithToArmSME)
add_subdirectory(ArithToEmitC)
add_subdirectory(ArithToLLVM)
+add_subdirectory(ArithToLUT)
add_subdirectory(ArithToSPIRV)
add_subdirectory(ArmNeon2dToIntr)
add_subdirectory(ArmSMEToSCF)
diff --git a/mlir/test/Conversion/ArithToLUT/extf-to-lut.mlir b/mlir/test/Conversion/ArithToLUT/extf-to-lut.mlir
new file mode 100644
index 0000000000000..2fce8572df790
--- /dev/null
+++ b/mlir/test/Conversion/ArithToLUT/extf-to-lut.mlir
@@ -0,0 +1,36 @@
+// RUN: mlir-opt %s --convert-arith-fp8-extf-to-lut | FileCheck %s
+
+// CHECK-DAG: memref.global "private" constant @__extf_lut_f8E4M3FN : memref<256xf32>
+// CHECK-DAG: memref.global "private" constant @__extf_lut_f8E5M2 : memref<256xf32>
+
+// Single f8E4M3FN input — table should appear once, op replaced by LUT sequence.
+func.func @extf_f8E4M3FN(%a: f8E4M3FN) -> f32 {
+ // CHECK-LABEL: @extf_f8E4M3FN
+ // CHECK: memref.get_global @__extf_lut_f8E4M3FN : memref<256xf32>
+ // CHECK: arith.bitcast {{.*}} : f8E4M3FN to i8
+ // CHECK: arith.extui {{.*}} : i8 to i32
+ // CHECK: arith.index_cast {{.*}} : i32 to index
+ // CHECK: memref.load {{.*}}[{{.*}}] : memref<256xf32>
+ // CHECK-NOT: arith.extf
+ %r = arith.extf %a : f8E4M3FN to f32
+ return %r : f32
+}
+
+// Two f8E4M3FN inputs — same table reused, not duplicated.
+func.func @extf_f8E4M3FN_twice(%a: f8E4M3FN, %b: f8E4M3FN) -> f32 {
+ // CHECK-LABEL: @extf_f8E4M3FN_twice
+ // CHECK-COUNT-2: memref.get_global @__extf_lut_f8E4M3FN
+ %ra = arith.extf %a : f8E4M3FN to f32
+ %rb = arith.extf %b : f8E4M3FN to f32
+ %sum = arith.addf %ra, %rb : f32
+ return %sum : f32
+}
+
+// Second f8 format — separate table emitted.
+func.func @extf_f8E5M2(%a: f8E5M2) -> f32 {
+ // CHECK-LABEL: @extf_f8E5M2
+ // CHECK: memref.get_global @__extf_lut_f8E5M2 : memref<256xf32>
+ // CHECK-NOT: arith.extf
+ %r = arith.extf %a : f8E5M2 to f32
+ return %r : f32
+}
diff --git a/mlir/test/Conversion/ArithToLUT/linalg-generic-fp8-addf.mlir b/mlir/test/Conversion/ArithToLUT/linalg-generic-fp8-addf.mlir
new file mode 100644
index 0000000000000..200a8262a9669
--- /dev/null
+++ b/mlir/test/Conversion/ArithToLUT/linalg-generic-fp8-addf.mlir
@@ -0,0 +1,33 @@
+// RUN: mlir-opt %s --convert-arith-fp8-extf-to-lut | FileCheck %s
+
+// linalg.generic over tensor<1024xf8E4M3FN>: both extf ops inside the body
+// should be replaced by LUT-lookup sequences; one shared table is emitted.
+
+// CHECK-DAG: memref.global "private" constant @__extf_lut_f8E4M3FN : memref<256xf32>
+
+func.func @linalg_addf_fp8(
+ %a: tensor<1024xf8E4M3FN>,
+ %b: tensor<1024xf8E4M3FN>) -> tensor<1024xf32> {
+ // CHECK-LABEL: @linalg_addf_fp8
+ // CHECK: linalg.generic
+ // CHECK-COUNT-2: memref.get_global @__extf_lut_f8E4M3FN : memref<256xf32>
+ // CHECK: arith.addf {{.*}} : f32
+ // CHECK-NOT: arith.extf
+ %init = tensor.empty() : tensor<1024xf32>
+ %result = linalg.generic {
+ indexing_maps = [
+ affine_map<(d0) -> (d0)>,
+ affine_map<(d0) -> (d0)>,
+ affine_map<(d0) -> (d0)>
+ ],
+ iterator_types = ["parallel"]
+ } ins(%a, %b : tensor<1024xf8E4M3FN>, tensor<1024xf8E4M3FN>)
+ outs(%init : tensor<1024xf32>) {
+ ^bb0(%in_a: f8E4M3FN, %in_b: f8E4M3FN, %out: f32):
+ %ra = arith.extf %in_a : f8E4M3FN to f32
+ %rb = arith.extf %in_b : f8E4M3FN to f32
+ %sum = arith.addf %ra, %rb : f32
+ linalg.yield %sum : f32
+ } -> tensor<1024xf32>
+ return %result : tensor<1024xf32>
+}
diff --git a/mlir/test/Integration/Dialect/Arith/CPU/test-lut-extf-f8.mlir b/mlir/test/Integration/Dialect/Arith/CPU/test-lut-extf-f8.mlir
new file mode 100644
index 0000000000000..2b89887493114
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Arith/CPU/test-lut-extf-f8.mlir
@@ -0,0 +1,59 @@
+// Verify that the LUT-based f8E4M3FN → f32 lowering produces the correct f32
+// value for a selection of well-known bit patterns.
+
+// RUN: mlir-opt %s \
+// RUN: --convert-arith-fp8-extf-to-lut \
+// RUN: --finalize-memref-to-llvm \
+// RUN: --convert-arith-to-llvm \
+// RUN: --convert-vector-to-llvm \
+// RUN: --convert-func-to-llvm \
+// RUN: --reconcile-unrealized-casts \
+// RUN: | mlir-runner -e entry --entry-point-result=void \
+// RUN: --shared-libs=%mlir_c_runner_utils \
+// RUN: | FileCheck %s --match-full-lines
+
+func.func @check(%bits: i8) {
+ %f8 = arith.bitcast %bits : i8 to f8E4M3FN
+ %f32 = arith.extf %f8 : f8E4M3FN to f32
+ vector.print %f32 : f32
+ return
+}
+
+func.func @entry() {
+ // +0.0 (bit pattern 0x00)
+ %b0 = arith.constant 0 : i8
+ // CHECK: 0
+ func.call @check(%b0) : (i8) -> ()
+
+ // -0.0 (bit pattern 0x80)
+ %b1 = arith.constant -128 : i8
+ // CHECK: -0
+ func.call @check(%b1) : (i8) -> ()
+
+ // 1.0 (bit pattern 0x38: exp=7, mant=0)
+ %b2 = arith.constant 56 : i8
+ // CHECK: 1
+ func.call @check(%b2) : (i8) -> ()
+
+ // -1.0 (bit pattern 0xB8)
+ %b3 = arith.constant -72 : i8
+ // CHECK: -1
+ func.call @check(%b3) : (i8) -> ()
+
+ // 2.0 (bit pattern 0x40: exp=8, mant=0)
+ %b4 = arith.constant 64 : i8
+ // CHECK: 2
+ func.call @check(%b4) : (i8) -> ()
+
+ // 0.5 (bit pattern 0x30: exp=6, mant=0)
+ %b5 = arith.constant 48 : i8
+ // CHECK: 0.5
+ func.call @check(%b5) : (i8) -> ()
+
+ // max finite: 448.0 (bit pattern 0x7E: exp=15, mant=0b110)
+ %b6 = arith.constant 126 : i8
+ // CHECK: 448
+ func.call @check(%b6) : (i8) -> ()
+
+ return
+}
diff --git a/mlir/unittests/Conversion/ArithToLUT/ArithToLUTTest.cpp b/mlir/unittests/Conversion/ArithToLUT/ArithToLUTTest.cpp
new file mode 100644
index 0000000000000..884cc5dc2497b
--- /dev/null
+++ b/mlir/unittests/Conversion/ArithToLUT/ArithToLUTTest.cpp
@@ -0,0 +1,177 @@
+//===- ArithToLUTTest.cpp - Exhaustive correctness test for the LUT pass --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// For every supported f8 format, iterate all 256 bit patterns, run each
+// through the LUT-lowered extf function via the MLIR JIT, and compare the
+// result bit-for-bit against the APFloat reference value computed the same
+// way buildExtFLUT does at compile time.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/ArithToLUT/ArithToLUT.h"
+#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
+#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
+#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
+#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/ExecutionEngine/ExecutionEngine.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/InitAllDialects.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/Support/TargetSelect.h"
+
+#include "gmock/gmock.h"
+
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+#include <functional>
+#include <string>
+
+// JIT is unavailable on some platforms.
+#ifdef __sparc__
+#define SKIP_WITHOUT_JIT(x) DISABLED_##x
+#else
+#define SKIP_WITHOUT_JIT(x) x
+#endif
+
+using namespace mlir;
+
+#if !defined(_WIN32) && !defined(_AIX)
+
+static struct LLVMInitializer {
+ LLVMInitializer() {
+ llvm::InitializeNativeTarget();
+ llvm::InitializeNativeTargetAsmPrinter();
+ }
+} initializer;
+
+static LogicalResult lowerToLLVM(ModuleOp module) {
+ PassManager pm(module->getName());
+ pm.addPass(createConvertArithFP8ExtFToLUT());
+ pm.addPass(createFinalizeMemRefToLLVMConversionPass());
+ pm.addNestedPass<func::FuncOp>(createArithToLLVMConversionPass());
+ pm.addPass(createConvertFuncToLLVMPass());
+ pm.addPass(createReconcileUnrealizedCastsPass());
+ return pm.run(module);
+}
+
+// Builds a module with function:
+// func @test(%arg0: i32) -> f32 { trunci i32 to i8; bitcast i8 to f8; extf f8 to f32 }
+static std::string makeModule(const char *f8TypeName) {
+ std::string s;
+ s += "func.func @test(%arg0: i32) -> f32 "
+ "attributes { llvm.emit_c_interface } {\n";
+ s += " %i8 = arith.trunci %arg0 : i32 to i8\n";
+ s += std::string(" %f8 = arith.bitcast %i8 : i8 to ") + f8TypeName + "\n";
+ s += std::string(" %f32 = arith.extf %f8 : ") + f8TypeName + " to f32\n";
+ s += " return %f32 : f32\n}\n";
+ return s;
+}
+
+struct F8Format {
+ const char *mlirName;
+ std::function<FloatType(MLIRContext *)> getType;
+};
+
+static void runExhaustiveTest(const F8Format &fmt) {
+ DialectRegistry registry;
+ registerAllDialects(registry);
+ registerBuiltinDialectTranslation(registry);
+ registerLLVMDialectTranslation(registry);
+ MLIRContext ctx(registry);
+
+ OwningOpRef<ModuleOp> module =
+ parseSourceString<ModuleOp>(makeModule(fmt.mlirName), &ctx);
+ ASSERT_TRUE(!!module) << "parse failed for " << fmt.mlirName;
+ ASSERT_TRUE(succeeded(lowerToLLVM(*module)))
+ << "lowering failed for " << fmt.mlirName;
+
+ auto jitOrError = ExecutionEngine::create(*module);
+ ASSERT_TRUE(!!jitOrError) << "JIT creation failed for " << fmt.mlirName;
+ auto jit = std::move(jitOrError.get());
+
+ const llvm::fltSemantics &sem = fmt.getType(&ctx).getFloatSemantics();
+
+ for (int i = 0; i < 256; ++i) {
+ float got = 0.0f;
+ int32_t input = i;
+ llvm::Error err =
+ jit->invoke("test", input, ExecutionEngine::Result<float>(got));
+ ASSERT_FALSE(err) << llvm::toString(std::move(err));
+
+ llvm::APFloat ref(sem, llvm::APInt(8, static_cast<uint64_t>(i)));
+ bool lossy = false;
+ ref.convert(llvm::APFloat::IEEEsingle(),
+ llvm::APFloat::rmNearestTiesToEven, &lossy);
+ float expected = ref.convertToFloat();
+
+ if (std::isnan(expected)) {
+ EXPECT_TRUE(std::isnan(got))
+ << fmt.mlirName << " pattern " << i
+ << ": expected NaN, got " << got;
+ } else {
+ uint32_t gotBits = 0, expBits = 0;
+ std::memcpy(&gotBits, &got, 4);
+ std::memcpy(&expBits, &expected, 4);
+ EXPECT_EQ(gotBits, expBits)
+ << fmt.mlirName << " pattern " << i
+ << ": expected " << expected << ", got " << got;
+ }
+ }
+}
+
+TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E4M3FN)) {
+ runExhaustiveTest({"f8E4M3FN", [](MLIRContext *ctx) -> FloatType {
+ return Float8E4M3FNType::get(ctx);
+ }});
+}
+
+TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E5M2)) {
+ runExhaustiveTest({"f8E5M2", [](MLIRContext *ctx) -> FloatType {
+ return Float8E5M2Type::get(ctx);
+ }});
+}
+
+TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E4M3FNUZ)) {
+ runExhaustiveTest({"f8E4M3FNUZ", [](MLIRContext *ctx) -> FloatType {
+ return Float8E4M3FNUZType::get(ctx);
+ }});
+}
+
+TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E5M2FNUZ)) {
+ runExhaustiveTest({"f8E5M2FNUZ", [](MLIRContext *ctx) -> FloatType {
+ return Float8E5M2FNUZType::get(ctx);
+ }});
+}
+
+TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E4M3B11FNUZ)) {
+ runExhaustiveTest({"f8E4M3B11FNUZ", [](MLIRContext *ctx) -> FloatType {
+ return Float8E4M3B11FNUZType::get(ctx);
+ }});
+}
+
+TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E3M4)) {
+ runExhaustiveTest({"f8E3M4", [](MLIRContext *ctx) -> FloatType {
+ return Float8E3M4Type::get(ctx);
+ }});
+}
+
+TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E4M3)) {
+ runExhaustiveTest({"f8E4M3", [](MLIRContext *ctx) -> FloatType {
+ return Float8E4M3Type::get(ctx);
+ }});
+}
+
+#endif // !_WIN32 && !_AIX
diff --git a/mlir/unittests/Conversion/ArithToLUT/CMakeLists.txt b/mlir/unittests/Conversion/ArithToLUT/CMakeLists.txt
new file mode 100644
index 0000000000000..249f007096136
--- /dev/null
+++ b/mlir/unittests/Conversion/ArithToLUT/CMakeLists.txt
@@ -0,0 +1,27 @@
+if(MLIR_ENABLE_EXECUTION_ENGINE)
+ set(LLVM_LINK_COMPONENTS
+ nativecodegen
+ native
+ orcjit
+ support
+ )
+
+ add_mlir_unittest(MLIRArithToLUTTests
+ ArithToLUTTest.cpp
+ )
+
+ mlir_target_link_libraries(MLIRArithToLUTTests
+ PRIVATE
+ MLIRArithToLUT
+ MLIRArithToLLVM
+ MLIRFuncToLLVM
+ MLIRMemRefToLLVM
+ MLIRReconcileUnrealizedCasts
+ MLIRRegisterAllDialects
+ )
+
+ target_link_libraries(MLIRArithToLUTTests
+ PRIVATE
+ MLIRExecutionEngine
+ )
+endif()
diff --git a/mlir/unittests/Conversion/CMakeLists.txt b/mlir/unittests/Conversion/CMakeLists.txt
index 2dee5e7dac90c..af22373a83c26 100644
--- a/mlir/unittests/Conversion/CMakeLists.txt
+++ b/mlir/unittests/Conversion/CMakeLists.txt
@@ -1 +1,2 @@
add_subdirectory(PDLToPDLInterp)
+add_subdirectory(ArithToLUT)
>From 1b163344f1e9fb3914f640b5a8b2f259c87fd5e8 Mon Sep 17 00:00:00 2001
From: Javed Absar <javed.absar at gmail.com>
Date: Thu, 7 May 2026 10:18:05 -0400
Subject: [PATCH 2/3] Fix comments
---
mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp | 14 --------------
1 file changed, 14 deletions(-)
diff --git a/mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp b/mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp
index ad03c70149fdb..ac1615543d788 100644
--- a/mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp
+++ b/mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp
@@ -51,10 +51,6 @@ namespace mlir {
using namespace mlir;
-//===----------------------------------------------------------------------===//
-// Helpers
-//===----------------------------------------------------------------------===//
-
// Returns true for the f8 float types that need LUT-based extf lowering.
static bool isSupportedF8Type(Type t) {
return isa<Float8E4M3FNType, Float8E5M2Type, Float8E4M3FNUZType,
@@ -62,8 +58,6 @@ static bool isSupportedF8Type(Type t) {
Float8E4M3Type>(t);
}
-// Returns a stable, symbol-safe name for the global LUT of the given f8 type.
-
// Returns a name for the global LUT by appending the MLIR textual
// // representation of the given f8 type to a fixed prefix.
static std::string lutSymbolName(FloatType srcType) {
@@ -121,10 +115,6 @@ getOrCreateLUT(ModuleOp module, FloatType srcType,
return global;
}
-//===----------------------------------------------------------------------===//
-// Rewrite pattern
-//===----------------------------------------------------------------------===//
-
struct ExtFToLUTPattern : public OpRewritePattern<arith::ExtFOp> {
ExtFToLUTPattern(MLIRContext *ctx,
llvm::DenseMap<Type, memref::GlobalOp> &lutCache)
@@ -172,10 +162,6 @@ struct ExtFToLUTPattern : public OpRewritePattern<arith::ExtFOp> {
llvm::DenseMap<Type, memref::GlobalOp> &lutCache;
};
-//===----------------------------------------------------------------------===//
-// Pass
-//===----------------------------------------------------------------------===//
-
namespace {
struct ConvertArithFP8ExtFToLUTPass
: public impl::ConvertArithFP8ExtFToLUTBase<ConvertArithFP8ExtFToLUTPass> {
>From 3546193779b444bc9540a77e736a9ab11dd16692 Mon Sep 17 00:00:00 2001
From: Javed Absar <javed.absar at gmail.com>
Date: Sat, 9 May 2026 11:56:03 -0400
Subject: [PATCH 3/3] address review comments.
---
.../mlir/Conversion/ArithToLUT/ArithToLUT.h | 5 +-
mlir/include/mlir/Conversion/Passes.td | 45 +++--
mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp | 172 ++++++++++++------
mlir/lib/Conversion/ArithToLUT/CMakeLists.txt | 1 +
.../Conversion/ArithToLUT/extf-to-lut.mlir | 31 +++-
.../ArithToLUT/linalg-generic-fp8-addf.mlir | 33 ----
.../Dialect/Arith/CPU/test-lut-extf-f8.mlir | 2 +-
.../Conversion/ArithToLUT/ArithToLUTTest.cpp | 64 ++++---
8 files changed, 211 insertions(+), 142 deletions(-)
delete mode 100644 mlir/test/Conversion/ArithToLUT/linalg-generic-fp8-addf.mlir
diff --git a/mlir/include/mlir/Conversion/ArithToLUT/ArithToLUT.h b/mlir/include/mlir/Conversion/ArithToLUT/ArithToLUT.h
index 09545ebf3f7ac..7e7049fecf03a 100644
--- a/mlir/include/mlir/Conversion/ArithToLUT/ArithToLUT.h
+++ b/mlir/include/mlir/Conversion/ArithToLUT/ArithToLUT.h
@@ -1,4 +1,5 @@
-//===- ArithToLUT.h - Arith FP8 extf to LUT conversion -------*- C++ -*-===//
+//===- ArithToLUT.h - Arith narrow-float extf to LUT conversion --*- C++
+//-*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -15,7 +16,7 @@
namespace mlir {
class Pass;
-#define GEN_PASS_DECL_CONVERTARITHFP8EXTFTOLUT
+#define GEN_PASS_DECL_CONVERTARITHEXTFTOLUT
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 2d3214f670a08..caab74e87c9a7 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -190,26 +190,35 @@ def ArithToLLVMConversionPass : Pass<"convert-arith-to-llvm"> {
// ArithToLUT
//===----------------------------------------------------------------------===//
-def ConvertArithFP8ExtFToLUT : Pass<"convert-arith-fp8-extf-to-lut", "ModuleOp"> {
- let summary = "Replace arith.extf from f8 types to f32 with a memref LUT lookup";
+def ConvertArithExtFToLUT : Pass<"convert-arith-extf-to-lut", "ModuleOp"> {
+ let summary = "Replace arith.extf from narrow float types to f32 with a memref LUT lookup";
let description = [{
- Replaces each `arith.extf %v : f8X to f32` with a load from a 256-entry
- f32 global constant table indexed by the f8 bit pattern (0–255). One table
- is emitted per distinct f8 source format; the values are precomputed at
- compile time via APFloat so they match the format's exact IEEE semantics.
- This LUT approach is alternative to direct arith operations to convert to f32.
-
- The expansion sequence is:
- %tbl = memref.get_global @__extf_lut_<fmt> : memref<256xf32>
- %i8 = arith.bitcast %v : f8X -> i8
- %ui32 = arith.extui %i8 : i8 -> i32
- %idx = arith.index_cast %ui32 : i32 -> index
- %res = memref.load %tbl[%idx] : memref<256xf32>
-
- Supported source types: f8E4M3FN, f8E5M2, f8E4M3FNUZ, f8E5M2FNUZ,
- f8E4M3B11FNUZ, f8E3M4, f8E4M3.
+ Replaces each `arith.extf %v : narrowFP to f32` (scalar or vector) with a
+ lookup into a precomputed f32 global constant table indexed by the source
+ bit pattern. One table is emitted per distinct source format; values are
+ precomputed at compile time via APFloat so they match exact IEEE semantics.
+ Table size is 2^N entries for an N-bit source type (e.g. 16 for f4, 64 for
+ f6, 256 for f8).
+
+ Scalar expansion:
+ %tbl = memref.get_global @__extf_lut_<fmt> : memref<Mxf32>
+ %iN = arith.bitcast %v : fNX -> iN
+ %idx = arith.index_castui %iN : iN -> index
+ %res = memref.load %tbl[%idx] : memref<Mxf32>
+
+ Vector expansion (uses vector.gather for all lanes in one op):
+ %tbl = memref.get_global @__extf_lut_<fmt> : memref<Mxf32>
+ %iNvec = arith.bitcast %v : vector<NxFP> -> vector<NxiN>
+ %offsets = arith.extui %iNvec : vector<NxiN> -> vector<Nxi32>
+ %mask = vector.constant_mask [N] : vector<Nxi1>
+ %pass = arith.constant dense<0> : vector<Nxf32>
+ %res = vector.gather %tbl[%c0][%offsets], %mask, %pass
+
+ Supported source types: f4E2M1FN, f6E2M3FN, f6E3M2FN, f8E4M3FN, f8E5M2,
+ f8E4M3FNUZ, f8E5M2FNUZ, f8E4M3B11FNUZ, f8E3M4, f8E4M3.
}];
- let dependentDialects = ["arith::ArithDialect", "memref::MemRefDialect"];
+ let dependentDialects = ["arith::ArithDialect", "memref::MemRefDialect",
+ "vector::VectorDialect"];
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp b/mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp
index ac1615543d788..73eee2ecc7370 100644
--- a/mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp
+++ b/mlir/lib/Conversion/ArithToLUT/ArithToLUT.cpp
@@ -1,4 +1,4 @@
-//===- ArithToLUT.cpp - Replace arith.extf f8→f32 with LUT lookup ---------===//
+//===- ArithToLUT.cpp - Replace arith.extf narrow-float→f32 with LUT ------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@@ -6,8 +6,9 @@
//
//===----------------------------------------------------------------------===//
//
-// For each arith.extf %v : f8X to f32, this pass emits a 256-entry f32 global
-// constant (one per distinct f8 format) and replaces the op with LUT:
+// For each arith.extf %v : narrowFP to f32 (scalar or vector), this pass emits
+// a 2^N-entry f32 global constant (one per distinct source format) and replaces
+// the op with a LUT lookup. Scalars use memref.load; vectors use vector.gather.
// e.g.
// ```
// ...
@@ -21,11 +22,10 @@
// ...
// func.func @foo (...) {
// ...
-// %tbl = memref.get_global @__extf_lut_f8E4M3FN : memref<256xf32>
-// %i8 = arith.bitcast %v : f8X -> i8
-// %ui32 = arith.extui %i8 : i8 -> i32
-// %idx = arith.index_cast %ui32 : i32 -> index
-// %res = memref.load %tbl[%idx] : memref<256xf32>
+// %tbl = memref.get_global @__extf_lut_f8E4M3FN : memref<256xf32>
+// %i8 = arith.bitcast %v : f8X -> i8
+// %idx = arith.index_castui %i8 : i8 -> index
+// %res = memref.load %tbl[%idx] : memref<256xf32>
// ```
//===----------------------------------------------------------------------===//
@@ -33,11 +33,12 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Pass/Pass.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/WalkPatternRewriteDriver.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/APInt.h"
#include "llvm/ADT/DenseMap.h"
@@ -45,21 +46,22 @@
#include "llvm/ADT/StringRef.h"
namespace mlir {
-#define GEN_PASS_DEF_CONVERTARITHFP8EXTFTOLUT
+#define GEN_PASS_DEF_CONVERTARITHEXTFTOLUT
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir
using namespace mlir;
-// Returns true for the f8 float types that need LUT-based extf lowering.
-static bool isSupportedF8Type(Type t) {
- return isa<Float8E4M3FNType, Float8E5M2Type, Float8E4M3FNUZType,
+// Returns true for the narrow float types that support LUT-based extf lowering.
+static bool isSupportedNarrowFloatType(Type t) {
+ return isa<Float4E2M1FNType, Float6E2M3FNType, Float6E3M2FNType,
+ Float8E4M3FNType, Float8E5M2Type, Float8E4M3FNUZType,
Float8E5M2FNUZType, Float8E4M3B11FNUZType, Float8E3M4Type,
Float8E4M3Type>(t);
}
// Returns a name for the global LUT by appending the MLIR textual
-// // representation of the given f8 type to a fixed prefix.
+// representation of the given f8 type to a fixed prefix.
static std::string lutSymbolName(FloatType srcType) {
std::string name = "__extf_lut_";
llvm::raw_string_ostream os(name);
@@ -67,13 +69,15 @@ static std::string lutSymbolName(FloatType srcType) {
return name;
}
-// Precomputes 256 f32 values by enumerating every 8-bit pattern for srcType.
-static SmallVector<float, 256> buildExtFLUT(FloatType srcType) {
+// Precomputes 2^N f32 values by enumerating every N-bit pattern for srcType.
+static SmallVector<float> buildExtFLUT(FloatType srcType) {
const llvm::fltSemantics &sem = srcType.getFloatSemantics();
- SmallVector<float, 256> table;
- table.reserve(256);
- for (unsigned i = 0; i < 256; ++i) {
- APFloat val(sem, APInt(8, i));
+ unsigned bitWidth = srcType.getWidth();
+ unsigned numEntries = 1u << bitWidth;
+ SmallVector<float> table;
+ table.reserve(numEntries);
+ for (unsigned i = 0; i < numEntries; ++i) {
+ APFloat val(sem, APInt(bitWidth, i));
bool losesInfo = false;
val.convert(APFloat::IEEEsingle(), APFloat::rmNearestTiesToEven,
&losesInfo);
@@ -84,27 +88,47 @@ static SmallVector<float, 256> buildExtFLUT(FloatType srcType) {
// Inserts (or returns existing) memref.global constant for the given f8 type.
static memref::GlobalOp
-getOrCreateLUT(ModuleOp module, FloatType srcType,
+getOrCreateLUT(Operation *symbolTableOp, FloatType srcType,
llvm::DenseMap<Type, memref::GlobalOp> &cache) {
auto it = cache.find(srcType);
if (it != cache.end())
return it->second;
- std::string symName = lutSymbolName(srcType);
- if (auto existing = module.lookupSymbol<memref::GlobalOp>(symName))
- return existing;
+ // Choose a symbol name for the LUT global, handling potential collisions.
+ //
+ // Start with the canonical name (e.g. "__extf_lut_f8E4M3FN") and look it up:
+ // - Not found: the name is free; exit the loop and create a new global.
+ // - Found a memref.global: a compatible LUT already exists (e.g. the pass
+ // was applied twice, or the user pre-defined the table); reuse it.
+ // - Found a non-memref.global symbol: a symbol of a different op kind
+ // already holds this name (e.g. a function named "__extf_lut_f8E4M3FN").
+ // Retry with a numeric suffix ("__extf_lut_f8E4M3FN_0", "_1", …) until
+ // we find a free slot.
+ std::string baseSymName = lutSymbolName(srcType);
+ std::string symName = baseSymName;
+ for (unsigned suffix = 0;; ++suffix) {
+ auto *existing = SymbolTable::lookupSymbolIn(symbolTableOp, symName);
+ if (!existing)
+ break;
+ if (auto globalOp = dyn_cast<memref::GlobalOp>(existing)) {
+ cache[srcType] = globalOp;
+ return globalOp;
+ }
+ symName = baseSymName + "_" + std::to_string(suffix);
+ }
- OpBuilder builder(module.getContext());
- builder.setInsertionPointToStart(module.getBody());
+ OpBuilder builder(symbolTableOp->getContext());
+ builder.setInsertionPointToStart(&symbolTableOp->getRegion(0).front());
auto f32Ty = builder.getF32Type();
- auto memrefTy = MemRefType::get({256}, f32Ty);
- auto tensorTy = RankedTensorType::get({256}, f32Ty);
+ int64_t numEntries = 1LL << srcType.getWidth();
+ auto memrefTy = MemRefType::get({numEntries}, f32Ty);
+ auto tensorTy = RankedTensorType::get({numEntries}, f32Ty);
- SmallVector<float, 256> values = buildExtFLUT(srcType);
+ SmallVector<float> values = buildExtFLUT(srcType);
auto denseAttr = DenseElementsAttr::get(tensorTy, ArrayRef<float>(values));
auto global = memref::GlobalOp::create(
- builder, module.getLoc(),
+ builder, symbolTableOp->getLoc(),
/*sym_name=*/symName,
/*sym_visibility=*/builder.getStringAttr("private"),
/*type=*/memrefTy,
@@ -124,36 +148,69 @@ struct ExtFToLUTPattern : public OpRewritePattern<arith::ExtFOp> {
PatternRewriter &rewriter) const override {
Type srcTy = op.getIn().getType();
Type dstTy = op.getType();
+ Location loc = op.getLoc();
- if (!isSupportedF8Type(srcTy) || !isa<Float32Type>(dstTy))
+ // Scalar narrowFP → f32
+ if (isSupportedNarrowFloatType(srcTy) && isa<Float32Type>(dstTy)) {
+ auto srcFloatTy = cast<FloatType>(srcTy);
+ Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(op);
+ memref::GlobalOp global =
+ getOrCreateLUT(symbolTableOp, srcFloatTy, lutCache);
+ auto memrefTy = cast<MemRefType>(global.getType());
+
+ Value tbl = memref::GetGlobalOp::create(rewriter, loc, memrefTy,
+ global.getSymName());
+ Value iNval = arith::BitcastOp::create(
+ rewriter, loc, rewriter.getIntegerType(srcFloatTy.getWidth()),
+ op.getIn());
+ Value idx = arith::IndexCastUIOp::create(rewriter, loc,
+ rewriter.getIndexType(), iNval);
+ Value result = memref::LoadOp::create(rewriter, loc, tbl, idx);
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+
+ // Vector<NxNarrowFP> → vector<NxF32>
+ auto srcVecTy = dyn_cast<VectorType>(srcTy);
+ auto dstVecTy = dyn_cast<VectorType>(dstTy);
+ if (!srcVecTy || !dstVecTy)
+ return failure();
+ if (!isSupportedNarrowFloatType(srcVecTy.getElementType()) ||
+ !isa<Float32Type>(dstVecTy.getElementType()))
return failure();
- auto srcFloatTy = cast<FloatType>(srcTy);
- auto module = op->getParentOfType<ModuleOp>();
- memref::GlobalOp global = getOrCreateLUT(module, srcFloatTy, lutCache);
-
- Location loc = op.getLoc();
+ auto srcElemTy = cast<FloatType>(srcVecTy.getElementType());
+ Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(op);
+ memref::GlobalOp global =
+ getOrCreateLUT(symbolTableOp, srcElemTy, lutCache);
auto memrefTy = cast<MemRefType>(global.getType());
- // %tbl = memref.get_global @__extf_lut_<fmt>
Value tbl = memref::GetGlobalOp::create(rewriter, loc, memrefTy,
global.getSymName());
- // %i8 = arith.bitcast %in : f8X -> i8
- Value i8val = arith::BitcastOp::create(rewriter, loc,
- rewriter.getIntegerType(8),
- op.getIn());
-
- // %ui32 = arith.extui %i8 : i8 -> i32
- Value ui32val =
- arith::ExtUIOp::create(rewriter, loc, rewriter.getI32Type(), i8val);
-
- // %idx = arith.index_cast %ui32 : i32 -> index
- Value idx = arith::IndexCastOp::create(rewriter, loc,
- rewriter.getIndexType(), ui32val);
-
- // %res = memref.load %tbl[%idx]
- Value result = memref::LoadOp::create(rewriter, loc, tbl, idx);
-
+ // bitcast vector<NxNarrowFP> -> vector<NxiW>
+ unsigned bitWidth = srcElemTy.getWidth();
+ auto iNVecTy =
+ VectorType::get(srcVecTy.getShape(), rewriter.getIntegerType(bitWidth));
+ Value iNvec = arith::BitcastOp::create(rewriter, loc, iNVecTy, op.getIn());
+
+ // extui vector<NxiW> -> vector<Nxi32> (gather offsets)
+ auto i32VecTy = VectorType::get(srcVecTy.getShape(), rewriter.getI32Type());
+ Value offsets = arith::ExtUIOp::create(rewriter, loc, i32VecTy, iNvec);
+
+ // all-true mask
+ auto maskVecTy = VectorType::get(srcVecTy.getShape(), rewriter.getI1Type());
+ Value mask = vector::ConstantMaskOp::create(
+ rewriter, loc, maskVecTy, vector::ConstantMaskKind::AllTrue);
+
+ // passthru: dense<0.0> vector<Nxf32>
+ Value passthru = arith::ConstantOp::create(
+ rewriter, loc, DenseElementsAttr::get(dstVecTy, ArrayRef<float>{0.0f}));
+
+ // base index into the LUT
+ Value c0 = arith::ConstantIndexOp::create(rewriter, loc, 0);
+
+ Value result = vector::GatherOp::create(
+ rewriter, loc, dstVecTy, tbl, ValueRange{c0}, offsets, mask, passthru);
rewriter.replaceOp(op, result);
return success();
}
@@ -163,11 +220,11 @@ struct ExtFToLUTPattern : public OpRewritePattern<arith::ExtFOp> {
};
namespace {
-struct ConvertArithFP8ExtFToLUTPass
- : public impl::ConvertArithFP8ExtFToLUTBase<ConvertArithFP8ExtFToLUTPass> {
+struct ConvertArithExtFToLUTPass
+ : public impl::ConvertArithExtFToLUTBase<ConvertArithExtFToLUTPass> {
void runOnOperation() override {
- ModuleOp module = getOperation();
+ ModuleOp moduleOp = getOperation();
MLIRContext *ctx = &getContext();
// Cache so each format gets exactly one global, inserted once.
@@ -176,8 +233,7 @@ struct ConvertArithFP8ExtFToLUTPass
RewritePatternSet patterns(ctx);
patterns.add<ExtFToLUTPattern>(ctx, lutCache);
- if (failed(applyPatternsGreedily(module, std::move(patterns))))
- signalPassFailure();
+ walkAndApplyPatterns(moduleOp, std::move(patterns));
}
};
} // namespace
diff --git a/mlir/lib/Conversion/ArithToLUT/CMakeLists.txt b/mlir/lib/Conversion/ArithToLUT/CMakeLists.txt
index e94e4a1835e11..fa4ed1e8344a1 100644
--- a/mlir/lib/Conversion/ArithToLUT/CMakeLists.txt
+++ b/mlir/lib/Conversion/ArithToLUT/CMakeLists.txt
@@ -15,4 +15,5 @@ add_mlir_conversion_library(MLIRArithToLUT
MLIRMemRefDialect
MLIRPass
MLIRTransforms
+ MLIRVectorDialect
)
diff --git a/mlir/test/Conversion/ArithToLUT/extf-to-lut.mlir b/mlir/test/Conversion/ArithToLUT/extf-to-lut.mlir
index 2fce8572df790..a95733083b8b7 100644
--- a/mlir/test/Conversion/ArithToLUT/extf-to-lut.mlir
+++ b/mlir/test/Conversion/ArithToLUT/extf-to-lut.mlir
@@ -1,15 +1,27 @@
-// RUN: mlir-opt %s --convert-arith-fp8-extf-to-lut | FileCheck %s
+// RUN: mlir-opt %s --convert-arith-extf-to-lut | FileCheck %s
+// CHECK-DAG: memref.global "private" constant @__extf_lut_f4E2M1FN : memref<16xf32>
// CHECK-DAG: memref.global "private" constant @__extf_lut_f8E4M3FN : memref<256xf32>
// CHECK-DAG: memref.global "private" constant @__extf_lut_f8E5M2 : memref<256xf32>
+// f4E2M1FN — 16-entry LUT.
+func.func @extf_f4E2M1FN(%a: f4E2M1FN) -> f32 {
+ // CHECK-LABEL: @extf_f4E2M1FN
+ // CHECK: memref.get_global @__extf_lut_f4E2M1FN : memref<16xf32>
+ // CHECK: arith.bitcast {{.*}} : f4E2M1FN to i4
+ // CHECK: arith.index_castui {{.*}} : i4 to index
+ // CHECK: memref.load {{.*}}[{{.*}}] : memref<16xf32>
+ // CHECK-NOT: arith.extf
+ %r = arith.extf %a : f4E2M1FN to f32
+ return %r : f32
+}
+
// Single f8E4M3FN input — table should appear once, op replaced by LUT sequence.
func.func @extf_f8E4M3FN(%a: f8E4M3FN) -> f32 {
// CHECK-LABEL: @extf_f8E4M3FN
// CHECK: memref.get_global @__extf_lut_f8E4M3FN : memref<256xf32>
// CHECK: arith.bitcast {{.*}} : f8E4M3FN to i8
- // CHECK: arith.extui {{.*}} : i8 to i32
- // CHECK: arith.index_cast {{.*}} : i32 to index
+ // CHECK: arith.index_castui {{.*}} : i8 to index
// CHECK: memref.load {{.*}}[{{.*}}] : memref<256xf32>
// CHECK-NOT: arith.extf
%r = arith.extf %a : f8E4M3FN to f32
@@ -34,3 +46,16 @@ func.func @extf_f8E5M2(%a: f8E5M2) -> f32 {
%r = arith.extf %a : f8E5M2 to f32
return %r : f32
}
+
+// Vector f8E4M3FN → f32: LUT looked up via vector.gather.
+func.func @extf_vec_f8E4M3FN(%a: vector<4xf8E4M3FN>) -> vector<4xf32> {
+ // CHECK-LABEL: @extf_vec_f8E4M3FN
+ // CHECK: memref.get_global @__extf_lut_f8E4M3FN : memref<256xf32>
+ // CHECK: arith.bitcast {{.*}} : vector<4xf8E4M3FN> to vector<4xi8>
+ // CHECK: arith.extui {{.*}} : vector<4xi8> to vector<4xi32>
+ // CHECK: vector.constant_mask [4] : vector<4xi1>
+ // CHECK: vector.gather {{.*}} : memref<256xf32>
+ // CHECK-NOT: arith.extf
+ %r = arith.extf %a : vector<4xf8E4M3FN> to vector<4xf32>
+ return %r : vector<4xf32>
+}
diff --git a/mlir/test/Conversion/ArithToLUT/linalg-generic-fp8-addf.mlir b/mlir/test/Conversion/ArithToLUT/linalg-generic-fp8-addf.mlir
deleted file mode 100644
index 200a8262a9669..0000000000000
--- a/mlir/test/Conversion/ArithToLUT/linalg-generic-fp8-addf.mlir
+++ /dev/null
@@ -1,33 +0,0 @@
-// RUN: mlir-opt %s --convert-arith-fp8-extf-to-lut | FileCheck %s
-
-// linalg.generic over tensor<1024xf8E4M3FN>: both extf ops inside the body
-// should be replaced by LUT-lookup sequences; one shared table is emitted.
-
-// CHECK-DAG: memref.global "private" constant @__extf_lut_f8E4M3FN : memref<256xf32>
-
-func.func @linalg_addf_fp8(
- %a: tensor<1024xf8E4M3FN>,
- %b: tensor<1024xf8E4M3FN>) -> tensor<1024xf32> {
- // CHECK-LABEL: @linalg_addf_fp8
- // CHECK: linalg.generic
- // CHECK-COUNT-2: memref.get_global @__extf_lut_f8E4M3FN : memref<256xf32>
- // CHECK: arith.addf {{.*}} : f32
- // CHECK-NOT: arith.extf
- %init = tensor.empty() : tensor<1024xf32>
- %result = linalg.generic {
- indexing_maps = [
- affine_map<(d0) -> (d0)>,
- affine_map<(d0) -> (d0)>,
- affine_map<(d0) -> (d0)>
- ],
- iterator_types = ["parallel"]
- } ins(%a, %b : tensor<1024xf8E4M3FN>, tensor<1024xf8E4M3FN>)
- outs(%init : tensor<1024xf32>) {
- ^bb0(%in_a: f8E4M3FN, %in_b: f8E4M3FN, %out: f32):
- %ra = arith.extf %in_a : f8E4M3FN to f32
- %rb = arith.extf %in_b : f8E4M3FN to f32
- %sum = arith.addf %ra, %rb : f32
- linalg.yield %sum : f32
- } -> tensor<1024xf32>
- return %result : tensor<1024xf32>
-}
diff --git a/mlir/test/Integration/Dialect/Arith/CPU/test-lut-extf-f8.mlir b/mlir/test/Integration/Dialect/Arith/CPU/test-lut-extf-f8.mlir
index 2b89887493114..dfc65374c1ead 100644
--- a/mlir/test/Integration/Dialect/Arith/CPU/test-lut-extf-f8.mlir
+++ b/mlir/test/Integration/Dialect/Arith/CPU/test-lut-extf-f8.mlir
@@ -2,7 +2,7 @@
// value for a selection of well-known bit patterns.
// RUN: mlir-opt %s \
-// RUN: --convert-arith-fp8-extf-to-lut \
+// RUN: --convert-arith-extf-to-lut \
// RUN: --finalize-memref-to-llvm \
// RUN: --convert-arith-to-llvm \
// RUN: --convert-vector-to-llvm \
diff --git a/mlir/unittests/Conversion/ArithToLUT/ArithToLUTTest.cpp b/mlir/unittests/Conversion/ArithToLUT/ArithToLUTTest.cpp
index 884cc5dc2497b..9cc4b4cc175a9 100644
--- a/mlir/unittests/Conversion/ArithToLUT/ArithToLUTTest.cpp
+++ b/mlir/unittests/Conversion/ArithToLUT/ArithToLUTTest.cpp
@@ -6,8 +6,8 @@
//
//===----------------------------------------------------------------------===//
//
-// For every supported f8 format, iterate all 256 bit patterns, run each
-// through the LUT-lowered extf function via the MLIR JIT, and compare the
+// For every supported narrow-float format, iterate all 2^N bit patterns, run
+// each through the LUT-lowered extf function via the MLIR JIT, and compare the
// result bit-for-bit against the APFloat reference value computed the same
// way buildExtFLUT does at compile time.
//
@@ -59,7 +59,7 @@ static struct LLVMInitializer {
static LogicalResult lowerToLLVM(ModuleOp module) {
PassManager pm(module->getName());
- pm.addPass(createConvertArithFP8ExtFToLUT());
+ pm.addPass(createConvertArithExtFToLUT());
pm.addPass(createFinalizeMemRefToLLVMConversionPass());
pm.addNestedPass<func::FuncOp>(createArithToLLVMConversionPass());
pm.addPass(createConvertFuncToLLVMPass());
@@ -68,24 +68,28 @@ static LogicalResult lowerToLLVM(ModuleOp module) {
}
// Builds a module with function:
-// func @test(%arg0: i32) -> f32 { trunci i32 to i8; bitcast i8 to f8; extf f8 to f32 }
-static std::string makeModule(const char *f8TypeName) {
+// func @test(%arg0: i32) -> f32 { trunci i32 to iN; bitcast iN to fN; extf fN
+// to f32 }
+static std::string makeModule(const char *floatTypeName, unsigned bitWidth) {
+ std::string iN = "i" + std::to_string(bitWidth);
std::string s;
s += "func.func @test(%arg0: i32) -> f32 "
"attributes { llvm.emit_c_interface } {\n";
- s += " %i8 = arith.trunci %arg0 : i32 to i8\n";
- s += std::string(" %f8 = arith.bitcast %i8 : i8 to ") + f8TypeName + "\n";
- s += std::string(" %f32 = arith.extf %f8 : ") + f8TypeName + " to f32\n";
+ s += " %" + iN + " = arith.trunci %arg0 : i32 to " + iN + "\n";
+ s += " %fp = arith.bitcast %" + iN + " : " + iN + " to " + floatTypeName +
+ "\n";
+ s += std::string(" %f32 = arith.extf %fp : ") + floatTypeName + " to f32\n";
s += " return %f32 : f32\n}\n";
return s;
}
-struct F8Format {
+struct NarrowFPFormat {
const char *mlirName;
+ unsigned bitWidth;
std::function<FloatType(MLIRContext *)> getType;
};
-static void runExhaustiveTest(const F8Format &fmt) {
+static void runExhaustiveTest(const NarrowFPFormat &fmt) {
DialectRegistry registry;
registerAllDialects(registry);
registerBuiltinDialectTranslation(registry);
@@ -93,7 +97,7 @@ static void runExhaustiveTest(const F8Format &fmt) {
MLIRContext ctx(registry);
OwningOpRef<ModuleOp> module =
- parseSourceString<ModuleOp>(makeModule(fmt.mlirName), &ctx);
+ parseSourceString<ModuleOp>(makeModule(fmt.mlirName, fmt.bitWidth), &ctx);
ASSERT_TRUE(!!module) << "parse failed for " << fmt.mlirName;
ASSERT_TRUE(succeeded(lowerToLLVM(*module)))
<< "lowering failed for " << fmt.mlirName;
@@ -103,73 +107,79 @@ static void runExhaustiveTest(const F8Format &fmt) {
auto jit = std::move(jitOrError.get());
const llvm::fltSemantics &sem = fmt.getType(&ctx).getFloatSemantics();
+ unsigned numEntries = 1u << fmt.bitWidth;
- for (int i = 0; i < 256; ++i) {
+ for (unsigned i = 0; i < numEntries; ++i) {
float got = 0.0f;
- int32_t input = i;
+ int32_t input = static_cast<int32_t>(i);
llvm::Error err =
jit->invoke("test", input, ExecutionEngine::Result<float>(got));
ASSERT_FALSE(err) << llvm::toString(std::move(err));
- llvm::APFloat ref(sem, llvm::APInt(8, static_cast<uint64_t>(i)));
+ llvm::APFloat ref(sem, llvm::APInt(fmt.bitWidth, static_cast<uint64_t>(i)));
bool lossy = false;
- ref.convert(llvm::APFloat::IEEEsingle(),
- llvm::APFloat::rmNearestTiesToEven, &lossy);
+ ref.convert(llvm::APFloat::IEEEsingle(), llvm::APFloat::rmNearestTiesToEven,
+ &lossy);
float expected = ref.convertToFloat();
if (std::isnan(expected)) {
EXPECT_TRUE(std::isnan(got))
- << fmt.mlirName << " pattern " << i
- << ": expected NaN, got " << got;
+ << fmt.mlirName << " pattern " << i << ": expected NaN, got " << got;
} else {
uint32_t gotBits = 0, expBits = 0;
std::memcpy(&gotBits, &got, 4);
std::memcpy(&expBits, &expected, 4);
EXPECT_EQ(gotBits, expBits)
- << fmt.mlirName << " pattern " << i
- << ": expected " << expected << ", got " << got;
+ << fmt.mlirName << " pattern " << i << ": expected " << expected
+ << ", got " << got;
}
}
}
+TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF4E2M1FN)) {
+ runExhaustiveTest({"f4E2M1FN", 4, [](MLIRContext *ctx) -> FloatType {
+ return Float4E2M1FNType::get(ctx);
+ }});
+}
+
TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E4M3FN)) {
- runExhaustiveTest({"f8E4M3FN", [](MLIRContext *ctx) -> FloatType {
+ runExhaustiveTest({"f8E4M3FN", 8, [](MLIRContext *ctx) -> FloatType {
return Float8E4M3FNType::get(ctx);
}});
}
TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E5M2)) {
- runExhaustiveTest({"f8E5M2", [](MLIRContext *ctx) -> FloatType {
+ runExhaustiveTest({"f8E5M2", 8, [](MLIRContext *ctx) -> FloatType {
return Float8E5M2Type::get(ctx);
}});
}
TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E4M3FNUZ)) {
- runExhaustiveTest({"f8E4M3FNUZ", [](MLIRContext *ctx) -> FloatType {
+ runExhaustiveTest({"f8E4M3FNUZ", 8, [](MLIRContext *ctx) -> FloatType {
return Float8E4M3FNUZType::get(ctx);
}});
}
TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E5M2FNUZ)) {
- runExhaustiveTest({"f8E5M2FNUZ", [](MLIRContext *ctx) -> FloatType {
+ runExhaustiveTest({"f8E5M2FNUZ", 8, [](MLIRContext *ctx) -> FloatType {
return Float8E5M2FNUZType::get(ctx);
}});
}
TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E4M3B11FNUZ)) {
- runExhaustiveTest({"f8E4M3B11FNUZ", [](MLIRContext *ctx) -> FloatType {
+ runExhaustiveTest({"f8E4M3B11FNUZ", 8, [](MLIRContext *ctx) -> FloatType {
return Float8E4M3B11FNUZType::get(ctx);
}});
}
TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E3M4)) {
- runExhaustiveTest({"f8E3M4", [](MLIRContext *ctx) -> FloatType {
+ runExhaustiveTest({"f8E3M4", 8, [](MLIRContext *ctx) -> FloatType {
return Float8E3M4Type::get(ctx);
}});
}
TEST(ArithToLUT, SKIP_WITHOUT_JIT(ExtFAllPatternsF8E4M3)) {
- runExhaustiveTest({"f8E4M3", [](MLIRContext *ctx) -> FloatType {
+ runExhaustiveTest({"f8E4M3", 8, [](MLIRContext *ctx) -> FloatType {
return Float8E4M3Type::get(ctx);
}});
}
More information about the Mlir-commits
mailing list