//===- LibDeviceFuncsRaisingPass.cpp - Raise libdevice.bc func calls ------===//
//
// 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/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "src/enzyme_ad/jax/Dialect/Ops.h"
#include "src/enzyme_ad/jax/Passes/Passes.h"
#include "src/enzyme_ad/jax/Passes/SelectPatterns.h"

#include "Enzyme/EnzymeCallMarkers.h"
#include "Enzyme/MLIR/Dialect/Dialect.h"
#include "Enzyme/MLIR/Dialect/Ops.h"

#include "mlir/Conversion/LLVMCommon/VectorPattern.h"

#include "llvm/ADT/TypeSwitch.h"

namespace mlir {
namespace enzyme {
#define GEN_PASS_DEF_LIBDEVICEFUNCSRAISINGPASS
#include "src/enzyme_ad/jax/Passes/Passes.h.inc"
} // namespace enzyme
} // namespace mlir

using namespace mlir;
using namespace mlir::enzyme;

template <typename SourceOp, typename TargetOp,
          template <typename, typename> typename AttrConvert =
              AttrConvertPassThrough>
class VectorConvertFromLLVMPattern : public OpRewritePattern<SourceOp> {
public:
  using OpRewritePattern<SourceOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(SourceOp op,
                                PatternRewriter &rewriter) const override {
    static_assert(
        std::is_base_of<OpTrait::OneResult<SourceOp>, SourceOp>::value,
        "expected single result op");
    // Determine attributes for the target op
    AttrConvert<SourceOp, TargetOp> attrConvert(op);

    auto operands = op->getOperands();
    auto llvmNDVectorTy = operands[0].getType();
    if (isa<LLVM::LLVMArrayType, mlir::VectorType>(llvmNDVectorTy)) {
      return failure();
    }
    // What comes out counts elements as much as what goes in, and an op free to
    // change one without the other -- a bitcast of an integer to a vector of
    // floats -- cannot be said with one that holds its shape.
    if (llvm::any_of(op->getResultTypes(), [](Type type) {
          return isa<LLVM::LLVMArrayType, mlir::VectorType>(type);
        })) {
      return failure();
    }

    Operation *newOp = rewriter.create(
        op->getLoc(), rewriter.getStringAttr(TargetOp::getOperationName()),
        operands, op->getResultTypes(), attrConvert.getAttrs());

    rewriter.replaceOp(op, newOp->getResult(0));
    return success();
  }
};

arith::IntegerOverflowFlags
convertArithOverflowFlagsFromLLVM(LLVM::IntegerOverflowFlags llvmFlags) {
  arith::IntegerOverflowFlags arithFlags{};
  const std::pair<arith::IntegerOverflowFlags, LLVM::IntegerOverflowFlags>
      flags[] = {
          {arith::IntegerOverflowFlags::nsw, LLVM::IntegerOverflowFlags::nsw},
          {arith::IntegerOverflowFlags::nuw, LLVM::IntegerOverflowFlags::nuw}};
  for (auto [arithFlag, llvmFlag] : flags) {
    if (bitEnumContainsAny(llvmFlags, llvmFlag))
      arithFlags = arithFlags | arithFlag;
  }
  return arithFlags;
}

template <typename SourceOp, typename TargetOp>
class AttrConvertOverflowFromLLVM {
public:
  AttrConvertOverflowFromLLVM(SourceOp srcOp) {
    // Copy the source attributes, minus the flags we are translating. The
    // accessor rather than the attribute dictionary, since the flags are an
    // inherent attribute and so live in the op's properties.
    convertedAttr = NamedAttrList{srcOp->getAttrs()};
    convertedAttr.erase(SourceOp::getOverflowFlagsAttrName());
    LLVM::IntegerOverflowFlags llvmFlags = srcOp.getOverflowFlags();
    if (llvmFlags != LLVM::IntegerOverflowFlags::none)
      convertedAttr.set(TargetOp::getIntegerOverflowAttrName(),
                        arith::IntegerOverflowFlagsAttr::get(
                            srcOp->getContext(),
                            convertArithOverflowFlagsFromLLVM(llvmFlags)));
  }

  ArrayRef<NamedAttribute> getAttrs() const { return convertedAttr.getAttrs(); }

private:
  NamedAttrList convertedAttr;
};

arith::FastMathFlags
convertArithFastMathFlagsFromLLVM(LLVM::FastmathFlags llvmFMF) {
  arith::FastMathFlags arithFMF{};
  const std::pair<arith::FastMathFlags, LLVM::FastmathFlags> flags[] = {
      {arith::FastMathFlags::nnan, LLVM::FastmathFlags::nnan},
      {arith::FastMathFlags::ninf, LLVM::FastmathFlags::ninf},
      {arith::FastMathFlags::nsz, LLVM::FastmathFlags::nsz},
      {arith::FastMathFlags::arcp, LLVM::FastmathFlags::arcp},
      {arith::FastMathFlags::contract, LLVM::FastmathFlags::contract},
      {arith::FastMathFlags::afn, LLVM::FastmathFlags::afn},
      {arith::FastMathFlags::reassoc, LLVM::FastmathFlags::reassoc}};
  for (auto [arithFlag, llvmFlag] : flags) {
    if (bitEnumContainsAny(llvmFMF, llvmFlag))
      arithFMF = arithFMF | arithFlag;
  }
  return arithFMF;
}

arith::FastMathFlagsAttr
convertArithFastMathAttrFromLLVM(LLVM::FastmathFlagsAttr fmfAttr) {
  auto arithFMF = fmfAttr.getValue();
  return arith::FastMathFlagsAttr::get(
      fmfAttr.getContext(), convertArithFastMathFlagsFromLLVM(arithFMF));
}

// Attribute converter that populates a NamedAttrList by removing the fastmath
// attribute from the source operation attributes, and replacing it with an
// equivalent LLVM fastmath attribute.
template <typename SourceOp, typename TargetOp>
class AttrConvertFastMathFromLLVM {
public:
  AttrConvertFastMathFromLLVM(SourceOp srcOp) {
    // Copy the source attributes.
    convertedAttr = NamedAttrList{srcOp->getAttrs()};
    // Get the name of the arith fastmath attribute.
    StringRef arithFMFAttrName = SourceOp::getFastmathAttrName();
    // Remove the source fastmath attribute.
    auto arithFMFAttr = dyn_cast_if_present<LLVM::FastmathFlagsAttr>(
        convertedAttr.erase(arithFMFAttrName));
    if (arithFMFAttr &&
        arithFMFAttr.getValue() != mlir::LLVM::FastmathFlags::none) {
      StringRef targetAttrName = TargetOp::getFastMathAttrName();
      convertedAttr.set(targetAttrName,
                        convertArithFastMathAttrFromLLVM(arithFMFAttr));
    }
  }

  ArrayRef<NamedAttribute> getAttrs() const { return convertedAttr.getAttrs(); }

private:
  NamedAttrList convertedAttr;
};

namespace {
template <typename TargetOp>
class CallToOpRaising : public OpRewritePattern<LLVM::CallOp> {
public:
  CallToOpRaising(MLIRContext *context, StringRef funcNameStr)
      : OpRewritePattern<LLVM::CallOp>(context),
        funcName(StringAttr::get(context, funcNameStr)) {}

  LogicalResult matchAndRewrite(LLVM::CallOp op,
                                PatternRewriter &rewriter) const override {
    CallInterfaceCallable callable = op.getCallableForCallee();
    auto callee = dyn_cast<SymbolRefAttr>(callable);
    if (!callee)
      return failure();

    if (callee.getLeafReference() != funcName)
      return failure();

    rewriter.replaceOpWithNewOp<TargetOp>(op, op->getResults().getTypes(),
                                          op->getOperands());
    return success();
  }

private:
  StringAttr funcName;
};

// The name of the enzyme_* marker global `v` reads, where it reads one. The
// markers reach here as a load of the address of a global, sometimes through a
// cast into another address space.
std::optional<StringRef> getEnzymeMarker(Value v) {
  auto loadOp = dyn_cast_if_present<LLVM::LoadOp>(v.getDefiningOp());
  if (!loadOp)
    return std::nullopt;

  Value address = loadOp.getAddr();
  if (auto castOp =
          dyn_cast_if_present<LLVM::AddrSpaceCastOp>(address.getDefiningOp()))
    address = castOp.getArg();

  auto addressOf =
      dyn_cast_if_present<LLVM::AddressOfOp>(address.getDefiningOp());
  if (!addressOf)
    return std::nullopt;

  StringRef name = addressOf.getGlobalName();
  if (!name.starts_with("enzyme_"))
    return std::nullopt;
  return name;
}

// The activity the shared grammar names, said the way the MLIR ops say it.
enzyme::Activity toMLIRActivity(DIFFE_TYPE activity) {
  switch (activity) {
  case DIFFE_TYPE::CONSTANT:
    return enzyme::Activity::enzyme_const;
  case DIFFE_TYPE::DUP_ARG:
    return enzyme::Activity::enzyme_dup;
  case DIFFE_TYPE::DUP_NONEED:
    return enzyme::Activity::enzyme_dupnoneed;
  case DIFFE_TYPE::OUT_DIFF:
    return enzyme::Activity::enzyme_active;
  }
  llvm_unreachable("unknown marker activity");
}

// Markers that say something about the call rather than about one argument.
struct EnzymeCallFlags {
  bool runtimeActivity = false;
  bool strongZero = false;
};

// Given an LLVM dialect __enzyme_* call with optional enzyme_const/dup/active
// configurations, recover:
//    1) the actual arguments to go into the enzyme.*diff op
//    2) the activities (annotated or inferred) for the arguments and return
//       value(s)
//    3) the markers that configure the call as a whole
void parseEnzymeCall(FunctionOpInterface funcToDiff, ValueRange eoperands,
                     SmallVectorImpl<Value> &arguments,
                     SmallVectorImpl<enzyme::Activity> &argActivities,
                     SmallVectorImpl<enzyme::Activity> &retActivities,
                     EnzymeCallFlags &flags) {
  auto markerAt = [&](unsigned i) { return getEnzymeMarker(eoperands[i]); };

  // Where the shadows are has to be settled before walking, since it decides
  // how every argument is read. The first eoperand is a pointer to funcToDiff,
  // so the arguments start after it.
  enzyme_markers::InterleaveSplit split =
      enzyme_markers::findEnzymeInterleave(1, eoperands.size(), markerAt);
  unsigned shadowIdx = split.shadowStart;

  // An activity marker holds until the next one, so one marker can cover a
  // whole group of arguments. That only applies where the shadows were
  // interleaved: written one by one, an unmarked argument is read off its type
  // as it always was.
  enzyme::Activity sticky = enzyme::Activity::enzyme_dup;

  auto takeArgument = [&](unsigned argIdx, enzyme::Activity activity,
                          unsigned &cursor) {
    argActivities.push_back(activity);
    arguments.push_back(eoperands[argIdx]);
    if (activity != enzyme::Activity::enzyme_dup &&
        activity != enzyme::Activity::enzyme_dupnoneed)
      return;
    if (split.interleaved)
      arguments.push_back(eoperands[shadowIdx++]);
    else
      arguments.push_back(eoperands[++cursor]);
  };

  for (unsigned argIdx = 1; argIdx < split.primalEnd; ++argIdx) {
    Value autodiffArg = eoperands[argIdx];

    if (auto name = markerAt(argIdx)) {
      auto marker = enzyme_markers::lookupEnzymeMarker(*name);
      if (!marker) {
        // Reading an unknown marker as an argument is how a derivative goes
        // quietly wrong, so say so instead.
        funcToDiff.emitError()
            << "unknown enzyme marker '" << *name << "' in a call to it";
        return;
      }

      // Whatever the marker takes for itself is not an argument.
      argIdx += marker->extraOperands;

      switch (marker->role) {
      case enzyme_markers::MarkerRole::Activity: {
        enzyme::Activity activity = toMLIRActivity(*marker->activity);
        sticky = activity;
        if (argIdx + 1 >= split.primalEnd)
          break;
        ++argIdx;
        takeArgument(argIdx, activity, argIdx);
        break;
      }
      case enzyme_markers::MarkerRole::CallFlag:
        if (*name == "enzyme_runtime_activity")
          flags.runtimeActivity = true;
        else if (*name == "enzyme_strong_zero")
          flags.strongZero = true;
        break;
      case enzyme_markers::MarkerRole::Interleave:
      case enzyme_markers::MarkerRole::Modifier:
        break;
      }
      continue;
    }

    if (split.interleaved) {
      takeArgument(argIdx, sticky, argIdx);
      continue;
    }

    // Use default activities based on the types
    arguments.push_back(autodiffArg);
    argActivities.push_back(
        llvm::TypeSwitch<Type, enzyme::Activity>(autodiffArg.getType())
            .Case<FloatType, ComplexType>(
                [](auto type) { return enzyme::Activity::enzyme_active; })
            .Case<LLVM::LLVMPointerType, MemRefType>([&](auto type) {
              // Skip the shadow
              arguments.push_back(eoperands[argIdx + 1]);
              argIdx += 1;
              return enzyme::Activity::enzyme_dup;
            })
            .Default([](Type type) { return enzyme::Activity::enzyme_const; }));
  }

  // Past the shadows there may be more whole-call markers.
  for (unsigned i = shadowIdx; split.interleaved && i < eoperands.size(); ++i) {
    if (auto name = markerAt(i)) {
      if (*name == "enzyme_runtime_activity")
        flags.runtimeActivity = true;
      else if (*name == "enzyme_strong_zero")
        flags.strongZero = true;
    }
  }

  for (const auto &[i, retType] :
       llvm::enumerate(funcToDiff.getResultTypes())) {
    retActivities.push_back(
        llvm::TypeSwitch<Type, enzyme::Activity>(retType)
            .Case<FloatType, ComplexType>(
                [](auto type) { return enzyme::Activity::enzyme_active; })
            .Default([](Type type) { return enzyme::Activity::enzyme_const; }));
  }
}

ArrayAttr getActivityArrayAttr(MLIRContext *ctx,
                               ArrayRef<enzyme::Activity> activities) {
  auto attrs = llvm::map_to_vector<8>(
      activities, [ctx](enzyme::Activity actv) -> Attribute {
        return enzyme::ActivityAttr::get(ctx, actv);
      });
  return ArrayAttr::get(ctx, attrs);
}

// The enzyme entry points are compiler markers, not functions that unwind:
// raised, they become an op with no unwind edge to keep. An invoke of one
// becomes the call and the branch it meant, and the block's other exception
// handling stays as it was.
class EnzymeInvokeToCall : public OpRewritePattern<LLVM::InvokeOp> {
public:
  using OpRewritePattern<LLVM::InvokeOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(LLVM::InvokeOp op,
                                PatternRewriter &rewriter) const override {
    auto callee = dyn_cast<SymbolRefAttr>(op.getCallableForCallee());
    if (!callee)
      return failure();
    if (!callee.getLeafReference().strref().contains("__enzyme_"))
      return failure();
    auto call =
        LLVM::CallOp::create(rewriter, op.getLoc(), op.getResultTypes(),
                             callee.getLeafReference(), op.getCalleeOperands());
    LLVM::BrOp::create(rewriter, op.getLoc(), op.getNormalDestOperands(),
                       op.getNormalDest());
    rewriter.replaceOp(op, call->getResults());
    return success();
  }
};

// The function `op` names, or null where the operand is not the address of one
// this can see.
FunctionOpInterface getEnzymeCallTarget(LLVM::CallOp op, StringRef intrinsic,
                                        FlatSymbolRefAttr &symbol) {
  Operation::operand_range operands = op.getArgOperands();
  if (operands.empty())
    return nullptr;
  auto targetAddr =
      dyn_cast_if_present<LLVM::AddressOfOp>(operands.front().getDefiningOp());
  if (!targetAddr) {
    op.emitError() << "first operand of " << intrinsic
                   << " was not an llvm.mlir.addressof op";
    return nullptr;
  }
  symbol = targetAddr.getGlobalNameAttr();
  auto mod = op->getParentOfType<ModuleOp>();
  return dyn_cast_or_null<FunctionOpInterface>(mod.lookupSymbol(symbol));
}

class EnzymeAutodiffOpRaising : public OpRewritePattern<LLVM::CallOp> {
public:
  using OpRewritePattern<LLVM::CallOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(LLVM::CallOp op,
                                PatternRewriter &rewriter) const override {
    CallInterfaceCallable callable = op.getCallableForCallee();
    auto callee = dyn_cast<SymbolRefAttr>(callable);
    if (!callee)
      return failure();
    if (!callee.getLeafReference().strref().contains("__enzyme_autodiff"))
      return failure();
    Operation::operand_range operands = op.getArgOperands();
    FlatSymbolRefAttr funcToDiffSymbol;
    auto funcToDiff =
        getEnzymeCallTarget(op, "__enzyme_autodiff", funcToDiffSymbol);
    if (!funcToDiff)
      return failure();

    // Infer the argument activities
    SmallVector<Value> arguments;
    SmallVector<enzyme::Activity> argActivities, retActivities;
    EnzymeCallFlags flags;
    argActivities.reserve(funcToDiff.getNumArguments());
    retActivities.reserve(funcToDiff.getNumResults());
    parseEnzymeCall(funcToDiff, operands, arguments, argActivities,
                    retActivities, flags);

    // An active return is differentiated from a seed, and enzyme.autodiff
    // takes one operand for each. __enzyme_autodiff does not name them: the C
    // interface seeds every active return with one, which is what makes its
    // result the gradient rather than a directional derivative. Say the one.
    for (auto [retType, activity] :
         llvm::zip_equal(funcToDiff.getResultTypes(), retActivities)) {
      if (activity != enzyme::Activity::enzyme_active &&
          activity != enzyme::Activity::enzyme_activenoneed)
        continue;
      auto floatType = dyn_cast<FloatType>(retType);
      if (!floatType) {
        op.emitError() << "cannot seed an active return of type " << retType
                       << " in __enzyme_autodiff";
        return failure();
      }
      Value seed =
          LLVM::ConstantOp::create(rewriter, op.getLoc(), floatType,
                                   rewriter.getFloatAttr(floatType, 1.0));
      arguments.push_back(seed);
    }

    MLIRContext *ctx = rewriter.getContext();
    bool atomicAdd = false;
    if (auto llFunc = dyn_cast<LLVM::LLVMFuncOp>(funcToDiff.getOperation())) {
      auto targetFeatures = llFunc.getTargetCpu();
      if (targetFeatures && targetFeatures->starts_with("sm_"))
        atomicAdd = true;
    }
    rewriter.replaceOpWithNewOp<enzyme::AutoDiffOp>(
        op, op.getResultTypes(), funcToDiffSymbol.getValue(), arguments,
        getActivityArrayAttr(ctx, argActivities),
        getActivityArrayAttr(ctx, retActivities),
        /*width=*/1, flags.strongZero, atomicAdd);
    return success();
  }
};

// Forward mode differs from the reverse above in what it says and in what it
// needs said: the shadows come in with the arguments and nothing seeds a
// return, so there is only the op to build.
class EnzymeFwddiffOpRaising : public OpRewritePattern<LLVM::CallOp> {
public:
  using OpRewritePattern<LLVM::CallOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(LLVM::CallOp op,
                                PatternRewriter &rewriter) const override {
    CallInterfaceCallable callable = op.getCallableForCallee();
    auto callee = dyn_cast<SymbolRefAttr>(callable);
    if (!callee)
      return failure();
    if (!callee.getLeafReference().strref().contains("__enzyme_fwddiff"))
      return failure();
    Operation::operand_range operands = op.getArgOperands();
    FlatSymbolRefAttr funcToDiffSymbol;
    auto funcToDiff =
        getEnzymeCallTarget(op, "__enzyme_fwddiff", funcToDiffSymbol);
    if (!funcToDiff)
      return failure();

    SmallVector<Value> arguments;
    SmallVector<enzyme::Activity> argActivities, retActivities;
    EnzymeCallFlags flags;
    argActivities.reserve(funcToDiff.getNumArguments());
    retActivities.reserve(funcToDiff.getNumResults());
    parseEnzymeCall(funcToDiff, operands, arguments, argActivities,
                    retActivities, flags);

    // A return with no shadow to put anywhere is not differentiated.
    for (auto &activity : retActivities)
      if (activity == enzyme::Activity::enzyme_active)
        activity = enzyme::Activity::enzyme_dup;

    if (flags.runtimeActivity)
      op.emitWarning() << "enzyme_runtime_activity is not modelled by the MLIR "
                          "pipeline; the activities given are taken as they "
                          "stand";

    MLIRContext *ctx = rewriter.getContext();
    rewriter.replaceOpWithNewOp<enzyme::ForwardDiffOp>(
        op, op.getResultTypes(), funcToDiffSymbol.getValue(), arguments,
        getActivityArrayAttr(ctx, argActivities),
        getActivityArrayAttr(ctx, retActivities),
        /*width=*/1, flags.strongZero);
    return success();
  }
};

template <typename TargetOp>
class CallToOpIntAdaptRaising : public OpRewritePattern<LLVM::CallOp> {
public:
  CallToOpIntAdaptRaising(MLIRContext *context, StringRef funcNameStr)
      : OpRewritePattern<LLVM::CallOp>(context),
        funcName(StringAttr::get(context, funcNameStr)) {}

  LogicalResult matchAndRewrite(LLVM::CallOp op,
                                PatternRewriter &rewriter) const override {
    CallInterfaceCallable callable = op.getCallableForCallee();
    auto callee = dyn_cast<SymbolRefAttr>(callable);
    if (!callee)
      return failure();

    if (callee.getLeafReference() != funcName)
      return failure();

    auto newOp = TargetOp::create(
        rewriter, op->getLoc(), op->getOperand(0).getType(), op->getOperands());
    auto sourceType = cast<IntegerType>(op->getOperand(0).getType());
    auto targetType = cast<IntegerType>(op->getResultTypes()[0]);
    if (targetType.getWidth() > sourceType.getWidth()) {
      rewriter.replaceOpWithNewOp<arith::ExtUIOp>(op, targetType,
                                                  newOp->getResult(0));
    } else if (targetType.getWidth() < sourceType.getWidth()) {
      rewriter.replaceOpWithNewOp<arith::TruncIOp>(op, targetType,
                                                   newOp->getResult(0));
    } else {
      rewriter.replaceOp(op, newOp);
    }

    return success();
  }

private:
  StringAttr funcName;
};

class IsFPClassRaising : public OpRewritePattern<LLVM::CallOp> {
public:
  IsFPClassRaising(MLIRContext *context)
      : OpRewritePattern<LLVM::CallOp>(context) {}

  LogicalResult matchAndRewrite(LLVM::CallOp op,
                                PatternRewriter &rewriter) const override {
    CallInterfaceCallable callable = op.getCallableForCallee();
    auto callee = dyn_cast<SymbolRefAttr>(callable);
    if (!callee)
      return failure();

    if (callee.getLeafReference() == "__nv_isnand" ||
        callee.getLeafReference() == "__nv_isnan" ||
        callee.getLeafReference() == "__nv_isnanf") {
      rewriter.replaceOpWithNewOp<LLVM::ZExtOp>(
          op, op->getResultTypes(),
          math::IsNaNOp::create(rewriter, op.getLoc(), op->getOperands()[0]));
      return success();
    }

    if (callee.getLeafReference() == "__nv_isfinited" ||
        callee.getLeafReference() == "__nv_isfinite" ||
        callee.getLeafReference() == "__nv_isfinitef") {
      rewriter.replaceOpWithNewOp<LLVM::ZExtOp>(
          op, op->getResultTypes(),
          math::IsFiniteOp::create(rewriter, op.getLoc(),
                                   op->getOperands()[0]));
      return success();
    }

    if (callee.getLeafReference() == "__nv_finited" ||
        callee.getLeafReference() == "__nv_finite" ||
        callee.getLeafReference() == "__nv_finitef") {
      rewriter.replaceOpWithNewOp<LLVM::ZExtOp>(
          op, op->getResultTypes(),
          math::IsFiniteOp::create(rewriter, op.getLoc(),
                                   op->getOperands()[0]));
      return success();
    }

    if (callee.getLeafReference() == "__nv_isinfd" ||
        callee.getLeafReference() == "__nv_isinf" ||
        callee.getLeafReference() == "__nv_isinff") {
      rewriter.replaceOpWithNewOp<LLVM::ZExtOp>(
          op, op->getResultTypes(),
          math::IsFiniteOp::create(rewriter, op.getLoc(),
                                   op->getOperands()[0]));
      return success();
    }

    return failure();
  }
};

class NVVMRcpRaising : public OpRewritePattern<NVVM::RcpApproxFtzF32Op> {
public:
  NVVMRcpRaising(MLIRContext *context)
      : mlir::OpRewritePattern<NVVM::RcpApproxFtzF32Op>(context) {}

  LogicalResult matchAndRewrite(NVVM::RcpApproxFtzF32Op op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    Type type = op.getResult().getType();
    Value one = arith::ConstantOp::create(rewriter, loc, type,
                                          rewriter.getFloatAttr(type, 1.0));
    auto fmfAttr = arith::FastMathFlagsAttr::get(op.getContext(),
                                                 arith::FastMathFlags::afn);
    rewriter.replaceOpWithNewOp<arith::DivFOp>(op, one, op->getOperands()[0],
                                               fmfAttr);
    return success();
  }
};

class RcpRaising : public OpRewritePattern<LLVM::CallOp> {
public:
  RcpRaising(MLIRContext *context) : OpRewritePattern<LLVM::CallOp>(context) {}

  LogicalResult matchAndRewrite(LLVM::CallOp op,
                                PatternRewriter &rewriter) const override {
    CallInterfaceCallable callable = op.getCallableForCallee();
    auto callee = dyn_cast<SymbolRefAttr>(callable);
    if (!callee)
      return failure();

    if (callee.getLeafReference() == "__nv_drcp_rn" ||
        callee.getLeafReference() == "__nv_frcp_rn") {
      Location loc = op.getLoc();
      Type type = op.getResultTypes()[0];
      Value one = arith::ConstantOp::create(rewriter, loc, type,
                                            rewriter.getFloatAttr(type, 1.0));
      rewriter.replaceOpWithNewOp<arith::DivFOp>(op, one, op->getOperands()[0]);
      return success();
    }

    return failure();
  }
};
} // namespace

template <typename TargetOp, typename Arg, typename... Args>
static void populateOpPatterns(MLIRContext *context,
                               RewritePatternSet &patterns, Arg &&arg,
                               Args &&...args) {
  patterns.add<CallToOpRaising<TargetOp>>(context, std::forward<Arg>(arg));
  if constexpr (sizeof...(Args) != 0)
    populateOpPatterns<TargetOp>(context, patterns,
                                 std::forward<Args>(args)...);
}

namespace {

// From
// https://github.com/llvm/llvm-project/blob/7d8b4eb0ead277f41ff69525ed807f9f6e227f37/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp#L31
// except we invert source and target
template <typename SourceOp, typename TargetOp>
using ConvertFastMath = AttrConvertFastMathFromLLVM<SourceOp, TargetOp>;

template <typename SourceOp, typename TargetOp,
          template <typename, typename> typename AttrConvert =
              AttrConvertPassThrough>
using InvVectorConvertFromLLVMPattern =
    VectorConvertFromLLVMPattern<TargetOp, SourceOp, AttrConvert>;

template <typename SourceOp, typename TargetOp>
using ConvertFMFMathFromLLVMPattern =
    VectorConvertFromLLVMPattern<TargetOp, SourceOp, ConvertFastMath>;

using AbsFOpLowering =
    ConvertFMFMathFromLLVMPattern<math::AbsFOp, LLVM::FAbsOp>;
using CeilOpLowering =
    ConvertFMFMathFromLLVMPattern<math::CeilOp, LLVM::FCeilOp>;
using CopySignOpLowering =
    ConvertFMFMathFromLLVMPattern<math::CopySignOp, LLVM::CopySignOp>;
using CosOpLowering = ConvertFMFMathFromLLVMPattern<math::CosOp, LLVM::CosOp>;
using CtPopFOpLowering =
    VectorConvertFromLLVMPattern<LLVM::CtPopOp, math::CtPopOp>;
using Exp2OpLowering =
    ConvertFMFMathFromLLVMPattern<math::Exp2Op, LLVM::Exp2Op>;
using ExpOpLowering = ConvertFMFMathFromLLVMPattern<math::ExpOp, LLVM::ExpOp>;
using FloorOpLowering =
    ConvertFMFMathFromLLVMPattern<math::FloorOp, LLVM::FFloorOp>;
using FmaOpLowering = ConvertFMFMathFromLLVMPattern<math::FmaOp, LLVM::FMAOp>;
using Log10OpLowering =
    ConvertFMFMathFromLLVMPattern<math::Log10Op, LLVM::Log10Op>;
using Log2OpLowering =
    ConvertFMFMathFromLLVMPattern<math::Log2Op, LLVM::Log2Op>;
using LogOpLowering = ConvertFMFMathFromLLVMPattern<math::LogOp, LLVM::LogOp>;
using PowFOpLowering = ConvertFMFMathFromLLVMPattern<math::PowFOp, LLVM::PowOp>;
using FPowIOpLowering =
    ConvertFMFMathFromLLVMPattern<math::FPowIOp, LLVM::PowIOp>;
using RoundEvenOpLowering =
    ConvertFMFMathFromLLVMPattern<math::RoundEvenOp, LLVM::RoundEvenOp>;
using RoundOpLowering =
    ConvertFMFMathFromLLVMPattern<math::RoundOp, LLVM::RoundOp>;
using RintOpLowering =
    ConvertFMFMathFromLLVMPattern<math::RoundEvenOp, LLVM::RintOp>;
using SinOpLowering = ConvertFMFMathFromLLVMPattern<math::SinOp, LLVM::SinOp>;
using SqrtOpLowering =
    ConvertFMFMathFromLLVMPattern<math::SqrtOp, LLVM::SqrtOp>;
using FTruncOpLowering =
    ConvertFMFMathFromLLVMPattern<math::TruncOp, LLVM::FTruncOp>;

using AddFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::AddFOp, LLVM::FAddOp,
                                    AttrConvertFastMathFromLLVM>;
using AddIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::AddIOp, LLVM::AddOp,
                                    AttrConvertOverflowFromLLVM>;
using AndIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::AndIOp, LLVM::AndOp>;
using BitcastOpLowering =
    InvVectorConvertFromLLVMPattern<arith::BitcastOp, LLVM::BitcastOp>;
using DivFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::DivFOp, LLVM::FDivOp,
                                    AttrConvertFastMathFromLLVM>;
using DivSIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::DivSIOp, LLVM::SDivOp>;
using DivUIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::DivUIOp, LLVM::UDivOp>;
using ExtFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::ExtFOp, LLVM::FPExtOp>;
using ExtSIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::ExtSIOp, LLVM::SExtOp>;
using ExtUIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::ExtUIOp, LLVM::ZExtOp>;
using FPToSIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::FPToSIOp, LLVM::FPToSIOp>;
using FPToUIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::FPToUIOp, LLVM::FPToUIOp>;
using MaximumFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::MaximumFOp, LLVM::MaximumOp,
                                    AttrConvertFastMathFromLLVM>;
using MaxNumFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::MaxNumFOp, LLVM::MaxNumOp,
                                    AttrConvertFastMathFromLLVM>;
using MaxSIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::MaxSIOp, LLVM::SMaxOp>;
using MaxUIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::MaxUIOp, LLVM::UMaxOp>;
using MinimumFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::MinimumFOp, LLVM::MinimumOp,
                                    AttrConvertFastMathFromLLVM>;
using MinNumFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::MinNumFOp, LLVM::MinNumOp,
                                    AttrConvertFastMathFromLLVM>;
using MinSIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::MinSIOp, LLVM::SMinOp>;
using MinUIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::MinUIOp, LLVM::UMinOp>;
using MulFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::MulFOp, LLVM::FMulOp,
                                    AttrConvertFastMathFromLLVM>;
using MulIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::MulIOp, LLVM::MulOp,
                                    AttrConvertOverflowFromLLVM>;
using NegFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::NegFOp, LLVM::FNegOp,
                                    AttrConvertFastMathFromLLVM>;
using OrIOpLowering = InvVectorConvertFromLLVMPattern<arith::OrIOp, LLVM::OrOp>;
using RemFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::RemFOp, LLVM::FRemOp,
                                    AttrConvertFastMathFromLLVM>;
using RemSIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::RemSIOp, LLVM::SRemOp>;
using RemUIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::RemUIOp, LLVM::URemOp>;
using SelectOpLowering =
    InvVectorConvertFromLLVMPattern<arith::SelectOp, LLVM::SelectOp>;
using ShLIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::ShLIOp, LLVM::ShlOp,
                                    AttrConvertOverflowFromLLVM>;
using ShRSIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::ShRSIOp, LLVM::AShrOp>;
using ShRUIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::ShRUIOp, LLVM::LShrOp>;
using SIToFPOpLowering =
    InvVectorConvertFromLLVMPattern<arith::SIToFPOp, LLVM::SIToFPOp>;
using SubFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::SubFOp, LLVM::FSubOp,
                                    AttrConvertFastMathFromLLVM>;
using SubIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::SubIOp, LLVM::SubOp,
                                    AttrConvertOverflowFromLLVM>;
// TODO the LLVM -> Arith conversion does
// arith::TruncFOp {rounding_mode=constrained} -> LLVM::ConstrainedFPTruncIntr
// arith::TruncFOp {} -> LLVM::FPTruncIntr
// for now we map both to plain truncfop
using TruncFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::TruncFOp, LLVM::FPTruncOp>;
using ConstrainedTruncFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::TruncFOp,
                                    LLVM::ConstrainedFPTruncIntr>;
using TruncIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::TruncIOp, LLVM::TruncOp>;
using UIToFPOpLowering =
    InvVectorConvertFromLLVMPattern<arith::UIToFPOp, LLVM::UIToFPOp>;
using XOrIOpLowering =
    InvVectorConvertFromLLVMPattern<arith::XOrIOp, LLVM::XOrOp>;

class CmpIOpLowering : public OpRewritePattern<LLVM::ICmpOp> {
public:
  using OpRewritePattern<LLVM::ICmpOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(LLVM::ICmpOp op,
                                PatternRewriter &rewriter) const override {
    if (isa<LLVM::LLVMPointerType>(op.getLhs().getType())) {
      return failure();
    }
    if (auto VT = dyn_cast<mlir::VectorType>(op.getLhs().getType())) {
      if (isa<LLVM::LLVMPointerType>(VT.getElementType())) {
        return failure();
      }
    }

    // Determine attributes for the target op
    AttrConvertPassThrough<LLVM::ICmpOp, arith::CmpIOp> attrConvert(op);

    auto operands = op->getOperands();
    auto llvmNDVectorTy = operands[0].getType();
    if (isa<LLVM::LLVMArrayType, mlir::VectorType>(llvmNDVectorTy)) {
      return failure();
    }

    Operation *newOp = rewriter.create(
        op->getLoc(), rewriter.getStringAttr(arith::CmpIOp::getOperationName()),
        operands, op->getResultTypes(), attrConvert.getAttrs());

    rewriter.replaceOp(op, newOp->getResult(0));
    return success();
  }
};

using CmpFOpLowering =
    InvVectorConvertFromLLVMPattern<arith::CmpFOp, LLVM::FCmpOp,
                                    AttrConvertFastMathFromLLVM>;
using CountLeadingZerosOpLowering =
    InvVectorConvertFromLLVMPattern<math::CountLeadingZerosOp,
                                    LLVM::CountLeadingZerosOp>;
using CountTrailingZerosOpLowering =
    InvVectorConvertFromLLVMPattern<math::CountTrailingZerosOp,
                                    LLVM::CountTrailingZerosOp>;
using CtPopOpLowering =
    InvVectorConvertFromLLVMPattern<math::CtPopOp, LLVM::CtPopOp>;

struct ConstantOpLowering : public OpRewritePattern<LLVM::ConstantOp> {
  using OpRewritePattern<LLVM::ConstantOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(LLVM::ConstantOp op,
                                PatternRewriter &rewriter) const override {
    if (isa<mlir::IntegerType, mlir::FloatType>(op.getResult().getType())) {
      rewriter.replaceOpWithNewOp<arith::ConstantOp>(
          op, op->getResultTypes(), op->getOperands(), op->getAttrs());
      return success();
    }
    return failure();
  }
};

struct RemoveFreeze : public OpRewritePattern<LLVM::FreezeOp> {
  using OpRewritePattern<LLVM::FreezeOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(LLVM::FreezeOp op,
                                PatternRewriter &rewriter) const override {
    rewriter.replaceOp(op, op.getOperand());
    return success();
  }
};

template <typename From, typename To, mlir::gpu::Dimension dim>
struct GPUConvert : public OpRewritePattern<From> {
  using OpRewritePattern<From>::OpRewritePattern;

  LogicalResult matchAndRewrite(From op,
                                PatternRewriter &rewriter) const override {
    auto res = To::create(rewriter, op.getLoc(), dim);
    SmallVector<Operation *> toReplace;
    for (auto u : op->getUsers()) {
      if (auto ext = dyn_cast<arith::ExtUIOp>(u)) {
        toReplace.push_back(ext);
        continue;
      }
      if (auto ext = dyn_cast<arith::ExtSIOp>(u)) {
        toReplace.push_back(ext);
        continue;
      }
    }
    for (auto e : toReplace)
      rewriter.replaceOpWithNewOp<arith::IndexCastUIOp>(
          e, e->getResultTypes()[0], res);
    rewriter.replaceOpWithNewOp<arith::IndexCastUIOp>(op, op.getType(), res);
    return success();
  }
};

struct BarrierConvert : public OpRewritePattern<LLVM::CallIntrinsicOp> {
  using OpRewritePattern<LLVM::CallIntrinsicOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(LLVM::CallIntrinsicOp op,
                                PatternRewriter &rewriter) const override {
    if (op.getIntrin() != "llvm.nvvm.barrier.cta.sync.aligned.all")
      return failure();

    if (!matchPattern(op.getArgs()[0], m_Zero()))
      return failure();

    rewriter.replaceOpWithNewOp<gpu::BarrierOp>(op);
    return success();
  }
};

struct NVVMRsqrtApproxRaising : public OpRewritePattern<LLVM::CallIntrinsicOp> {
  using OpRewritePattern<LLVM::CallIntrinsicOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(LLVM::CallIntrinsicOp op,
                                PatternRewriter &rewriter) const override {
    if (op.getIntrin() != "llvm.nvvm.rsqrt.approx.f" &&
        op.getIntrin() != "llvm.nvvm.rsqrt.approx.d")
      return failure();

    auto fmfAttr = arith::FastMathFlagsAttr::get(op.getContext(),
                                                 arith::FastMathFlags::afn);
    rewriter.replaceOp(op, math::RsqrtOp::create(rewriter, op.getLoc(),
                                                 op.getArgs()[0], fmfAttr));
    return success();
  }
};

struct ReadOnlyAllocaElim : public OpRewritePattern<LLVM::AllocaOp> {
  ReadOnlyAllocaElim(MLIRContext *context)
      : OpRewritePattern<LLVM::AllocaOp>(context, /*benefit=*/1) {}

  LogicalResult matchAndRewrite(LLVM::AllocaOp alloca,
                                PatternRewriter &rewriter) const override {
    Value ptr = alloca.getResult();
    SmallVector<Operation *> deadUsers;

    // Check all users of the alloca
    for (Operation *user : ptr.getUsers()) {
      // Allow lifetime markers
      if (isa<LLVM::LifetimeStartOp, LLVM::LifetimeEndOp>(user)) {
        deadUsers.push_back(user);
        continue;
      }

      if (auto memcpy = dyn_cast<LLVM::MemcpyOp>(user)) {
        // If stores into allocation, keep it
        if (memcpy.getDst() == ptr)
          return failure();

        deadUsers.push_back(user);
      } else {
        // Found non-read/lifetime user
        return failure();
      }
    }

    for (Operation *user : llvm::reverse(deadUsers)) {
      rewriter.eraseOp(user);
    }
    rewriter.eraseOp(alloca);

    return success();
  }
};

class HalfMathRaising : public OpRewritePattern<LLVM::CallOp> {
public:
  HalfMathRaising(MLIRContext *context)
      : OpRewritePattern<LLVM::CallOp>(context) {}

  enum class MathOp { Log, Div, Abs, Mul, Add, Sqrt, Sub, Exp, Neg, Lt, None };

  LogicalResult matchAndRewrite(LLVM::CallOp op,
                                PatternRewriter &rewriter) const override {
    CallInterfaceCallable callable = op.getCallableForCallee();
    auto callee = dyn_cast<SymbolRefAttr>(callable);
    if (!callee)
      return failure();

    MathOp mathOp =
        llvm::StringSwitch<MathOp>(callee.getLeafReference().getValue())
            .Case("_ZL4hlog6__half", MathOp::Log)
            .Case("_ZL6__hdiv6__halfS_", MathOp::Div)
            .Case("_ZL22__internal_device_hdiv13__nv_bfloat16S_", MathOp::Div)
            .Case("_ZL6__habs6__half", MathOp::Abs)
            .Case("_ZL6__hmul6__halfS_", MathOp::Mul)
            .Case("_ZL27__internal_sm80_device_hmul13__nv_bfloat16S_",
                  MathOp::Mul)
            .Case("_ZL6__hadd6__halfS_", MathOp::Add)
            .Case("_ZL6__hadd13__nv_bfloat16S_", MathOp::Add)
            .Case("_ZL5hsqrt6__half", MathOp::Sqrt)
            .Case("_ZL6__hsub6__halfS_", MathOp::Sub)
            .Case("_ZL27__internal_sm80_device_hsub13__nv_bfloat16S_",
                  MathOp::Sub)
            .Case("_ZL4hexp6__half", MathOp::Exp)
            .Case("_ZL6__hneg6__half", MathOp::Neg)
            .Case("_ZL22__internal_device_hneg13__nv_bfloat16", MathOp::Neg)
            .Case("_ZL5__hlt6__halfS_", MathOp::Lt)
            .Default(MathOp::None);

    if (mathOp == MathOp::None)
      return failure();

    Location loc = op.getLoc();
    SmallVector<Value> newArgs;
    Type fltType = callee.getLeafReference().getValue().contains("bfloat16")
                       ? rewriter.getBF16Type()
                       : rewriter.getF16Type();

    for (Value arg : op.getOperands()) {
      if (isa<LLVM::LLVMPointerType>(arg.getType())) {
        newArgs.push_back(LLVM::LoadOp::create(rewriter, loc, fltType, arg));
      } else if (isa<LLVM::LLVMStructType>(arg.getType())) {
        Value val = LLVM::ExtractValueOp::create(rewriter, loc, arg, 0);
        newArgs.push_back(
            arith::BitcastOp::create(rewriter, loc, fltType, val));
      } else if (isa<IntegerType>(arg.getType())) {
        newArgs.push_back(
            arith::BitcastOp::create(rewriter, loc, fltType, arg));
      } else {
        newArgs.push_back(arg);
      }
    }

    Value res;
    switch (mathOp) {
    case MathOp::Log:
      res = math::LogOp::create(rewriter, loc, newArgs[0]);
      break;
    case MathOp::Div:
      res = arith::DivFOp::create(rewriter, loc, newArgs[0], newArgs[1]);
      break;
    case MathOp::Abs:
      res = math::AbsFOp::create(rewriter, loc, newArgs[0]);
      break;
    case MathOp::Mul:
      res = arith::MulFOp::create(rewriter, loc, newArgs[0], newArgs[1]);
      break;
    case MathOp::Add:
      res = arith::AddFOp::create(rewriter, loc, newArgs[0], newArgs[1]);
      break;
    case MathOp::Sqrt:
      res = math::SqrtOp::create(rewriter, loc, newArgs[0]);
      break;
    case MathOp::Sub:
      res = arith::SubFOp::create(rewriter, loc, newArgs[0], newArgs[1]);
      break;
    case MathOp::Exp:
      res = math::ExpOp::create(rewriter, loc, newArgs[0]);
      break;
    case MathOp::Neg:
      res = arith::NegFOp::create(rewriter, loc, newArgs[0]);
      break;
    case MathOp::Lt:
      res = arith::CmpFOp::create(rewriter, loc, arith::CmpFPredicate::OLT,
                                  newArgs[0], newArgs[1]);
      break;
    case MathOp::None:
      llvm_unreachable("Invalid math function");
    }

    Type resType = op.getResultTypes()[0];
    if (mathOp == MathOp::Lt) {
      if (isa<IntegerType>(resType) && resType.getIntOrFloatBitWidth() > 1) {
        res = arith::ExtUIOp::create(rewriter, loc, resType, res);
      }
    } else {
      if (isa<IntegerType>(resType)) {
        res = arith::BitcastOp::create(rewriter, loc, resType, res);
      } else if (auto structTy = dyn_cast<LLVM::LLVMStructType>(resType)) {
        res =
            arith::BitcastOp::create(rewriter, loc, rewriter.getI16Type(), res);
        res = LLVM::InsertValueOp::create(
            rewriter, loc, structTy,
            LLVM::UndefOp::create(rewriter, loc, structTy), res,
            rewriter.getDenseI64ArrayAttr(0));
      } else if (!isa<FloatType>(resType)) {
        res = arith::BitcastOp::create(rewriter, loc, resType, res);
      }
    }

    rewriter.replaceOp(op, res);
    return success();
  }
};

class BF16HalfToFloatRaising : public OpRewritePattern<LLVM::CallOp> {
public:
  BF16HalfToFloatRaising(MLIRContext *context)
      : OpRewritePattern<LLVM::CallOp>(context) {}

  LogicalResult matchAndRewrite(LLVM::CallOp op,
                                PatternRewriter &rewriter) const override {
    CallInterfaceCallable callable = op.getCallableForCallee();
    auto callee = dyn_cast<SymbolRefAttr>(callable);
    if (!callee)
      return failure();

    StringRef funcName = callee.getLeafReference();
    if (funcName == "__half2float" || funcName == "_ZL12__half2float6__half") {
      Value input = op.getOperand(0);
      Location loc = op.getLoc();
      if (isa<LLVM::LLVMPointerType>(input.getType())) {
        input =
            LLVM::LoadOp::create(rewriter, loc, rewriter.getF16Type(), input);
      } else if (isa<LLVM::LLVMStructType>(input.getType())) {
        input = LLVM::ExtractValueOp::create(rewriter, loc, input, 0);
      }
      if (isa<IntegerType>(input.getType())) {
        input = arith::BitcastOp::create(rewriter, loc, rewriter.getF16Type(),
                                         input);
      }
      rewriter.replaceOpWithNewOp<arith::ExtFOp>(op, op.getResultTypes()[0],
                                                 input);
      return success();
    }
    if (funcName == "__bfloat162float" ||
        funcName == "_ZL16__bfloat162float13__nv_bfloat16" ||
        funcName == "_ZL25__internal_bfloat162floatt" ||
        funcName == "_ZL32__internal_device_bfloat162floatt") {
      Value input = op.getOperand(0);
      Location loc = op.getLoc();
      if (isa<LLVM::LLVMPointerType>(input.getType())) {
        input =
            LLVM::LoadOp::create(rewriter, loc, rewriter.getBF16Type(), input);
      } else if (isa<LLVM::LLVMStructType>(input.getType())) {
        input = LLVM::ExtractValueOp::create(rewriter, loc, input, 0);
      }
      if (isa<IntegerType>(input.getType())) {
        input = arith::BitcastOp::create(rewriter, loc, rewriter.getBF16Type(),
                                         input);
      }
      rewriter.replaceOpWithNewOp<arith::ExtFOp>(op, op.getResultTypes()[0],
                                                 input);
      return success();
    }
    if (funcName == "__float2bfloat16" ||
        funcName == "_ZL16__float2bfloat16f") {
      Value input = op.getOperand(0);
      Location loc = op.getLoc();
      Type resType = op.getResultTypes()[0];
      Value res =
          arith::TruncFOp::create(rewriter, loc, rewriter.getBF16Type(), input);
      if (isa<IntegerType>(resType)) {
        res = arith::BitcastOp::create(rewriter, loc, resType, res);
      } else if (auto structTy = dyn_cast<LLVM::LLVMStructType>(resType)) {
        res =
            arith::BitcastOp::create(rewriter, loc, rewriter.getI16Type(), res);
        res = LLVM::InsertValueOp::create(
            rewriter, loc, structTy,
            LLVM::UndefOp::create(rewriter, loc, structTy), res,
            rewriter.getDenseI64ArrayAttr(0));
      }
      rewriter.replaceOp(op, res);
      return success();
    }
    if (funcName == "_ZL12__float2halff" ||
        funcName == "_ZL21__internal_float2halffRjS_") {
      Value input = op.getOperand(0);
      Location loc = op.getLoc();
      Type resType = op.getResultTypes()[0];
      Value res =
          arith::TruncFOp::create(rewriter, loc, rewriter.getF16Type(), input);
      if (isa<IntegerType>(resType)) {
        res = arith::BitcastOp::create(rewriter, loc, resType, res);
      } else if (auto structTy = dyn_cast<LLVM::LLVMStructType>(resType)) {
        res =
            arith::BitcastOp::create(rewriter, loc, rewriter.getI16Type(), res);
        res = LLVM::InsertValueOp::create(
            rewriter, loc, structTy,
            LLVM::UndefOp::create(rewriter, loc, structTy), res,
            rewriter.getDenseI64ArrayAttr(0));
      }
      rewriter.replaceOp(op, res);
      return success();
    }

    return failure();
  }
};

class InlineAsmHalfRaising : public OpRewritePattern<LLVM::InlineAsmOp> {
public:
  InlineAsmHalfRaising(MLIRContext *context)
      : OpRewritePattern<LLVM::InlineAsmOp>(context) {}

  LogicalResult matchAndRewrite(LLVM::InlineAsmOp op,
                                PatternRewriter &rewriter) const override {
    StringRef asmStr = op.getAsmString();
    if (asmStr == "{  cvt.rn.f16.f32 $0, $1;}\n") {
      Value input = op.getOperand(0);
      Location loc = op.getLoc();
      Type resType = op.getResultTypes()[0];
      Value res =
          arith::TruncFOp::create(rewriter, loc, rewriter.getF16Type(), input);
      if (isa<IntegerType>(resType)) {
        res = arith::BitcastOp::create(rewriter, loc, resType, res);
      }
      rewriter.replaceOp(op, res);
      return success();
    }
    if (asmStr == "{  cvt.f32.f16 $0, $1;}\n") {
      Value input = op.getOperand(0);
      Location loc = op.getLoc();
      if (isa<IntegerType>(input.getType())) {
        input = arith::BitcastOp::create(rewriter, loc, rewriter.getF16Type(),
                                         input);
      }
      Value res =
          arith::ExtFOp::create(rewriter, loc, rewriter.getF32Type(), input);
      rewriter.replaceOp(op, res);
      return success();
    }
    return failure();
  }
};

class FMulAddRaising : public RewritePattern {
public:
  FMulAddRaising(MLIRContext *context)
      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}

  LogicalResult matchAndRewrite(Operation *op,
                                PatternRewriter &rewriter) const override {
    if (op->getName().getStringRef() != "llvm.intr.fmuladd")
      return failure();

    if (op->getNumOperands() != 3 || op->getNumResults() != 1)
      return failure();

    Value a = op->getOperand(0);
    Value b = op->getOperand(1);
    Value c = op->getOperand(2);

    rewriter.replaceOpWithNewOp<math::FmaOp>(op, op->getResultTypes()[0], a, b,
                                             c);
    return success();
  }
};

} // namespace

void mlir::enzyme::populateLibDeviceFuncsToOpsPatterns(
    MLIRContext *context, RewritePatternSet &patterns) {
  // XXX: Keep in sync with
  // mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp.

  auto *converter = context;

  patterns.add<IsFPClassRaising>(context);
  patterns.add<RcpRaising>(context);
  patterns.add<NVVMRcpRaising>(context);
  patterns.add<BF16HalfToFloatRaising>(context);
  patterns.add<HalfMathRaising>(context);
  patterns.add<InlineAsmHalfRaising>(context);
  patterns.add<CallToOpIntAdaptRaising<math::CountLeadingZerosOp>>(context,
                                                                   "__nv_clz");
  patterns.add<CallToOpIntAdaptRaising<math::CountLeadingZerosOp>>(
      context, "__nv_clzll");
  patterns.add<CallToOpIntAdaptRaising<math::CtPopOp>>(context, "__nv_popc");
  patterns.add<CallToOpIntAdaptRaising<math::CtPopOp>>(context, "__nv_popcll");

  populateOpPatterns<arith::RemFOp>(converter, patterns, "__nv_fmodf",
                                    "__nv_fmod", "fmodf", "fmod");
  populateOpPatterns<math::AbsFOp>(converter, patterns, "__nv_fabsf",
                                   "__nv_fabs", "fabsf", "fabs");
  populateOpPatterns<math::AcosOp>(converter, patterns, "__nv_acosf",
                                   "__nv_acos", "acosf", "acos");
  populateOpPatterns<math::AcoshOp>(converter, patterns, "__nv_acoshf",
                                    "__nv_acosh", "acoshf", "acosh");
  populateOpPatterns<math::AsinOp>(converter, patterns, "__nv_asinf",
                                   "__nv_asin", "asinf", "asin");
  populateOpPatterns<math::AsinhOp>(converter, patterns, "__nv_asinhf",
                                    "__nv_asinh", "asinhf", "asinh");
  populateOpPatterns<math::AtanOp>(converter, patterns, "__nv_atanf",
                                   "__nv_atan", "atanf", "atan");
  populateOpPatterns<math::Atan2Op>(converter, patterns, "__nv_atan2f",
                                    "__nv_atan2", "atan2f", "atan2");
  populateOpPatterns<math::AtanhOp>(converter, patterns, "__nv_atanhf",
                                    "__nv_atanh", "atanhf", "atanh");
  populateOpPatterns<math::CbrtOp>(converter, patterns, "__nv_cbrtf",
                                   "__nv_cbrt", "cbrtf", "cbrt");
  populateOpPatterns<math::CeilOp>(converter, patterns, "__nv_ceilf",
                                   "__nv_ceil", "ceilf", "ceil");
  populateOpPatterns<math::CopySignOp>(converter, patterns, "__nv_copysignf",
                                       "__nv_copysign", "copysignf",
                                       "copysign");
  populateOpPatterns<math::CosOp>(converter, patterns, "__nv_cosf", "__nv_cos",
                                  "__nv_fast_cosf", "cosf", "cos");
  populateOpPatterns<math::CoshOp>(converter, patterns, "__nv_coshf",
                                   "__nv_cosh", "coshf", "cosh");
  populateOpPatterns<math::ErfOp>(converter, patterns, "__nv_erff", "__nv_erf",
                                  "erff", "erf");
  populateOpPatterns<math::ExpOp>(converter, patterns, "__nv_expf", "__nv_exp",
                                  "__nv_fast_expf", "expf", "exp");
  populateOpPatterns<math::Exp2Op>(converter, patterns, "__nv_exp2f",
                                   "__nv_exp2", "exp2f", "exp2");
  populateOpPatterns<math::ExpM1Op>(converter, patterns, "__nv_expm1f",
                                    "__nv_expm1", "expm1f", "expm1");
  populateOpPatterns<math::FloorOp>(converter, patterns, "__nv_floorf",
                                    "__nv_floor", "floorf", "floor");
  populateOpPatterns<math::FmaOp>(converter, patterns, "__nv_fmaf", "__nv_fma",
                                  "fmaf", "fma");
  populateOpPatterns<math::LogOp>(converter, patterns, "__nv_logf", "__nv_log",
                                  "__nv_fast_logf", "logf", "log");
  populateOpPatterns<math::Log10Op>(converter, patterns, "__nv_log10f",
                                    "__nv_log10", "__nv_fast_log10f", "log10f",
                                    "log10");
  populateOpPatterns<math::Log1pOp>(converter, patterns, "__nv_log1pf",
                                    "__nv_log1p", "log1pf", "log1p");
  populateOpPatterns<math::Log2Op>(converter, patterns, "__nv_log2f",
                                   "__nv_log2", "__nv_fast_log2f", "log2f",
                                   "log2");
  populateOpPatterns<math::PowFOp>(converter, patterns, "__nv_powf", "__nv_pow",
                                   "__nv_fast_powf", "powf", "pow");
  populateOpPatterns<arith::DivFOp>(converter, patterns, "__nv_fdividef",
                                    "__nv_fdivide", "__nv_fast_fdividef");
  populateOpPatterns<math::RoundOp>(converter, patterns, "__nv_roundf",
                                    "__nv_round", "roundf", "round");
  populateOpPatterns<math::RoundEvenOp>(converter, patterns, "__nv_rintf",
                                        "__nv_rint", "rintf", "rint");
  populateOpPatterns<math::RsqrtOp>(converter, patterns, "__nv_rsqrtf",
                                    "__nv_rsqrt");
  populateOpPatterns<math::SinOp>(converter, patterns, "__nv_sinf", "__nv_sin",
                                  "__nv_fast_sinf", "sinf", "sin");
  populateOpPatterns<math::SinhOp>(converter, patterns, "__nv_sinhf",
                                   "__nv_sinh", "sinhf", "sinh");
  populateOpPatterns<math::SqrtOp>(converter, patterns, "__nv_sqrtf",
                                   "__nv_sqrt", "sqrtf", "sqrt");
  populateOpPatterns<math::TanOp>(converter, patterns, "__nv_tanf", "__nv_tan",
                                  "__nv_fast_tanf", "tanf", "tan");
  populateOpPatterns<math::TanhOp>(converter, patterns, "__nv_tanhf",
                                   "__nv_tanh", "tanhf", "tanh");
  populateOpPatterns<math::FPowIOp>(converter, patterns, "__nv_powif",
                                    "__nv_powi");
  populateOpPatterns<math::AbsIOp>(converter, patterns, "__nv_abs",
                                   "__nv_llabs");
  populateOpPatterns<arith::MaxNumFOp>(converter, patterns, "__nv_fmax",
                                       "__nv_fmaxf");
  populateOpPatterns<arith::MinNumFOp>(converter, patterns, "__nv_fmin",
                                       "__nv_fminf");
  populateOpPatterns<math::TruncOp>(converter, patterns, "__nv_trunc",
                                    "__nv_truncf");
  populateOpPatterns<enzymexla::TGammaOp>(converter, patterns, "__nv_tgamma",
                                          "__nv_tgammaf");
  populateOpPatterns<enzymexla::LGammaOp>(converter, patterns, "__nv_lgamma",
                                          "__nv_lgammaf");
}

void populateLLVMToMathPatterns(MLIRContext *context,
                                RewritePatternSet &patterns) {
  auto *converter = context;
  // From
  // https://github.com/llvm/llvm-project/blob/7d8b4eb0ead277f41ff69525ed807f9f6e227f37/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp#L306
  // patterns.add<FTruncOpLowering>(converter);
  patterns.add<AbsFOpLowering,
               // AbsIOpLowering,
               CeilOpLowering, CopySignOpLowering, CosOpLowering,
               CountLeadingZerosOpLowering, CountTrailingZerosOpLowering,
               CtPopFOpLowering, Exp2OpLowering,
               // ExpM1OpLowering,
               ExpOpLowering, FPowIOpLowering, FloorOpLowering, FmaOpLowering,
               Log10OpLowering, Log2OpLowering, LogOpLowering, PowFOpLowering,
               RoundEvenOpLowering, RoundOpLowering, RintOpLowering,
               // RsqrtOpLowering,
               SinOpLowering, SqrtOpLowering, FTruncOpLowering>(converter);
  patterns.add<FMulAddRaising>(converter);

  patterns
      .add<GPUConvert<NVVM::ThreadIdXOp, gpu::ThreadIdOp, gpu::Dimension::x>>(
          converter);
  patterns
      .add<GPUConvert<NVVM::ThreadIdYOp, gpu::ThreadIdOp, gpu::Dimension::y>>(
          converter);
  patterns
      .add<GPUConvert<NVVM::ThreadIdZOp, gpu::ThreadIdOp, gpu::Dimension::z>>(
          converter);

  patterns.add<GPUConvert<NVVM::BlockIdXOp, gpu::BlockIdOp, gpu::Dimension::x>>(
      converter);
  patterns.add<GPUConvert<NVVM::BlockIdYOp, gpu::BlockIdOp, gpu::Dimension::y>>(
      converter);
  patterns.add<GPUConvert<NVVM::BlockIdZOp, gpu::BlockIdOp, gpu::Dimension::z>>(
      converter);

  patterns.add<BarrierConvert>(converter);
  patterns.add<NVVMRsqrtApproxRaising>(converter);

  patterns
      .add<GPUConvert<NVVM::BlockDimXOp, gpu::BlockDimOp, gpu::Dimension::x>>(
          converter);
  patterns
      .add<GPUConvert<NVVM::BlockDimYOp, gpu::BlockDimOp, gpu::Dimension::y>>(
          converter);
  patterns
      .add<GPUConvert<NVVM::BlockDimZOp, gpu::BlockDimOp, gpu::Dimension::z>>(
          converter);

  patterns.add<GPUConvert<NVVM::GridDimXOp, gpu::GridDimOp, gpu::Dimension::x>>(
      converter);
  patterns.add<GPUConvert<NVVM::GridDimYOp, gpu::GridDimOp, gpu::Dimension::y>>(
      converter);
  patterns.add<GPUConvert<NVVM::GridDimZOp, gpu::GridDimOp, gpu::Dimension::z>>(
      converter);

  patterns.add<CmpFOpLowering, CmpIOpLowering>(converter);
  patterns.add<ReadOnlyAllocaElim>(converter);
  patterns
      .add<AddFOpLowering, AddIOpLowering, AndIOpLowering,
           // AddUIExtendedOpLowering,
           BitcastOpLowering, ConstantOpLowering, DivFOpLowering,
           DivSIOpLowering, DivUIOpLowering, ExtFOpLowering, ExtSIOpLowering,
           ExtUIOpLowering, FPToSIOpLowering, FPToUIOpLowering,
           // IndexCastOpSILowering,
           // IndexCastOpUILowering,
           MaximumFOpLowering, MaxNumFOpLowering, MaxSIOpLowering,
           MaxUIOpLowering, MinimumFOpLowering, MinNumFOpLowering,
           MinSIOpLowering, MinUIOpLowering, MulFOpLowering, MulIOpLowering,
           // MulSIExtendedOpLowering,
           // MulUIExtendedOpLowering,
           NegFOpLowering, OrIOpLowering, RemFOpLowering, RemSIOpLowering,
           RemUIOpLowering, SelectOpLowering, ShLIOpLowering, ShRSIOpLowering,
           ShRUIOpLowering, SIToFPOpLowering, SubFOpLowering, SubIOpLowering,
           TruncFOpLowering, ConstrainedTruncFOpLowering, TruncIOpLowering,
           UIToFPOpLowering, XOrIOpLowering>(converter);
  populateSelectExtractPatterns(patterns);
}

namespace {
struct LibDeviceFuncsRaisingPass
    : public enzyme::impl::LibDeviceFuncsRaisingPassBase<
          LibDeviceFuncsRaisingPass> {
  using LibDeviceFuncsRaisingPassBase::LibDeviceFuncsRaisingPassBase;

  void runOnOperation() override {
    RewritePatternSet patterns(getOperation()->getContext());
    populateLLVMToMathPatterns(getOperation()->getContext(), patterns);
    populateLibDeviceFuncsToOpsPatterns(getOperation()->getContext(), patterns);
    if (remove_freeze)
      patterns.add<RemoveFreeze>(getOperation()->getContext());
    patterns.add<EnzymeInvokeToCall, EnzymeAutodiffOpRaising,
                 EnzymeFwddiffOpRaising>(getOperation()->getContext());
    GreedyRewriteConfig config;
    // We disable region simplification to avoid inadvertently merging
    // llvm.cond_br now that there is an index type.
    config.setRegionSimplificationLevel(
        mlir::GreedySimplifyRegionLevel::Disabled);
    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns),
                                     config))) {
      emitError(getOperation()->getLoc()) << "failed to raise __nv functions";
      return signalPassFailure();
    }
  }
};
} // namespace
