diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMGPU/BUILD.bazel index dc6a9d4f8fa8..56e8b81c474d 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/BUILD.bazel +++ b/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", diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt index a38916616b89..94b01d6d8440 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt +++ b/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" diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToNVVM.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToNVVM.cpp index 6ccd0f665aea..3855d3f08d7b 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToNVVM.cpp +++ b/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); diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/FP8Lowering.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/FP8Lowering.cpp new file mode 100644 index 000000000000..e8c7fed698b8 --- /dev/null +++ b/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 + +#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(getElementTypeOrSelf(type)); + auto elementAttr = rewriter.getIntegerAttr(elementType, value); + Attribute attr = elementAttr; + if (auto shapedType = dyn_cast(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(format.exponentBias) - + static_cast(format.mantissaBits) + + static_cast(leadingBit); + uint32_t fraction = (mantissa - (1u << leadingBit)) << (23u - leadingBit); + return (static_cast(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(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 { + 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(sourceType)) { + format = {/*exponentBits=*/4, /*mantissaBits=*/3, + /*exponentBias=*/7, /*hasInfinity=*/false}; + } else if (isa(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(op, convertedResultType, + resultBits); + return success(); + } +}; + +} // namespace + +void populateFP8ToNVVMConversionPatterns(LLVMTypeConverter &typeConverter, + RewritePatternSet &patterns) { + patterns.add(typeConverter, patterns.getContext()); +} + +} // namespace mlir::iree_compiler diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/FP8Lowering.h b/compiler/src/iree/compiler/Codegen/LLVMGPU/FP8Lowering.h new file mode 100644 index 000000000000..9af4480adb21 --- /dev/null +++ b/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_ diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel index eb5ead95cb9c..72136b3710ce 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel +++ b/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", diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt index efc2d1968038..6bfc9cc8afa2 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt +++ b/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" diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/fp8_elementwise_conversion.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/fp8_elementwise_conversion.mlir new file mode 100644 index 000000000000..7e7e38763307 --- /dev/null +++ b/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, + #hal.pipeline.binding +]> + +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> + %output = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) + alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor> + %input_tensor = iree_tensor_ext.dispatch.tensor.load %input, offsets = [0], sizes = [1024], strides = [1] + : !iree_tensor_ext.dispatch.tensor> -> 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> + 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> + %output = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) + alignment(64) offset(%c0) + : !iree_tensor_ext.dispatch.tensor> + %input_tensor = iree_tensor_ext.dispatch.tensor.load %input, offsets = [0], sizes = [1024], strides = [1] + : !iree_tensor_ext.dispatch.tensor> -> 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> + 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