[Codegen][LLVMGPU] Lower FP8 elementwise extensions to NVVM

compiler/src/iree/compiler/Codegen/LLVMGPU/BUILD.bazel

@@ -123,6 +123,7 @@ iree_compiler_cc_library(
         "ConvertToLLVM.cpp",
         "ConvertToNVVM.cpp",
         "ConvertToROCDL.cpp",
+        "FP8Lowering.cpp",
         "KernelConfig.cpp",
         "LLVMGPU1DVectorCanonicalizations.cpp",
         "LLVMGPUAssignConstantOrdinals.cpp",
@@ -152,6 +153,7 @@ iree_compiler_cc_library(
     ],
     hdrs = [
         "ConvertToLLVM.h",
+        "FP8Lowering.h",
         "KernelConfig.h",
         "Passes.h",
         "ROCDLPasses.h",

compiler/src/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt

@@ -89,6 +89,7 @@ iree_cc_library(
     LLVMGPU
   HDRS
     "ConvertToLLVM.h"
+    "FP8Lowering.h"
     "KernelConfig.h"
     "Passes.h"
     "ROCDLPasses.h"
@@ -97,6 +98,7 @@ iree_cc_library(
     "ConvertToLLVM.cpp"
     "ConvertToNVVM.cpp"
     "ConvertToROCDL.cpp"
+    "FP8Lowering.cpp"
     "KernelConfig.cpp"
     "LLVMGPU1DVectorCanonicalizations.cpp"
     "LLVMGPUAssignConstantOrdinals.cpp"

compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToNVVM.cpp

@@ -10,6 +10,7 @@
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
 #include "iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.h"
+#include "iree/compiler/Codegen/LLVMGPU/FP8Lowering.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 #include "mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h"
@@ -163,6 +164,8 @@ struct ConvertToNVVMPass final
       populateFuncToLLVMConversionPatterns(converter, llvmPatterns);
       cf::populateControlFlowToLLVMConversionPatterns(converter, llvmPatterns);
       arith::populateCeilFloorDivExpandOpsPatterns(llvmPatterns);
+      // Lower FP8 extensions from their i8 storage representation.
+      populateFP8ToNVVMConversionPatterns(converter, llvmPatterns);
       arith::populateArithToLLVMConversionPatterns(converter, llvmPatterns);
       vector::populateVectorRankReducingFMAPattern(llvmPatterns);
       vector::populateVectorInsertExtractStridedSliceTransforms(llvmPatterns);

compiler/src/iree/compiler/Codegen/LLVMGPU/FP8Lowering.cpp

@@ -0,0 +1,194 @@
+//
+// Licensed 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
+
+#include "iree/compiler/Codegen/LLVMGPU/FP8Lowering.h"
+
+#include <cstdint>
+
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir::iree_compiler {
+namespace {
+
+struct FP8Format {
+  uint32_t exponentBits;
+  uint32_t mantissaBits;
+  uint32_t exponentBias;
+  bool hasInfinity;
+};
+
+static Value createIntegerConstant(ConversionPatternRewriter &rewriter,
+                                   Location loc, Type type, uint32_t value) {
+  auto elementType = cast<IntegerType>(getElementTypeOrSelf(type));
+  auto elementAttr = rewriter.getIntegerAttr(elementType, value);
+  Attribute attr = elementAttr;
+  if (auto shapedType = dyn_cast<ShapedType>(type)) {
+    attr = SplatElementsAttr::get(shapedType, elementAttr);
+  }
+  return LLVM::ConstantOp::create(rewriter, loc, type, attr);
+}
+
+static uint32_t getSubnormalF32Bits(uint32_t mantissa,
+                                    const FP8Format &format) {
+  if (mantissa == 0) {
+    return 0;
+  }
+
+  uint32_t leadingBit = 0;
+  for (uint32_t value = mantissa; value > 1; value >>= 1u) {
+    ++leadingBit;
+  }
+  int32_t exponent = 1 - static_cast<int32_t>(format.exponentBias) -
+                     static_cast<int32_t>(format.mantissaBits) +
+                     static_cast<int32_t>(leadingBit);
+  uint32_t fraction = (mantissa - (1u << leadingBit)) << (23u - leadingBit);
+  return (static_cast<uint32_t>(exponent + 127) << 23u) | fraction;
+}
+
+static Value createFP8ToF32Bits(ConversionPatternRewriter &rewriter,
+                                Location loc, Value input,
+                                const FP8Format &format) {
+  Type i32Type = rewriter.getI32Type();
+  if (auto vectorType = dyn_cast<VectorType>(input.getType())) {
+    i32Type = vectorType.cloneWith(std::nullopt, i32Type);
+  }
+  Value bits = LLVM::ZExtOp::create(rewriter, loc, i32Type, input);
+
+  Value sign =
+      LLVM::LShrOp::create(rewriter, loc, i32Type, bits,
+                           createIntegerConstant(rewriter, loc, i32Type, 7));
+  sign = LLVM::ShlOp::create(rewriter, loc, i32Type, sign,
+                             createIntegerConstant(rewriter, loc, i32Type, 31));
+
+  Value exponent = LLVM::LShrOp::create(
+      rewriter, loc, i32Type, bits,
+      createIntegerConstant(rewriter, loc, i32Type, format.mantissaBits));
+  uint32_t exponentMask = (1u << format.exponentBits) - 1;
+  exponent = LLVM::AndOp::create(
+      rewriter, loc, i32Type, exponent,
+      createIntegerConstant(rewriter, loc, i32Type, exponentMask));
+
+  uint32_t mantissaMask = (1u << format.mantissaBits) - 1;
+  Value mantissa = LLVM::AndOp::create(
+      rewriter, loc, i32Type, bits,
+      createIntegerConstant(rewriter, loc, i32Type, mantissaMask));
+
+  Value exponentF32 = LLVM::AddOp::create(
+      rewriter, loc, i32Type, exponent,
+      createIntegerConstant(rewriter, loc, i32Type, 127 - format.exponentBias));
+  exponentF32 =
+      LLVM::ShlOp::create(rewriter, loc, i32Type, exponentF32,
+                          createIntegerConstant(rewriter, loc, i32Type, 23));
+  Value mantissaF32 = LLVM::ShlOp::create(
+      rewriter, loc, i32Type, mantissa,
+      createIntegerConstant(rewriter, loc, i32Type, 23 - format.mantissaBits));
+  Value normal = LLVM::OrOp::create(
+      rewriter, loc, i32Type, sign,
+      LLVM::OrOp::create(rewriter, loc, i32Type, exponentF32, mantissaF32));
+
+  // Build the (small) subnormal lookup with selects. This avoids relying on a
+  // target-specific exponent-scaling instruction and works for scalar and
+  // vector values alike.
+  Value subnormal = sign;
+  for (uint32_t value = 1; value <= mantissaMask; ++value) {
+    Value valueBits = LLVM::OrOp::create(
+        rewriter, loc, i32Type, sign,
+        createIntegerConstant(rewriter, loc, i32Type,
+                              getSubnormalF32Bits(value, format)));
+    Value isValue = LLVM::ICmpOp::create(
+        rewriter, loc, LLVM::ICmpPredicate::eq, mantissa,
+        createIntegerConstant(rewriter, loc, i32Type, value));
+    subnormal =
+        LLVM::SelectOp::create(rewriter, loc, isValue, valueBits, subnormal);
+  }
+
+  Value isExponentZero =
+      LLVM::ICmpOp::create(rewriter, loc, LLVM::ICmpPredicate::eq, exponent,
+                           createIntegerConstant(rewriter, loc, i32Type, 0));
+  Value result =
+      LLVM::SelectOp::create(rewriter, loc, isExponentZero, subnormal, normal);
+
+  Value isMaxExponent = LLVM::ICmpOp::create(
+      rewriter, loc, LLVM::ICmpPredicate::eq, exponent,
+      createIntegerConstant(rewriter, loc, i32Type, exponentMask));
+  if (format.hasInfinity) {
+    Value infinity = LLVM::OrOp::create(
+        rewriter, loc, i32Type, sign,
+        createIntegerConstant(rewriter, loc, i32Type, 0x7f800000));
+    Value nan = LLVM::OrOp::create(
+        rewriter, loc, i32Type, infinity,
+        LLVM::ShlOp::create(rewriter, loc, i32Type, mantissa,
+                            createIntegerConstant(rewriter, loc, i32Type,
+                                                  23 - format.mantissaBits)));
+    Value isMantissaZero =
+        LLVM::ICmpOp::create(rewriter, loc, LLVM::ICmpPredicate::eq, mantissa,
+                             createIntegerConstant(rewriter, loc, i32Type, 0));
+    Value special =
+        LLVM::SelectOp::create(rewriter, loc, isMantissaZero, infinity, nan);
+    result =
+        LLVM::SelectOp::create(rewriter, loc, isMaxExponent, special, result);
+  } else {
+    Value isNaN = LLVM::AndOp::create(
+        rewriter, loc, isMaxExponent.getType(), isMaxExponent,
+        LLVM::ICmpOp::create(
+            rewriter, loc, LLVM::ICmpPredicate::eq, mantissa,
+            createIntegerConstant(rewriter, loc, i32Type, mantissaMask)));
+    Value nan = LLVM::OrOp::create(
+        rewriter, loc, i32Type, sign,
+        createIntegerConstant(rewriter, loc, i32Type, 0x7fc00000));
+    result = LLVM::SelectOp::create(rewriter, loc, isNaN, nan, result);
+  }
+
+  return result;
+}
+
+struct LowerFP8ExtFOp final : OpConversionPattern<arith::ExtFOp> {
+  LowerFP8ExtFOp(const LLVMTypeConverter &typeConverter, MLIRContext *context)
+      : OpConversionPattern(typeConverter, context, PatternBenefit(2)) {}
+
+  LogicalResult
+  matchAndRewrite(arith::ExtFOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type sourceType = getElementTypeOrSelf(op.getIn().getType());
+    Type resultType = getElementTypeOrSelf(op.getOut().getType());
+    if (!resultType.isF32()) {
+      return failure();
+    }
+
+    FP8Format format;
+    if (isa<Float8E4M3FNType>(sourceType)) {
+      format = {/*exponentBits=*/4, /*mantissaBits=*/3,
+                /*exponentBias=*/7, /*hasInfinity=*/false};
+    } else if (isa<Float8E5M2Type>(sourceType)) {
+      format = {/*exponentBits=*/5, /*mantissaBits=*/2,
+                /*exponentBias=*/15, /*hasInfinity=*/true};
+    } else {
+      return failure();
+    }
+
+    Type convertedResultType = getTypeConverter()->convertType(op.getType());
+    if (!convertedResultType) {
+      return failure();
+    }
+    Value resultBits =
+        createFP8ToF32Bits(rewriter, op.getLoc(), adaptor.getIn(), format);
+    rewriter.replaceOpWithNewOp<LLVM::BitcastOp>(op, convertedResultType,
+                                                 resultBits);
+    return success();
+  }
+};
+
+} // namespace
+
+void populateFP8ToNVVMConversionPatterns(LLVMTypeConverter &typeConverter,
+                                         RewritePatternSet &patterns) {
+  patterns.add<LowerFP8ExtFOp>(typeConverter, patterns.getContext());
+}
+
+} // namespace mlir::iree_compiler

compiler/src/iree/compiler/Codegen/LLVMGPU/FP8Lowering.h

@@ -0,0 +1,22 @@
+//
+// Licensed 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 IREE_COMPILER_CODEGEN_LLVMGPU_FP8LOWERING_H_
+#define IREE_COMPILER_CODEGEN_LLVMGPU_FP8LOWERING_H_
+
+namespace mlir {
+class LLVMTypeConverter;
+class RewritePatternSet;
+
+namespace iree_compiler {
+
+// Lowers FP8-to-FP32 extensions from i8 storage.
+void populateFP8ToNVVMConversionPatterns(LLVMTypeConverter &typeConverter,
+                                         RewritePatternSet &patterns);
+
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_COMPILER_CODEGEN_LLVMGPU_FP8LOWERING_H_

compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel

@@ -20,6 +20,7 @@ iree_lit_test_suite(
             "config_tile_and_fuse_sm120.mlir",
             "config_tile_and_fuse_sm80.mlir",
             "config_vector_distribute_sm80.mlir",
+            "fp8_elementwise_conversion.mlir",
             "pipeline_full_smoketests.mlir",
             "pipeline_tile_and_fuse_mma_sync.mlir",
             "pipeline_vector_distribute_mma_sync.mlir",

compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt

@@ -18,6 +18,7 @@ iree_lit_test_suite(
     "config_tile_and_fuse_sm80.mlir"
     "config_tile_and_fuse_sm89.mlir"
     "config_vector_distribute_sm80.mlir"
+    "fp8_elementwise_conversion.mlir"
     "pipeline_full_smoketests.mlir"
     "pipeline_tile_and_fuse_mma_sync.mlir"
     "pipeline_vector_distribute_mma_sync.mlir"

compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/fp8_elementwise_conversion.mlir

@@ -0,0 +1,85 @@
+// RUN: iree-opt --iree-gpu-test-target=sm_80 --pass-pipeline="builtin.module(hal.executable(hal.executable.variant(iree-codegen-configuration-preprocessing-pipeline, builtin.module(iree-codegen-llvmgpu-configuration-pipeline, iree-codegen-llvmgpu-nvvm-lowering-pipeline), iree-codegen-translation-postprocessing-pipeline)))" %s | FileCheck %s
+
+// Verify that elementwise FP8-to-FP32 conversion does not leave FP8 values or
+// unrealized casts in the final CUDA LLVM module. This is a general CUDA Core
+// lowering and must not depend on FP8 Tensor Core support.
+
+#executable_target_cuda_nvptx_fb = #hal.executable.target<"cuda", "cuda-nvptx-fb">
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+  #hal.pipeline.binding<storage_buffer, ReadOnly>,
+  #hal.pipeline.binding<storage_buffer>
+]>
+
+hal.executable private @fp8_to_f32 {
+  hal.executable.variant public @cuda_nvptx_fb target(#executable_target_cuda_nvptx_fb) {
+    hal.executable.export public @fp8_to_f32 ordinal(0) layout(#pipeline_layout)
+        count(%device: !hal.device) -> (index, index, index) {
+      %x, %y, %z = iree_tensor_ext.dispatch.workgroup_count_from_slice()
+      hal.return %x, %y, %z : index, index, index
+    }
+    builtin.module {
+      func.func @convert_e4m3fn() {
+        %c0 = arith.constant 0 : index
+        %input = hal.interface.binding.subspan layout(#pipeline_layout) binding(0)
+            alignment(64) offset(%c0) flags(ReadOnly)
+            : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1024xf8E4M3FN>>
+        %output = hal.interface.binding.subspan layout(#pipeline_layout) binding(1)
+            alignment(64) offset(%c0)
+            : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<1024xf32>>
+        %input_tensor = iree_tensor_ext.dispatch.tensor.load %input, offsets = [0], sizes = [1024], strides = [1]
+            : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1024xf8E4M3FN>> -> tensor<1024xf8E4M3FN>
+        %empty = tensor.empty() : tensor<1024xf32>
+        %result = linalg.generic {
+          indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+          iterator_types = ["parallel"]
+        } ins(%input_tensor : tensor<1024xf8E4M3FN>) outs(%empty : tensor<1024xf32>) {
+        ^bb0(%value: f8E4M3FN, %unused: f32):
+          %extended = arith.extf %value : f8E4M3FN to f32
+          linalg.yield %extended : f32
+        } -> tensor<1024xf32>
+        iree_tensor_ext.dispatch.tensor.store %result, %output, offsets = [0], sizes = [1024], strides = [1]
+            : tensor<1024xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<1024xf32>>
+        return
+      }
+
+      func.func @convert_e5m2() {
+        %c0 = arith.constant 0 : index
+        %input = hal.interface.binding.subspan layout(#pipeline_layout) binding(0)
+            alignment(64) offset(%c0) flags(ReadOnly)
+            : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1024xf8E5M2>>
+        %output = hal.interface.binding.subspan layout(#pipeline_layout) binding(1)
+            alignment(64) offset(%c0)
+            : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<1024xf32>>
+        %input_tensor = iree_tensor_ext.dispatch.tensor.load %input, offsets = [0], sizes = [1024], strides = [1]
+            : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1024xf8E5M2>> -> tensor<1024xf8E5M2>
+        %empty = tensor.empty() : tensor<1024xf32>
+        %result = linalg.generic {
+          indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+          iterator_types = ["parallel"]
+        } ins(%input_tensor : tensor<1024xf8E5M2>) outs(%empty : tensor<1024xf32>) {
+        ^bb0(%value: f8E5M2, %unused: f32):
+          %extended = arith.extf %value : f8E5M2 to f32
+          linalg.yield %extended : f32
+        } -> tensor<1024xf32>
+        iree_tensor_ext.dispatch.tensor.store %result, %output, offsets = [0], sizes = [1024], strides = [1]
+            : tensor<1024xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<1024xf32>>
+        return
+      }
+    }
+  }
+}
+
+// CHECK-LABEL: hal.executable private @fp8_to_f32
+// CHECK: builtin.module
+// CHECK-LABEL: llvm.func @convert_e4m3fn
+// CHECK: llvm.zext
+// CHECK: llvm.bitcast
+// CHECK-NOT: builtin.unrealized_conversion_cast
+// CHECK-NOT: f8E4M3FN
+// CHECK-NOT: f8E5M2
+// CHECK-LABEL: llvm.func @convert_e5m2
+// CHECK: llvm.zext
+// CHECK: llvm.bitcast
+// CHECK-NOT: builtin.unrealized_conversion_cast
+// CHECK-NOT: f8E4M3FN
+// CHECK-NOT: f8E5M2