168 lines · cpp
1//===-- MathToXeVM.cpp - conversion from Math to XeVM ---------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Conversion/MathToXeVM/MathToXeVM.h"10#include "mlir/Conversion/ArithCommon/AttrToLLVMConverter.h"11#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"12#include "mlir/Dialect/LLVMIR/LLVMDialect.h"13#include "mlir/Dialect/Math/IR/Math.h"14#include "mlir/IR/BuiltinDialect.h"15#include "mlir/Pass/Pass.h"16#include "llvm/Support/FormatVariadic.h"17 18namespace mlir {19#define GEN_PASS_DEF_CONVERTMATHTOXEVM20#include "mlir/Conversion/Passes.h.inc"21} // namespace mlir22 23using namespace mlir;24 25#define DEBUG_TYPE "math-to-xevm"26 27/// Convert math ops marked with `fast` (`afn`) to native OpenCL intrinsics.28template <typename Op>29struct ConvertNativeFuncPattern final : public OpConversionPattern<Op> {30 31 ConvertNativeFuncPattern(MLIRContext *context, StringRef nativeFunc,32 PatternBenefit benefit = 1)33 : OpConversionPattern<Op>(context, benefit), nativeFunc(nativeFunc) {}34 35 LogicalResult36 matchAndRewrite(Op op, typename Op::Adaptor adaptor,37 ConversionPatternRewriter &rewriter) const override {38 if (!isSPIRVCompatibleFloatOrVec(op.getType()))39 return failure();40 41 arith::FastMathFlags fastFlags = op.getFastmath();42 if (!arith::bitEnumContainsAll(fastFlags, arith::FastMathFlags::afn))43 return rewriter.notifyMatchFailure(op, "not a fastmath `afn` operation");44 45 SmallVector<Type, 1> operandTypes;46 for (auto operand : adaptor.getOperands()) {47 Type opTy = operand.getType();48 // This pass only supports operations on vectors that are already in SPIRV49 // supported vector sizes: Distributing unsupported vector sizes to SPIRV50 // supported vector sizes are done in other blocking optimization passes.51 if (!isSPIRVCompatibleFloatOrVec(opTy))52 return rewriter.notifyMatchFailure(53 op, llvm::formatv("incompatible operand type: '{0}'", opTy));54 operandTypes.push_back(opTy);55 }56 57 auto moduleOp = op->template getParentWithTrait<OpTrait::SymbolTable>();58 auto funcOpRes = LLVM::lookupOrCreateFn(59 rewriter, moduleOp, getMangledNativeFuncName(operandTypes),60 operandTypes, op.getType());61 assert(!failed(funcOpRes));62 LLVM::LLVMFuncOp funcOp = funcOpRes.value();63 64 auto callOp = rewriter.replaceOpWithNewOp<LLVM::CallOp>(65 op, funcOp, adaptor.getOperands());66 // Preserve fastmath flags in our MLIR op when converting to llvm function67 // calls, in order to allow further fastmath optimizations: We thus need to68 // convert arith fastmath attrs into attrs recognized by llvm.69 arith::AttrConvertFastMathToLLVM<Op, LLVM::CallOp> fastAttrConverter(op);70 mlir::NamedAttribute fastAttr = fastAttrConverter.getAttrs()[0];71 callOp->setAttr(fastAttr.getName(), fastAttr.getValue());72 return success();73 }74 75 inline bool isSPIRVCompatibleFloatOrVec(Type type) const {76 if (type.isFloat())77 return true;78 if (auto vecType = dyn_cast<VectorType>(type)) {79 if (!vecType.getElementType().isFloat())80 return false;81 // SPIRV distinguishes between vectors and matrices: OpenCL native math82 // intrsinics are not compatible with matrices.83 ArrayRef<int64_t> shape = vecType.getShape();84 if (shape.size() != 1)85 return false;86 // SPIRV only allows vectors of size 2, 3, 4, 8, 16.87 if (shape[0] == 2 || shape[0] == 3 || shape[0] == 4 || shape[0] == 8 ||88 shape[0] == 16)89 return true;90 }91 return false;92 }93 94 inline std::string95 getMangledNativeFuncName(const ArrayRef<Type> operandTypes) const {96 std::string mangledFuncName =97 "_Z" + std::to_string(nativeFunc.size()) + nativeFunc.str();98 99 auto appendFloatToMangledFunc = [&mangledFuncName](Type type) {100 if (type.isF32())101 mangledFuncName += "f";102 else if (type.isF16())103 mangledFuncName += "Dh";104 else if (type.isF64())105 mangledFuncName += "d";106 };107 108 for (auto type : operandTypes) {109 if (auto vecType = dyn_cast<VectorType>(type)) {110 mangledFuncName += "Dv" + std::to_string(vecType.getShape()[0]) + "_";111 appendFloatToMangledFunc(vecType.getElementType());112 } else113 appendFloatToMangledFunc(type);114 }115 116 return mangledFuncName;117 }118 119 const StringRef nativeFunc;120};121 122void mlir::populateMathToXeVMConversionPatterns(RewritePatternSet &patterns,123 bool convertArith) {124 patterns.add<ConvertNativeFuncPattern<math::ExpOp>>(patterns.getContext(),125 "__spirv_ocl_native_exp");126 patterns.add<ConvertNativeFuncPattern<math::CosOp>>(patterns.getContext(),127 "__spirv_ocl_native_cos");128 patterns.add<ConvertNativeFuncPattern<math::Exp2Op>>(129 patterns.getContext(), "__spirv_ocl_native_exp2");130 patterns.add<ConvertNativeFuncPattern<math::LogOp>>(patterns.getContext(),131 "__spirv_ocl_native_log");132 patterns.add<ConvertNativeFuncPattern<math::Log2Op>>(133 patterns.getContext(), "__spirv_ocl_native_log2");134 patterns.add<ConvertNativeFuncPattern<math::Log10Op>>(135 patterns.getContext(), "__spirv_ocl_native_log10");136 patterns.add<ConvertNativeFuncPattern<math::PowFOp>>(137 patterns.getContext(), "__spirv_ocl_native_powr");138 patterns.add<ConvertNativeFuncPattern<math::RsqrtOp>>(139 patterns.getContext(), "__spirv_ocl_native_rsqrt");140 patterns.add<ConvertNativeFuncPattern<math::SinOp>>(patterns.getContext(),141 "__spirv_ocl_native_sin");142 patterns.add<ConvertNativeFuncPattern<math::SqrtOp>>(143 patterns.getContext(), "__spirv_ocl_native_sqrt");144 patterns.add<ConvertNativeFuncPattern<math::TanOp>>(patterns.getContext(),145 "__spirv_ocl_native_tan");146 if (convertArith)147 patterns.add<ConvertNativeFuncPattern<arith::DivFOp>>(148 patterns.getContext(), "__spirv_ocl_native_divide");149}150 151namespace {152struct ConvertMathToXeVMPass153 : public impl::ConvertMathToXeVMBase<ConvertMathToXeVMPass> {154 using Base::Base;155 void runOnOperation() override;156};157} // namespace158 159void ConvertMathToXeVMPass::runOnOperation() {160 RewritePatternSet patterns(&getContext());161 populateMathToXeVMConversionPatterns(patterns, convertArith);162 ConversionTarget target(getContext());163 target.addLegalDialect<BuiltinDialect, LLVM::LLVMDialect>();164 if (failed(165 applyPartialConversion(getOperation(), target, std::move(patterns))))166 signalPassFailure();167}168