llvm/mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp

//===- ArithToAMDGPU.cpp - Arith to AMDGPU dialect conversion ---------===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "mlir/Conversion/ArithToAMDGPU/ArithToAMDGPU.h"

#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
#include "mlir/Dialect/AMDGPU/Utils/Chipset.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

namespace mlir {
#define GEN_PASS_DEF_ARITHTOAMDGPUCONVERSIONPASS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir

usingnamespacemlir;
usingnamespacemlir::amdgpu;

namespace {
struct ArithToAMDGPUConversionPass final
    : impl::ArithToAMDGPUConversionPassBase<ArithToAMDGPUConversionPass> { … };

struct ExtFOnFloat8RewritePattern final : OpRewritePattern<arith::ExtFOp> { … };

struct TruncFToFloat8RewritePattern final : OpRewritePattern<arith::TruncFOp> { … };

struct TruncfToFloat16RewritePattern final
    : public OpRewritePattern<arith::TruncFOp> { … };

} // end namespace

static Value castF32To(Type elementType, Value f32, Location loc,
                       PatternRewriter &rewriter) { … }

LogicalResult ExtFOnFloat8RewritePattern::match(arith::ExtFOp op) const { … }

void ExtFOnFloat8RewritePattern::rewrite(arith::ExtFOp op,
                                         PatternRewriter &rewriter) const { … }

static Value castToF32(Value value, Location loc, PatternRewriter &rewriter) { … }

// If `in` is a finite value, clamp it between the maximum and minimum values
// of `outElemType` so that subsequent conversion instructions don't
// overflow those out-of-range values to NaN. These semantics are commonly
// used in machine-learning contexts where failure to clamp would lead to
// excessive NaN production.
static Value clampInput(PatternRewriter &rewriter, Location loc,
                        Type outElemType, Value source) { … }

LogicalResult TruncFToFloat8RewritePattern::match(arith::TruncFOp op) const { … }

void TruncFToFloat8RewritePattern::rewrite(arith::TruncFOp op,
                                           PatternRewriter &rewriter) const { … }

LogicalResult TruncfToFloat16RewritePattern::match(arith::TruncFOp op) const { … }

void TruncfToFloat16RewritePattern::rewrite(arith::TruncFOp op,
                                            PatternRewriter &rewriter) const { … }

void mlir::arith::populateArithToAMDGPUConversionPatterns(
    RewritePatternSet &patterns, bool convertFP8Arithmetic,
    bool saturateFP8Truncf, bool allowPackedF16Rtz, Chipset chipset) { … }

void ArithToAMDGPUConversionPass::runOnOperation() { … }