mlir/lib/Conversion/MathToXeVM/MathToXeVM.cpp

//===-- MathToXeVM.cpp - conversion from Math to XeVM ---------------------===//
//
// 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/MathToXeVM/MathToXeVM.h"
#include "mlir/Analysis/DataLayoutAnalysis.h"
#include "mlir/Conversion/ArithAttrToLLVM/AttrToLLVMConverter.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Math/Transforms/Passes.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinDialect.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/Support/FormatVariadic.h"

#include "../GPUCommon/GPUOpsLowering.h"
#include "../GPUCommon/OpToFuncCallLowering.h"

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

using namespace mlir;

#define DEBUG_TYPE "math-to-xevm"

static bool isSizeOneVector(Type type) {
  auto vecType = dyn_cast<VectorType>(type);
  return vecType && vecType.getShape().size() == 1 &&
         vecType.getShape()[0] == 1 && vecType.getElementType().isFloat();
}

static bool isSPIRVCompatibleFloatOrVec(Type type) {
  if (type.isFloat())
    return true;
  if (auto vecType = dyn_cast<VectorType>(type)) {
    if (!vecType.getElementType().isFloat())
      return false;
    // SPIRV distinguishes between vectors and matrices: OpenCL native math
    // intrsinics are not compatible with matrices.
    ArrayRef<int64_t> shape = vecType.getShape();
    if (shape.size() != 1)
      return false;
    // SPIRV has no size-1 vector type; such degenerate vectors are handled
    // by unwrapping to the scalar intrinsic (see matchAndRewrite).
    if (shape[0] == 1)
      return true;
    // SPIRV only allows vectors of size 2, 3, 4, 8, 16.
    if (shape[0] == 2 || shape[0] == 3 || shape[0] == 4 || shape[0] == 8 ||
        shape[0] == 16)
      return true;
  }
  return false;
}

/// Convert math ops marked with `fast` (`afn`) to native OpenCL intrinsics.
template <typename Op>
struct ConvertNativeFuncPattern final : public OpConversionPattern<Op> {

  ConvertNativeFuncPattern(MLIRContext *context, StringRef nativeFunc,
                           PatternBenefit benefit = 1)
      : OpConversionPattern<Op>(context, benefit), nativeFunc(nativeFunc) {}

  inline std::string
  getMangledNativeFuncName(const ArrayRef<Type> operandTypes) const {
    std::string mangledFuncName =
        "_Z" + std::to_string(nativeFunc.size()) + nativeFunc.str();

    auto appendFloatToMangledFunc = [&mangledFuncName](Type type) {
      if (type.isF32())
        mangledFuncName += "f";
      else if (type.isF16())
        mangledFuncName += "Dh";
      else if (type.isF64())
        mangledFuncName += "d";
    };

    for (auto type : operandTypes) {
      if (auto vecType = dyn_cast<VectorType>(type)) {
        mangledFuncName += "Dv" + std::to_string(vecType.getShape()[0]) + "_";
        appendFloatToMangledFunc(vecType.getElementType());
      } else
        appendFloatToMangledFunc(type);
    }

    return mangledFuncName;
  }

  LogicalResult
  matchAndRewrite(Op op, typename Op::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (!isSPIRVCompatibleFloatOrVec(op.getType()))
      return failure();

    arith::FastMathFlags fastFlags = op.getFastmath();
    if (!arith::bitEnumContainsAll(fastFlags, arith::FastMathFlags::afn))
      return rewriter.notifyMatchFailure(op, "not a fastmath `afn` operation");

    Location loc = op.getLoc();

    // SPIRV has no size-1 vector type: such vectors are the degenerate result
    // of distributing/linearizing larger vectors down to a single element (e.g.
    // by the XeGPU lowering pipeline). They have no OpenCL vector intrinsic, so
    // unwrap them to the scalar element type and use the scalar intrinsic.
    SmallVector<Value, 1> operands(adaptor.getOperands());
    SmallVector<Type, 1> operandTypes;
    bool unwrapSizeOneVec = isSizeOneVector(op.getType());
    for (Value &operand : operands) {
      Type opTy = operand.getType();
      // This pass only supports operations on vectors that are already in SPIRV
      // supported vector sizes: Distributing unsupported vector sizes to SPIRV
      // supported vector sizes are done in other blocking optimization passes.
      if (!isSPIRVCompatibleFloatOrVec(opTy))
        return rewriter.notifyMatchFailure(
            op, llvm::formatv("incompatible operand type: '{0}'", opTy));
      if (unwrapSizeOneVec) {
        assert(isSizeOneVector(opTy) &&
               "expected all operands to be size-1 vectors");
        opTy = cast<VectorType>(opTy).getElementType();
        operand = vector::ExtractOp::create(rewriter, loc, operand,
                                            ArrayRef<int64_t>{0});
      }
      operandTypes.push_back(opTy);
    }

    Type resultType = unwrapSizeOneVec
                          ? cast<VectorType>(op.getType()).getElementType()
                          : op.getType();

    auto moduleOp = op->template getParentWithTrait<OpTrait::SymbolTable>();
    auto funcOpRes = LLVM::lookupOrCreateFn(
        rewriter, moduleOp, getMangledNativeFuncName(operandTypes),
        operandTypes, resultType);
    assert(!failed(funcOpRes));
    LLVM::LLVMFuncOp funcOp = funcOpRes.value();

    auto callOp = LLVM::CallOp::create(rewriter, loc, funcOp, operands);
    // Preserve fastmath flags in our MLIR op when converting to llvm function
    // calls, in order to allow further fastmath optimizations: We thus need to
    // convert arith fastmath attrs into attrs recognized by llvm.
    arith::AttrConvertFastMathToLLVM<Op, LLVM::CallOp> fastAttrConverter(op);
    callOp.setFastmathFlagsAttr(
        fastAttrConverter.getProperties().getFastmathFlags());
    callOp->setDiscardableAttrs(fastAttrConverter.getDiscardableAttrs());

    if (unwrapSizeOneVec) {
      // Re-wrap the scalar result back into a size-1 vector to preserve types.
      rewriter.replaceOpWithNewOp<vector::BroadcastOp>(op, op.getType(),
                                                       callOp.getResult());
    } else {
      rewriter.replaceOp(op, callOp);
    }
    return success();
  }

  const StringRef nativeFunc;
};

template <typename OpTy>
static void populateOCLExtSetOpPatterns(const LLVMTypeConverter &converter,
                                        RewritePatternSet &patterns,
                                        PatternBenefit benefit,
                                        StringRef opName) {
  std::string prefix = "__spirv_ocl_";
  std::string mangledName = "_Z" +
                            std::to_string(prefix.size() + opName.size()) +
                            prefix + opName.str();

  patterns.add<ScalarizeVectorOpLowering<OpTy>>(converter, benefit);
  patterns.add<OpToFuncCallLowering<OpTy>>(
      converter, mangledName + "f", mangledName + "d",
      /*f32ApproxFunc=*/"", /*f16Func=*/"",
      /*i32Func=*/"", benefit, LLVM::cconv::CConv::SPIR_FUNC);
}

void mlir::populateMathToScalarOCLExtSetConversionPatterns(
    const LLVMTypeConverter &converter, RewritePatternSet &patterns,
    PatternBenefit benefit) {
  populateOCLExtSetOpPatterns<math::AcosOp>(converter, patterns, benefit,
                                            "acos");
  populateOCLExtSetOpPatterns<math::AcoshOp>(converter, patterns, benefit,
                                             "acosh");
  populateOCLExtSetOpPatterns<math::AsinOp>(converter, patterns, benefit,
                                            "asin");
  populateOCLExtSetOpPatterns<math::AsinhOp>(converter, patterns, benefit,
                                             "asinh");
  populateOCLExtSetOpPatterns<math::AtanOp>(converter, patterns, benefit,
                                            "atan");
  populateOCLExtSetOpPatterns<math::Atan2Op>(converter, patterns, benefit,
                                             "atan2");
  populateOCLExtSetOpPatterns<math::AtanhOp>(converter, patterns, benefit,
                                             "atanh");
  populateOCLExtSetOpPatterns<math::CbrtOp>(converter, patterns, benefit,
                                            "cbrt");
  populateOCLExtSetOpPatterns<math::CopySignOp>(converter, patterns, benefit,
                                                "copysign");
  populateOCLExtSetOpPatterns<math::CosOp>(converter, patterns, benefit, "cos");
  populateOCLExtSetOpPatterns<math::CoshOp>(converter, patterns, benefit,
                                            "cosh");
  populateOCLExtSetOpPatterns<math::ErfOp>(converter, patterns, benefit, "erf");
  populateOCLExtSetOpPatterns<math::ErfcOp>(converter, patterns, benefit,
                                            "erfc");
  populateOCLExtSetOpPatterns<math::ExpOp>(converter, patterns, benefit, "exp");
  populateOCLExtSetOpPatterns<math::Exp2Op>(converter, patterns, benefit,
                                            "exp2");
  populateOCLExtSetOpPatterns<math::ExpM1Op>(converter, patterns, benefit,
                                             "expm1");
  populateOCLExtSetOpPatterns<math::LogOp>(converter, patterns, benefit, "log");
  populateOCLExtSetOpPatterns<math::Log10Op>(converter, patterns, benefit,
                                             "log10");
  populateOCLExtSetOpPatterns<math::Log1pOp>(converter, patterns, benefit,
                                             "log1p");
  populateOCLExtSetOpPatterns<math::Log2Op>(converter, patterns, benefit,
                                            "log2");
  populateOCLExtSetOpPatterns<math::PowFOp>(converter, patterns, benefit,
                                            "pow");
  populateOCLExtSetOpPatterns<math::RsqrtOp>(converter, patterns, benefit,
                                             "rsqrt");
  populateOCLExtSetOpPatterns<math::SinOp>(converter, patterns, benefit, "sin");
  populateOCLExtSetOpPatterns<math::SinhOp>(converter, patterns, benefit,
                                            "sinh");
  populateOCLExtSetOpPatterns<math::SqrtOp>(converter, patterns, benefit,
                                            "sqrt");
  populateOCLExtSetOpPatterns<math::TanOp>(converter, patterns, benefit, "tan");
  populateOCLExtSetOpPatterns<math::TanhOp>(converter, patterns, benefit,
                                            "tanh");
}

void mlir::populateMathToXeVMConversionPatterns(RewritePatternSet &patterns,
                                                bool convertArith,
                                                PatternBenefit benefit) {
  patterns.add<ConvertNativeFuncPattern<math::ExpOp>>(
      patterns.getContext(), "__spirv_ocl_native_exp", benefit);
  patterns.add<ConvertNativeFuncPattern<math::CosOp>>(
      patterns.getContext(), "__spirv_ocl_native_cos", benefit);
  patterns.add<ConvertNativeFuncPattern<math::Exp2Op>>(
      patterns.getContext(), "__spirv_ocl_native_exp2", benefit);
  patterns.add<ConvertNativeFuncPattern<math::LogOp>>(
      patterns.getContext(), "__spirv_ocl_native_log", benefit);
  patterns.add<ConvertNativeFuncPattern<math::Log2Op>>(
      patterns.getContext(), "__spirv_ocl_native_log2", benefit);
  patterns.add<ConvertNativeFuncPattern<math::Log10Op>>(
      patterns.getContext(), "__spirv_ocl_native_log10", benefit);
  patterns.add<ConvertNativeFuncPattern<math::PowFOp>>(
      patterns.getContext(), "__spirv_ocl_native_powr", benefit);
  patterns.add<ConvertNativeFuncPattern<math::RsqrtOp>>(
      patterns.getContext(), "__spirv_ocl_native_rsqrt", benefit);
  patterns.add<ConvertNativeFuncPattern<math::SinOp>>(
      patterns.getContext(), "__spirv_ocl_native_sin", benefit);
  patterns.add<ConvertNativeFuncPattern<math::SqrtOp>>(
      patterns.getContext(), "__spirv_ocl_native_sqrt", benefit);
  patterns.add<ConvertNativeFuncPattern<math::TanOp>>(
      patterns.getContext(), "__spirv_ocl_native_tan", benefit);
  if (convertArith)
    patterns.add<ConvertNativeFuncPattern<arith::DivFOp>>(
        patterns.getContext(), "__spirv_ocl_native_divide", benefit);
}

namespace {
struct ConvertMathToXeVMPass
    : public impl::ConvertMathToXeVMBase<ConvertMathToXeVMPass> {
  using Base::Base;
  void runOnOperation() override;
};
} // namespace

void ConvertMathToXeVMPass::runOnOperation() {
  Operation *op = getOperation();
  MLIRContext *ctx = op->getContext();

  // Simplify first, so the cheaper form is what the conversion below lowers. A
  // simplification rewrites a whole expression, so it must run before the
  // lowering turns the parts of that expression into calls. Only the ops these
  // patterns can match are given to the driver, and folding is off, so nothing
  // else is touched.
  {
    RewritePatternSet simplifications(ctx);
    populateMathAlgebraicSimplificationPatterns(simplifications);
    FrozenRewritePatternSet frozen(std::move(simplifications));
    SmallVector<Operation *> candidates;
    op->walk([&](Operation *nested) {
      if (frozen.getOpSpecificNativePatterns().contains(nested->getName()))
        candidates.push_back(nested);
    });
    GreedyRewriteConfig config;
    config.enableFolding(false);
    if (failed(applyOpPatternsGreedily(candidates, frozen, config)))
      return signalPassFailure();
  }

  const auto &dl = getAnalysis<DataLayoutAnalysis>();

  RewritePatternSet patterns(&getContext());
  LowerToLLVMOptions options(ctx, dl.getAtOrAbove(op));
  LLVMTypeConverter converter(ctx, options);
  ConversionTarget target(getContext());

  // The native (`afn`) patterns must outrank the precise OCL patterns: an op
  // marked `afn` gets the native intrinsic, and every other op falls through to
  // the precise OCL intrinsic.
  constexpr unsigned oclBenefit = 1;
  populateMathToXeVMConversionPatterns(patterns, convertArith, oclBenefit + 1);
  if (convertToOCL) {
    populateMathToScalarOCLExtSetConversionPatterns(converter, patterns,
                                                    oclBenefit);
    target
        .addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::Exp2Op, LLVM::LogOp,
                      LLVM::Log10Op, LLVM::Log2Op, LLVM::SinOp, LLVM::SqrtOp>();
  }
  target.addLegalDialect<BuiltinDialect, LLVM::LLVMDialect>();
  // The size-1-vector patterns unwrap to the scalar intrinsic via
  // vector.extract / vector.broadcast; these must be legal for the partial
  // conversion to succeed.
  target.addLegalOp<vector::ExtractOp, vector::BroadcastOp>();
  if (failed(
          applyPartialConversion(getOperation(), target, std::move(patterns))))
    signalPassFailure();
}