#include "Enzyme/MLIR/Dialect/Dialect.h"
#include "Enzyme/MLIR/Passes/EnzymeBatchPass.h"
#include "src/enzyme_ad/jax/Dialect/Dialect.h"
#include "src/enzyme_ad/jax/Dialect/Ops.h"
#include "src/enzyme_ad/jax/Passes/LinalgUtils.h"
#include "src/enzyme_ad/jax/Passes/Passes.h"
#include "src/enzyme_ad/jax/Utils.h"

#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "stablehlo/dialect/ChloOps.h"
#include "stablehlo/dialect/StablehloOps.h"

#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/SmallVector.h"
#include <mlir/IR/BuiltinAttributes.h>

#define DEBUG_TYPE "lower-enzymexla-blas"

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

using namespace mlir;
using namespace mlir::enzyme;
using namespace mlir::enzymexla;
using namespace mlir::stablehlo;

// Helper function to extract constant scalar value (real/imag parts)
static bool extractConstantScalar(Value val, double &realPart,
                                  double &imagPart) {
  DenseElementsAttr attr;
  if (!matchPattern(val, m_Constant(&attr)))
    return false;

  auto valType = cast<RankedTensorType>(val.getType());
  auto elemType = valType.getElementType();

  if (auto complexType = dyn_cast<ComplexType>(elemType)) {
    // Complex scalar
    auto complexVal = attr.getSplatValue<mlir::Complex<APFloat>>();
    realPart = complexVal.real().convertToDouble();
    imagPart = complexVal.imag().convertToDouble();
    return true;
  } else if (isa<FloatType>(elemType)) {
    // Real scalar
    realPart = attr.getSplatValue<APFloat>().convertToDouble();
    imagPart = 0.0;
    return true;
  }
  return false;
}

// Helper function to create operand and rank for scalar value
// Returns the operand to use and the rank (1 for empty placeholder, 0 for
// scalar)
static std::pair<Value, int64_t> createScalarOperand(PatternRewriter &rewriter,
                                                     Location loc,
                                                     Value originalVal,
                                                     bool useAttribute) {
  if (useAttribute) {
    // Create an empty 0-element tensor as placeholder
    auto emptyType = RankedTensorType::get(
        {0}, cast<RankedTensorType>(originalVal.getType()).getElementType());
    auto emptyTensor = stablehlo::ConstantOp::create(
        rewriter, loc,
        DenseElementsAttr::get(emptyType, ArrayRef<Attribute>{}));
    return {emptyTensor, 1};
  }
  return {originalVal, 0};
}

struct GemmOpLowering : public OpRewritePattern<enzymexla::GemmOp> {
  GemmOpLowering(MLIRContext *context, PatternBenefit benefit = 1)
      : OpRewritePattern(context, benefit) {}

  LogicalResult matchAndRewrite(enzymexla::GemmOp op,
                                PatternRewriter &rewriter) const override {
    auto ctx = op->getContext();
    LLVMTypeConverter typeConverter(ctx);

    auto alpha = op.getOperand(0);
    auto A = op.getOperand(1);
    auto B = op.getOperand(2);
    auto beta = op.getOperand(3);
    auto C = op.getOperand(4);

    auto type_A = cast<RankedTensorType>(A.getType());
    auto type_C = cast<RankedTensorType>(C.getType());
    auto rank = type_A.getRank();

    auto inner_dim_A = op.getTransa() == enzymexla::LapackTranspose::none
                           ? rank - 1
                           : rank - 2;
    auto inner_dim_B = op.getTransb() == enzymexla::LapackTranspose::none
                           ? rank - 2
                           : rank - 1;

    Value maybeConjA = A;
    if (op.getTransa() == enzymexla::LapackTranspose::adjoint) {
      maybeConjA = chlo::ConjOp::create(rewriter, op.getLoc(), A);
    }

    Value maybeConjB = B;
    if (op.getTransb() == enzymexla::LapackTranspose::adjoint) {
      maybeConjB = chlo::ConjOp::create(rewriter, op.getLoc(), B);
    }

    SmallVector<int64_t> leftBatchDims, rightBatchDims;
    for (int64_t i = 0; i < rank - 2; ++i) {
      leftBatchDims.push_back(i);
      rightBatchDims.push_back(i);
    }
    SmallVector<int64_t> leftContractDims = {inner_dim_A};
    SmallVector<int64_t> rightContractDims = {inner_dim_B};
    auto dotDimsAttr = stablehlo::DotDimensionNumbersAttr::get(
        ctx, leftBatchDims, rightBatchDims, leftContractDims,
        rightContractDims);
    auto dotOp = stablehlo::DotGeneralOp::create(rewriter, op.getLoc(), type_C,
                                                 maybeConjA, maybeConjB,
                                                 dotDimsAttr, nullptr, nullptr);
    auto scaledDot =
        stablehlo::MulOpCreate(rewriter, op->getLoc(), alpha, dotOp);

    auto scaledC = stablehlo::MulOpCreate(rewriter, op->getLoc(), beta, C);
    auto addOp =
        stablehlo::AddOpCreate(rewriter, op->getLoc(), scaledDot, scaledC);

    rewriter.replaceAllUsesWith(op.getResult(), addOp);
    rewriter.eraseOp(op);

    return success();
  }
};

struct SymmOpLowering : public OpRewritePattern<enzymexla::SymmOp> {

  using OpRewritePattern<enzymexla::SymmOp>::OpRewritePattern;

  std::string backend;
  int64_t blasIntWidth;
  SymmOpLowering(std::string backend, int64_t blasIntWidth,
                 MLIRContext *context, PatternBenefit benefit = 1)
      : OpRewritePattern(context, benefit), backend(backend),
        blasIntWidth(blasIntWidth) {}

  LogicalResult matchAndRewrite(enzymexla::SymmOp op,
                                PatternRewriter &rewriter) const override {
    auto AType = cast<RankedTensorType>(op.getA().getType());
    auto nBatchDims = AType.getRank() - 2;

    if (nBatchDims == 0) {
      if (backend == "cpu") {
        return matchAndRewriteCPU(op, rewriter);
      } else if (backend == "cuda") {
        return matchAndRewriteCUDA(op, rewriter);
      }
    }

    return matchAndRewriteFallback(op, rewriter);
  }

  LogicalResult matchAndRewriteCPU(enzymexla::SymmOp op,
                                   PatternRewriter &rewriter) const {
    auto ctx = op->getContext();
    LLVMTypeConverter typeConverter(ctx);

    Value a = op.getA();
    Value b = op.getB();
    Value c = op.getC();
    auto side_value = op.getSide() == enzymexla::LapackSide::left ? 'L' : 'R';
    auto uplo_value = op.getUplo() == enzymexla::LapackUplo::L ? 'L' : 'U';

    auto aType = cast<RankedTensorType>(a.getType());
    auto bType = cast<RankedTensorType>(b.getType());
    auto cType = cast<RankedTensorType>(c.getType());
    if (!aType || !bType || !cType)
      return rewriter.notifyMatchFailure(
          op, "operand types not ranked tensor types");

    if (!aType.hasRank() || !bType.hasRank() || !cType.hasRank())
      return rewriter.notifyMatchFailure(op, "expected ranked tensor types");

    if (aType.getRank() != 2 || bType.getRank() > 2 || cType.getRank() > 2)
      return rewriter.notifyMatchFailure(op,
                                         "only 2D matrices supported for symm");

    Type elementType = aType.getElementType();
    auto blasIntType = rewriter.getIntegerType(blasIntWidth);
    auto intType = RankedTensorType::get({}, blasIntType);
    auto uint8Type =
        RankedTensorType::get({}, rewriter.getIntegerType(8, false));
    auto llvmIntType = typeConverter.convertType(blasIntType);
    auto llvmPtrType = LLVM::LLVMPointerType::get(ctx);
    auto llvmVoidType = LLVM::LLVMVoidType::get(ctx);

    std::string blasFn;
    if (auto prefix = lapackPrecisionPrefix(elementType)) {
      blasFn = "enzymexla_blas_" + *prefix + "symm_";
    } else {
      op->emitOpError() << "Unsupported element type: " << elementType;
      return rewriter.notifyMatchFailure(op, "unsupported element type");
    }
    std::string blasFnWrapper = blasFn + "wrapper";

    auto moduleOp = op->getParentOfType<ModuleOp>();

    if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(blasFn)) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(moduleOp.getBody());
      auto funcType = LLVM::LLVMFunctionType::get(llvmVoidType,
                                                  {llvmPtrType, // side
                                                   llvmPtrType, // uplo
                                                   llvmPtrType, // m
                                                   llvmPtrType, // n
                                                   llvmPtrType, // alpha
                                                   llvmPtrType, // A
                                                   llvmPtrType, // lda
                                                   llvmPtrType, // B
                                                   llvmPtrType, // ldb
                                                   llvmPtrType, // beta
                                                   llvmPtrType, // C
                                                   llvmPtrType, // ldc
                                                   llvmIntType, llvmIntType},
                                                  false);
      LLVM::LLVMFuncOp::create(rewriter, op.getLoc(), blasFn, funcType,
                               LLVM::Linkage::External);
    }

    if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(blasFnWrapper)) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(moduleOp.getBody());

      auto funcType = LLVM::LLVMFunctionType::get(llvmVoidType,
                                                  {
                                                      llvmPtrType, // side
                                                      llvmPtrType, // uplo
                                                      llvmPtrType, // m
                                                      llvmPtrType, // n
                                                      llvmPtrType, // alpha
                                                      llvmPtrType, // A
                                                      llvmPtrType, // lda
                                                      llvmPtrType, // B
                                                      llvmPtrType, // ldb
                                                      llvmPtrType, // beta
                                                      llvmPtrType, // C
                                                      llvmPtrType, // ldc
                                                  },
                                                  false);

      auto funcOp =
          LLVM::LLVMFuncOp::create(rewriter, op.getLoc(), blasFnWrapper,
                                   funcType, LLVM::Linkage::Private);
      rewriter.setInsertionPointToStart(funcOp.addEntryBlock(rewriter));

      SmallVector<Value> args(funcOp.getArguments().begin(),
                              funcOp.getArguments().end());
      auto const1 =
          LLVM::ConstantOp::create(rewriter, op.getLoc(), llvmIntType,
                                   rewriter.getIntegerAttr(llvmIntType, 1));
      args.push_back(const1);
      args.push_back(const1);

      LLVM::CallOp::create(rewriter, op.getLoc(), TypeRange{},
                           SymbolRefAttr::get(ctx, blasFn), args);
      LLVM::ReturnOp::create(rewriter, op.getLoc(), ValueRange{});
    }

    static int64_t fn_counter = 0;
    std::string funcFnName = blasFnWrapper + "_" + std::to_string(fn_counter++);

    // Pass A, B and C in row-major format (isColMajor = false) to avoid layout
    // transforms. This requires flipping uplo and side parameters below,
    // similar to the CUDA FFI implementation in xla_ffi.cpp.
    // operandRanks: {uplo, side, m, n, alpha, A, lda, B, ldb, beta, C, ldc}

    SmallVector<bool> isColMajorArr(12, false);
    SmallVector<int64_t> operandRanks = {
        0, 0, 0, 0, 0, 2, 0, op.getB().getType().getRank(), 0, 0, 2, 0};
    SmallVector<int64_t> outputRanks = {2};
    auto operandLayouts =
        getSHLOLayout(rewriter, operandRanks, isColMajorArr, 2);
    auto resultLayouts =
        getSHLOLayout(rewriter, outputRanks, SmallVector<bool>{false}, 2);

    SmallVector<Attribute> aliases;
    aliases.push_back(
        stablehlo::OutputOperandAliasAttr::get(ctx, {}, 10, {})); /*C*/

    func::FuncOp shloFunc;

    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(moduleOp.getBody());

      SmallVector<Type> argTypes = {
          op.getA().getType(),     // A
          op.getB().getType(),     // B
          op.getC().getType(),     // C
          op.getAlpha().getType(), // alpha
          op.getBeta().getType(),  // beta
      };
      SmallVector<Type> retTypes = {op.getC().getType()};

      auto calleeType = rewriter.getFunctionType(argTypes, retTypes);
      shloFunc =
          func::FuncOp::create(rewriter, op.getLoc(), funcFnName, calleeType);
      shloFunc.setPrivate();

      auto &entryBlock = *shloFunc.addEntryBlock();
      rewriter.setInsertionPointToStart(&entryBlock);

      auto A = entryBlock.getArgument(0);
      auto B = entryBlock.getArgument(1);
      auto C = entryBlock.getArgument(2);
      auto alpha = entryBlock.getArgument(3);
      auto beta = entryBlock.getArgument(4);

      // We can swap side since (A*B)^T = B^T*A^T, where B^T is also the
      // column-major interpretation of B. Essentially we ask BLAS to compute
      // C^T in col-major, which is equivalent to C in row-major
      auto side = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), uint8Type,
          cast<ElementsAttr>(
              makeAttr(uint8Type, side_value == 'L' ? 'R' : 'L')));

      // We flip uplo here because A is passed in row-major format.
      // Row-major A is equivalent to A^T in column-major, and since A is
      // symmetric, this means we need to swap upper/lower triangular.
      auto uplo = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), uint8Type,
          cast<ElementsAttr>(
              makeAttr(uint8Type, uplo_value == 'U' ? 'L' : 'U')));

      // For row-major format, lda is the trailing dimension (columns in shape)
      // This is dimension 1 for a 2D matrix
      auto lda = stablehlo::ConvertOp::create(
          rewriter, op.getLoc(), intType,
          stablehlo::GetDimensionSizeOp::create(rewriter, op.getLoc(), A, 1));
      auto ldb = stablehlo::ConvertOp::create(
          rewriter, op.getLoc(), intType,
          stablehlo::GetDimensionSizeOp::create(rewriter, op.getLoc(), B, 1));
      auto ldc = stablehlo::ConvertOp::create(
          rewriter, op.getLoc(), intType,
          stablehlo::GetDimensionSizeOp::create(rewriter, op.getLoc(), C, 1));
      auto mSize = ldc;
      auto nSize = stablehlo::ConvertOp::create(
          rewriter, op.getLoc(), intType,
          stablehlo::GetDimensionSizeOp::create(rewriter, op.getLoc(), C, 0));

      auto jitCall = enzymexla::JITCallOp::create(
          rewriter, op.getLoc(), TypeRange{op.getC().getType()},
          mlir::FlatSymbolRefAttr::get(ctx, blasFnWrapper),
          ValueRange{side, uplo, mSize, nSize, alpha, A, lda, B, ldb, beta, C,
                     ldc},
          rewriter.getStringAttr(""),
          /*operand_layouts=*/operandLayouts,
          /*result_layouts=*/resultLayouts,
          /*arg_attrs=*/nullptr,
          /*res_attrs=*/nullptr,
          /*output_operand_aliases=*/rewriter.getArrayAttr(aliases),
          /*xla_side_effect_free=*/rewriter.getUnitAttr());

      func::ReturnOp::create(rewriter, op.getLoc(),
                             ValueRange{jitCall.getResult(0)});
    }

    auto callOp =
        func::CallOp::create(rewriter, op.getLoc(), shloFunc,
                             ValueRange{op.getA(), op.getB(), op.getC(),
                                        op.getAlpha(), op.getBeta()});

    rewriter.replaceOp(op, callOp);

    return success();
  }

  LogicalResult matchAndRewriteCUDA(enzymexla::SymmOp op,
                                    PatternRewriter &rewriter) const {
    Value A = op.getA();
    if (stablehlo::IsTensorFilled(A)) {
      // CUDA runtime performance for SYMM when the entire tensor is filled is
      // extremely suboptimal. So if the tensor is pre-populated, we can just
      // use the gemm-lowering
      return matchAndRewriteFallback(op, rewriter);
    }

    auto CType = cast<RankedTensorType>(op.getC().getType());
    auto rank = CType.getRank();

    // Try to extract alpha and beta as constants
    double alphaReal = 0.0, alphaImag = 0.0;
    double betaReal = 0.0, betaImag = 0.0;
    bool useAlphaAttr =
        extractConstantScalar(op.getAlpha(), alphaReal, alphaImag);
    bool useBetaAttr = extractConstantScalar(op.getBeta(), betaReal, betaImag);

    // Build operands list - use empty tensors for constant alpha/beta
    SmallVector<Value> operands;
    SmallVector<int64_t> operandRanks;
    SmallVector<bool> areColMajor;

    auto B = op.getB();
    auto C = op.getC();

    auto [alphaOperand, alphaRank] =
        createScalarOperand(rewriter, op.getLoc(), op.getAlpha(), useAlphaAttr);
    operands.push_back(alphaOperand);
    operandRanks.push_back(alphaRank);

    auto [betaOperand, betaRank] =
        createScalarOperand(rewriter, op.getLoc(), op.getBeta(), useBetaAttr);
    operands.push_back(betaOperand);
    operandRanks.push_back(betaRank);

    StringAttr customCallTarget;
    ArrayAttr aliases;

    SmallVector<NamedAttribute> configAttrs = {
        rewriter.getNamedAttr(
            "side",
            rewriter.getBoolAttr(op.getSide() != enzymexla::LapackSide::left)),
        rewriter.getNamedAttr(
            "uplo",
            rewriter.getBoolAttr(op.getUplo() == enzymexla::LapackUplo::U)),
        rewriter.getNamedAttr("use_alpha_attribute",
                              rewriter.getBoolAttr(useAlphaAttr)),
        rewriter.getNamedAttr("alpha_real",
                              rewriter.getF64FloatAttr(alphaReal)),
        rewriter.getNamedAttr("alpha_imag",
                              rewriter.getF64FloatAttr(alphaImag))};

    if (matchPattern(betaOperand, m_AnyZeroFloat()) ||
        matchPattern(betaOperand, m_Zero()) ||
        matchPattern(op.getC(), m_AnyZeroFloat()) ||
        matchPattern(op.getC(), m_Zero())) {
      customCallTarget =
          rewriter.getStringAttr("enzymejax_cublas_symm_no_c_ffi");
      operands = {A, B, alphaOperand};
      operandRanks = {rank, rank, alphaRank};
      aliases = rewriter.getArrayAttr({});
      areColMajor = {false, false, false};
    } else {
      customCallTarget = rewriter.getStringAttr("enzymejax_cublas_symm_ffi");
      operands = {A, B, C, alphaOperand, betaOperand};
      operandRanks = {rank, rank, rank, alphaRank, betaRank};

      configAttrs.push_back(rewriter.getNamedAttr(
          "use_beta_attribute", rewriter.getBoolAttr(useBetaAttr)));
      configAttrs.push_back(rewriter.getNamedAttr(
          "beta_real", rewriter.getF64FloatAttr(betaReal)));
      configAttrs.push_back(rewriter.getNamedAttr(
          "beta_imag", rewriter.getF64FloatAttr(betaImag)));

      aliases = rewriter.getArrayAttr(
          {stablehlo::OutputOperandAliasAttr::get(op.getContext(), {}, 2, {})});
      areColMajor = {false, false, false, true, true};
    }

    DictionaryAttr backendConfig = rewriter.getDictionaryAttr(configAttrs);

    auto customCall = stablehlo::CustomCallOp::create(
        rewriter, op.getLoc(), TypeRange{CType}, operands, customCallTarget,
        /*has_side_effect*/ nullptr,
        /*backend_config*/ backendConfig,
        /*api_version*/
        stablehlo::CustomCallApiVersionAttr::get(
            rewriter.getContext(),
            mlir::stablehlo::CustomCallApiVersion::API_VERSION_TYPED_FFI),
        /*calledcomputations*/ nullptr,
        /*operand_layouts*/
        getSHLOLayout(rewriter, operandRanks, areColMajor, rank),
        /*result_layouts*/
        getSHLOLayout(rewriter, {rank}, SmallVector<bool>(rank, false), rank),
        /*output_operand_aliases*/ aliases,
        /*result_tilings*/ nullptr);

    Value result = customCall.getResult(0);
    rewriter.replaceAllUsesWith(op.getResult(), result);

    return success();
  }

  LogicalResult matchAndRewriteFallback(enzymexla::SymmOp op,
                                        PatternRewriter &rewriter) const {
    auto AType = cast<RankedTensorType>(op.getA().getType());
    auto nBatchDims = AType.getRank() - 2;
    SmallVector<int64_t> batchDims(nBatchDims, 0);
    std::iota(batchDims.begin(), batchDims.end(), 0);

    Value A = op.getA();
    if (!stablehlo::IsTensorFilled(A)) {
      // If the tensor is not filled, we copy to the non-uplo region for safety
      A = stablehlo::copyTriangularPart(rewriter, A, op.getUplo());
      if (!A) {
        return failure();
      }
    }

    // fallback to emitting a stablehlo.dot_general that computes:
    //   alpha*A*B + beta*C if side = 'L'
    //   alpha*B*A + beta*C if side = 'R'
    stablehlo::DotDimensionNumbersAttr dotDims;
    dotDims = stablehlo::DotDimensionNumbersAttr::get(
        op.getContext(), batchDims, batchDims, {nBatchDims + 1}, {nBatchDims});

    stablehlo::DotGeneralOp dotGeneralOp;
    if (op.getSide() == enzymexla::LapackSide::left) {
      dotGeneralOp = stablehlo::DotGeneralOp::create(
          rewriter, op.getLoc(), cast<RankedTensorType>(op.getC().getType()), A,
          op.getB(), dotDims, nullptr, nullptr);
    } else {
      dotGeneralOp = stablehlo::DotGeneralOp::create(
          rewriter, op.getLoc(), cast<RankedTensorType>(op.getC().getType()),
          op.getB(), A, dotDims, nullptr, nullptr);
    }

    auto mul0 = stablehlo::MulOpCreate(rewriter, op->getLoc(), op.getAlpha(),
                                       dotGeneralOp);
    auto mul1 =
        stablehlo::MulOpCreate(rewriter, op->getLoc(), op.getBeta(), op.getC());
    auto res = stablehlo::AddOpCreate(rewriter, op->getLoc(), mul0, mul1);
    rewriter.replaceOp(op, res);
    return success();
  }
};

struct SyrkOpLowering : public OpRewritePattern<enzymexla::SyrkOp> {
  using OpRewritePattern<enzymexla::SyrkOp>::OpRewritePattern;

  SyrkOpLowering(std::string backend, int64_t blasIntWidth,
                 MLIRContext *context, PatternBenefit benefit = 1)
      : OpRewritePattern(context, benefit), backend(backend),
        blasIntWidth(blasIntWidth){};

  LogicalResult matchAndRewrite(enzymexla::SyrkOp op,
                                PatternRewriter &rewriter) const override {
    auto AType = cast<RankedTensorType>(op.getA().getType());
    auto nBatchDims = AType.getRank() - 2;

    if (nBatchDims == 0) {
      if (backend == "cpu") {
        return matchAndRewriteCPU(op, rewriter);
      } else if (backend == "cuda") {
        return matchAndRewriteCUDA(op, rewriter);
      }
    }

    return matchAndRewriteFallback(op, rewriter);
  }

  enum CopyMode { NOT_NEEDED, COPY, TRANSPOSE };

  void resolveUplo(enzymexla::SyrkOp op, enzymexla::LapackUplo &customCallUplo,
                   CopyMode &needsCopy) const {
    switch (op.getUplo()) {
    case enzymexla::LapackUplo::F:
      customCallUplo = standardizeUplo(op.getOutputUplo());
      needsCopy = op.getOutputUplo() == enzymexla::LapackUplo::F
                      ? CopyMode::COPY
                      : CopyMode::NOT_NEEDED;
      break;
    case enzymexla::LapackUplo::L:
      customCallUplo = op.getUplo();
      switch (op.getOutputUplo()) {
      case enzymexla::LapackUplo::F:
        needsCopy = CopyMode::COPY;
        break;
      case enzymexla::LapackUplo::L:
        needsCopy = CopyMode::NOT_NEEDED;
        break;
      case enzymexla::LapackUplo::U:
        needsCopy = CopyMode::TRANSPOSE;
        break;
      }
      break;
    case enzymexla::LapackUplo::U:
      customCallUplo = op.getUplo();
      switch (op.getOutputUplo()) {
      case enzymexla::LapackUplo::F:
        needsCopy = CopyMode::COPY;
        break;
      case enzymexla::LapackUplo::L:
        needsCopy = CopyMode::TRANSPOSE;
        break;
      case enzymexla::LapackUplo::U:
        needsCopy = CopyMode::NOT_NEEDED;
        break;
      }
      break;
    }
  }

  LogicalResult matchAndRewriteCPU(enzymexla::SyrkOp op,
                                   PatternRewriter &rewriter) const {
    auto ctx = op->getContext();
    LLVMTypeConverter typeConverter(ctx);

    auto AType = cast<RankedTensorType>(op.getA().getType());

    auto moduleOp = op->getParentOfType<ModuleOp>();

    auto blasIntType = rewriter.getIntegerType(blasIntWidth);
    auto intType = RankedTensorType::get({}, blasIntType);
    auto uint8Type =
        RankedTensorType::get({}, rewriter.getIntegerType(8, false));
    auto llvmPtrType = LLVM::LLVMPointerType::get(ctx);
    auto llvmVoidType = LLVM::LLVMVoidType::get(ctx);
    auto llvmIntType = typeConverter.convertType(blasIntType);

    std::string blasFn;
    auto prefix = lapackPrecisionPrefix(AType.getElementType());
    if (prefix) {
      blasFn = "enzymexla_blas_" + *prefix + "syrk_";
    } else {
      op->emitOpError() << "Unsupported element type: "
                        << AType.getElementType();
      return rewriter.notifyMatchFailure(op, "unsupported element type");
    }
    std::string blasFnWrapper = blasFn + "wrapper";

    // declare BLAS function declarations if not present
    if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(blasFn)) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(moduleOp.getBody());

      auto funcType = LLVM::LLVMFunctionType::get(
          llvmVoidType,
          {llvmPtrType, llvmPtrType, llvmPtrType, llvmPtrType, llvmPtrType,
           llvmPtrType, llvmPtrType, llvmPtrType, llvmPtrType, llvmPtrType,
           llvmIntType, llvmIntType},
          false);
      LLVM::LLVMFuncOp::create(rewriter, op.getLoc(), blasFn, funcType,
                               LLVM::Linkage::External);
    }

    if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(blasFnWrapper)) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(moduleOp.getBody());

      auto funcType = LLVM::LLVMFunctionType::get(
          llvmVoidType,
          {llvmPtrType, llvmPtrType, llvmPtrType, llvmPtrType, llvmPtrType,
           llvmPtrType, llvmPtrType, llvmPtrType, llvmPtrType, llvmPtrType},
          false);

      auto funcOp =
          LLVM::LLVMFuncOp::create(rewriter, op.getLoc(), blasFnWrapper,
                                   funcType, LLVM::Linkage::Private);
      rewriter.setInsertionPointToStart(funcOp.addEntryBlock(rewriter));

      SmallVector<Value> args(funcOp.getArguments().begin(),
                              funcOp.getArguments().end());
      auto const1 =
          LLVM::ConstantOp::create(rewriter, op.getLoc(), llvmIntType,
                                   rewriter.getIntegerAttr(llvmIntType, 1));
      args.push_back(const1);
      args.push_back(const1);

      LLVM::CallOp::create(rewriter, op.getLoc(), TypeRange{},
                           SymbolRefAttr::get(ctx, blasFn), args);
      LLVM::ReturnOp::create(rewriter, op.getLoc(), ValueRange{});
    }

    CopyMode needsCopy;
    enzymexla::LapackUplo customCallUplo;
    resolveUplo(op, customCallUplo, needsCopy);

    // generate the func.funcOp that calls the blas function
    static int64_t fn_counter = 0;
    std::string funcFnName = blasFnWrapper + "_" + std::to_string(fn_counter++);

    // Pass A and C in row-major format (isColMajor = false) to avoid layout
    // transforms. This requires flipping uplo and transpose parameters below,
    // similar to the CUDA FFI implementation in xla_ffi.cpp.
    // operandRanks: {uplo, trans, n, k, alpha, A, lda, beta, C, ldc}
    SmallVector<bool> isColMajorArr = {false, false, false, false, false,
                                       false, false, false, false, false};
    SmallVector<int64_t> operandRanks = {0, 0, 0, 0, 0, 2, 0, 0, 2, 0};
    SmallVector<int64_t> outputRanks = {2};
    auto operandLayouts =
        getSHLOLayout(rewriter, operandRanks, isColMajorArr, 2);
    auto resultLayouts = getSHLOLayout(rewriter, outputRanks, isColMajorArr, 2);

    SmallVector<Attribute> aliases;
    aliases.push_back(stablehlo::OutputOperandAliasAttr::get(
        ctx, std::vector<int64_t>{}, 8, std::vector<int64_t>{}));

    func::FuncOp shloFunc;

    {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(moduleOp.getBody());

      SmallVector<Type> argTypes = {AType, op.getC().getType(),
                                    op.getAlpha().getType(),
                                    op.getBeta().getType()};
      SmallVector<Type> retTypes = {op.getC().getType()};

      FunctionType calleeType = rewriter.getFunctionType(argTypes, retTypes);
      shloFunc =
          func::FuncOp::create(rewriter, op.getLoc(), funcFnName, calleeType);
      shloFunc.setPrivate();

      auto &entryBlock = *shloFunc.addEntryBlock();
      rewriter.setInsertionPointToStart(&entryBlock);

      auto A = entryBlock.getArgument(0);
      auto C = entryBlock.getArgument(1);
      auto alpha = entryBlock.getArgument(2);
      auto beta = entryBlock.getArgument(3);

      auto nSize = stablehlo::ConvertOp::create(
          rewriter, op.getLoc(), intType,
          stablehlo::GetDimensionSizeOp::create(
              rewriter, op.getLoc(), A,
              op.getTranspose() != enzymexla::LapackTranspose::none));
      auto kSize = stablehlo::ConvertOp::create(
          rewriter, op.getLoc(), intType,
          stablehlo::GetDimensionSizeOp::create(
              rewriter, op.getLoc(), A,
              op.getTranspose() == enzymexla::LapackTranspose::none));

      // For row-major format, lda is the trailing dimension (columns in shape)
      // This is dimension 1 for a 2D matrix
      auto lda = stablehlo::ConvertOp::create(
          rewriter, op.getLoc(), intType,
          stablehlo::GetDimensionSizeOp::create(rewriter, op.getLoc(), A, 1));
      auto ldc = stablehlo::ConvertOp::create(
          rewriter, op.getLoc(), intType,
          stablehlo::GetDimensionSizeOp::create(rewriter, op.getLoc(), C, 1));

      // We flip uplo here because C is passed in row-major format.
      // Row-major C is equivalent to C^T in column-major, and since C is
      // symmetric, this means we need to swap upper/lower triangular.
      auto uploConst = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), uint8Type,
          cast<ElementsAttr>(makeAttr(
              uint8Type,
              customCallUplo == enzymexla::LapackUplo::U ? 'L' : 'U')));
      // We intentionally flip transpose here, this allows us to pass in
      // the data as a row-major format without paying the cost of
      // layout transformation to a col-major (which CPU BLAS uses)
      auto transConst = stablehlo::ConstantOp::create(
          rewriter, op.getLoc(), uint8Type,
          cast<ElementsAttr>(makeAttr(
              uint8Type, op.getTranspose() == enzymexla::LapackTranspose::none
                             ? 'T'
                             : 'N')));

      // {uplo, trans, n, k, alpha, A, lda, beta, C, ldc}
      auto jitCall = enzymexla::JITCallOp::create(
          rewriter, op.getLoc(), TypeRange{op.getC().getType()},
          mlir::FlatSymbolRefAttr::get(ctx, blasFnWrapper),
          ValueRange{uploConst, transConst, nSize, kSize, alpha, A, lda, beta,
                     C, ldc},
          rewriter.getStringAttr(""),
          /*operand_layouts=*/operandLayouts,
          /*result_layouts=*/resultLayouts,
          /*arg_attrs=*/nullptr,
          /*res_attrs=*/nullptr,
          /*output_operand_aliases=*/rewriter.getArrayAttr(aliases),
          /*xla_side_effect_free=*/rewriter.getUnitAttr());

      func::ReturnOp::create(rewriter, op.getLoc(),
                             ValueRange{jitCall.getResult(0)});
    }

    auto callOp = func::CallOp::create(
        rewriter, op.getLoc(), shloFunc,
        ValueRange{op.getA(), op.getC(), op.getAlpha(), op.getBeta()});

    Value result = callOp.getResult(0);
    switch (needsCopy) {
    case CopyMode::COPY:
      result = stablehlo::copyTriangularPart(rewriter, result, customCallUplo);
      break;
    case CopyMode::TRANSPOSE:
      result = stablehlo::TransposeOpCreate(rewriter, op.getLoc(), result,
                                            ArrayRef<int64_t>{1, 0});
      break;
    case CopyMode::NOT_NEEDED:
      break;
    }
    rewriter.replaceAllUsesWith(op.getResult(), result);

    return success();
  }

  LogicalResult matchAndRewriteCUDA(enzymexla::SyrkOp op,
                                    PatternRewriter &rewriter) const {
    auto CType = cast<RankedTensorType>(op.getC().getType());
    auto rank = CType.getRank();

    // Try to extract alpha and beta as constants
    double alphaReal = 0.0, alphaImag = 0.0;
    double betaReal = 0.0, betaImag = 0.0;
    bool useAlphaAttr =
        extractConstantScalar(op.getAlpha(), alphaReal, alphaImag);
    bool useBetaAttr = extractConstantScalar(op.getBeta(), betaReal, betaImag);

    CopyMode needsCopy;
    enzymexla::LapackUplo customCallUplo;
    resolveUplo(op, customCallUplo, needsCopy);

    // Build operands list - use empty tensors for constant alpha/beta
    SmallVector<Value> operands;
    SmallVector<int64_t> operandRanks;
    SmallVector<bool> areColMajor;

    auto A = op.getA();
    auto C = op.getC();

    auto [alphaOperand, alphaRank] =
        createScalarOperand(rewriter, op.getLoc(), op.getAlpha(), useAlphaAttr);
    operands.push_back(alphaOperand);
    operandRanks.push_back(alphaRank);

    auto [betaOperand, betaRank] =
        createScalarOperand(rewriter, op.getLoc(), op.getBeta(), useBetaAttr);
    operands.push_back(betaOperand);
    operandRanks.push_back(betaRank);

    StringAttr customCallTarget;
    ArrayAttr aliases;

    SmallVector<NamedAttribute> configAttrs = {
        rewriter.getNamedAttr(
            "transpose",
            rewriter.getBoolAttr(op.getTranspose() !=
                                 enzymexla::LapackTranspose::none)),
        rewriter.getNamedAttr(
            "uplo",
            rewriter.getBoolAttr(customCallUplo == enzymexla::LapackUplo::U)),
        rewriter.getNamedAttr("use_alpha_attribute",
                              rewriter.getBoolAttr(useAlphaAttr)),
        rewriter.getNamedAttr("alpha_real",
                              rewriter.getF64FloatAttr(alphaReal)),
        rewriter.getNamedAttr("alpha_imag",
                              rewriter.getF64FloatAttr(alphaImag))};

    if (matchPattern(betaOperand, m_AnyZeroFloat()) ||
        matchPattern(betaOperand, m_Zero()) ||
        matchPattern(op.getC(), m_AnyZeroFloat()) ||
        matchPattern(op.getC(), m_Zero())) {
      customCallTarget =
          rewriter.getStringAttr("enzymejax_cublas_syrk_no_c_ffi");
      operands = {A, alphaOperand};
      operandRanks = {rank, alphaRank};
      aliases = rewriter.getArrayAttr({});
      areColMajor = {false, false};
    } else {
      customCallTarget = rewriter.getStringAttr("enzymejax_cublas_syrk_ffi");
      operands = {A, C, alphaOperand, betaOperand};
      operandRanks = {rank, rank, alphaRank, betaRank};

      configAttrs.push_back(rewriter.getNamedAttr(
          "use_beta_attribute", rewriter.getBoolAttr(useBetaAttr)));
      configAttrs.push_back(rewriter.getNamedAttr(
          "beta_real", rewriter.getF64FloatAttr(betaReal)));
      configAttrs.push_back(rewriter.getNamedAttr(
          "beta_imag", rewriter.getF64FloatAttr(betaImag)));

      aliases = rewriter.getArrayAttr(
          {stablehlo::OutputOperandAliasAttr::get(op.getContext(), {}, 1, {})});
      areColMajor = {false, false, true, true};
    }

    DictionaryAttr backendConfig = rewriter.getDictionaryAttr(configAttrs);

    auto customCall = stablehlo::CustomCallOp::create(
        rewriter, op.getLoc(), TypeRange{CType}, operands, customCallTarget,
        /*has_side_effect*/ nullptr,
        /*backend_config*/ backendConfig,
        /*api_version*/
        stablehlo::CustomCallApiVersionAttr::get(
            rewriter.getContext(),
            mlir::stablehlo::CustomCallApiVersion::API_VERSION_TYPED_FFI),
        /*calledcomputations*/ nullptr,
        /*operand_layouts*/
        getSHLOLayout(rewriter, operandRanks, areColMajor, rank),
        /*result_layouts*/
        getSHLOLayout(rewriter, {rank}, SmallVector<bool>(rank, false), rank),
        /*output_operand_aliases*/ aliases,
        /*result_tilings*/ nullptr);

    Value result = customCall.getResult(0);
    switch (needsCopy) {
    case CopyMode::COPY:
      result = stablehlo::copyTriangularPart(rewriter, result, customCallUplo);
      break;
    case CopyMode::TRANSPOSE:
      result = stablehlo::TransposeOpCreate(rewriter, op.getLoc(), result,
                                            ArrayRef<int64_t>{1, 0});
      break;
    case CopyMode::NOT_NEEDED:
      break;
    }
    rewriter.replaceAllUsesWith(op.getResult(), result);

    return success();
  }

  LogicalResult matchAndRewriteFallback(enzymexla::SyrkOp op,
                                        PatternRewriter &rewriter) const {
    auto AType = cast<RankedTensorType>(op.getA().getType());
    auto nBatchDims = AType.getRank() - 2;
    SmallVector<int64_t> batchDims(nBatchDims, 0);
    std::iota(batchDims.begin(), batchDims.end(), 0);

    Value C = op.getC();
    if (!stablehlo::IsTensorFilled(C)) {
      // If the tensor is not filled, we copy to the non-uplo region for safety
      C = stablehlo::copyTriangularPart(rewriter, C, op.getUplo());
      if (!C) {
        return failure();
      }
    }

    // fallback to emitting a stablehlo.dot_general that computes:
    //   alpha * A * A^T + beta * C
    //   alpha * A^T * A + beta * C
    stablehlo::DotDimensionNumbersAttr dotDims;
    switch (op.getTranspose()) {
    case enzymexla::LapackTranspose::none:
      dotDims = stablehlo::DotDimensionNumbersAttr::get(
          op.getContext(), batchDims, batchDims, {nBatchDims + 1},
          {nBatchDims + 1});
      break;
    case enzymexla::LapackTranspose::adjoint:
      LLVM_FALLTHROUGH;
    case enzymexla::LapackTranspose::transpose:
      dotDims = stablehlo::DotDimensionNumbersAttr::get(
          op.getContext(), batchDims, batchDims, {nBatchDims}, {nBatchDims});
      break;
    }

    auto AAT = stablehlo::DotGeneralOp::create(
        rewriter, op.getLoc(), cast<RankedTensorType>(op.getC().getType()),
        op.getA(), op.getA(), dotDims, nullptr, nullptr);

    auto aop =
        stablehlo::MulOpCreate(rewriter, op->getLoc(), op.getAlpha(), AAT);
    auto bop = stablehlo::MulOpCreate(rewriter, op->getLoc(), op.getBeta(), C);

    auto res = stablehlo::AddOpCreate(rewriter, op->getLoc(), aop, bop);
    rewriter.replaceOp(op, res);
    return success();
  }

private:
  std::string backend;
  int64_t blasIntWidth;
};

struct TrsmOpLowering : public OpRewritePattern<enzymexla::TrsmOp> {
  TrsmOpLowering(MLIRContext *context, PatternBenefit benefit = 1)
      : OpRewritePattern(context, benefit) {}

  LogicalResult matchAndRewrite(enzymexla::TrsmOp op,
                                PatternRewriter &rewriter) const override {
    auto ctx = op->getContext();
    LLVMTypeConverter typeConverter(ctx);

    auto alpha = op.getOperand(0);
    auto A = op.getOperand(1);
    auto B = op.getOperand(2);
    auto type_alpha = cast<RankedTensorType>(alpha.getType());
    auto type_A = cast<RankedTensorType>(A.getType());
    auto type_B = cast<RankedTensorType>(B.getType());
    auto type_element = type_alpha.getElementType();

    // (C1)
    if (type_A.getElementType() != type_element ||
        type_B.getElementType() != type_element) {
      return rewriter.notifyMatchFailure(
          op, "Element types of alpha, A and B must match");
    }
    auto rank = type_A.getRank();

    // (C2)
    if (type_A.getRank() != type_B.getRank()) {
      return rewriter.notifyMatchFailure(op, "Ranks of A and B must match");
    }

    // (C3)
    auto shape_A = type_A.getShape();
    auto shape_B = type_B.getShape();
    if (shape_A.drop_back(2) != shape_B.drop_back(2)) {
      return rewriter.notifyMatchFailure(
          op, "Batch dimensions of A and B must match");
    }

    if (shape_A[rank - 1] != shape_A[rank - 2]) {
      return rewriter.notifyMatchFailure(
          op, "Inner two dimensions of A must be square");
    }

    if (shape_A[rank - 1] !=
        shape_B[rank - (op.getSide() == enzymexla::LapackSide::left ? 2 : 1)]) {
      return rewriter.notifyMatchFailure(
          op, "Inner dimensions of A and B must match");
    }

    // (C4)
    if (op.getResult().getType() != type_B) {
      return rewriter.notifyMatchFailure(op, "Result type must match B's type");
    }

    auto scaledB = stablehlo::MulOpCreate(rewriter, op.getLoc(), B, alpha);

    auto transa = stablehlo::Transpose::NO_TRANSPOSE;
    switch (op.getTransa()) {
    case enzymexla::LapackTranspose::none:
      transa = stablehlo::Transpose::NO_TRANSPOSE;
      break;
    case enzymexla::LapackTranspose::transpose:
      transa = stablehlo::Transpose::TRANSPOSE;
      break;
    case enzymexla::LapackTranspose::adjoint:
      transa = stablehlo::Transpose::ADJOINT;
      break;
    }

    auto trisolve_op = stablehlo::TriangularSolveOp::create(
        rewriter, op.getLoc(), TypeRange{type_B}, A, scaledB,
        /*left_side=*/op.getSide() == enzymexla::LapackSide::left,
        /*lower=*/op.getUplo() == enzymexla::LapackUplo::L,
        /*unit_diagonal=*/op.getUnitDiagonal(),
        /*transpose_a=*/transa);

    // replace enzymexla.blas.trsm with the trisolve_op
    rewriter.replaceAllUsesWith(op.getResult(), trisolve_op.getResult());
    rewriter.eraseOp(op);

    return success();
  }
};

struct LowerEnzymeXLABLASPass
    : public enzyme::impl::LowerEnzymeXLABLASPassBase<LowerEnzymeXLABLASPass> {
  using Base::Base;

  void runOnOperation() override {
    auto context = getOperation()->getContext();

    GreedyRewriteConfig config;
    config.setUseTopDownTraversal(true);
    config.enableFolding();

    // We need to run SymmOp lowering first, since it conditionally lowers
    // to a custom call that needs us to detect potential Syrk Ops
    RewritePatternSet patternsSet1(context);
    patternsSet1.add<SymmOpLowering>(backend, blasIntWidth, context);
    if (failed(applyPatternsGreedily(getOperation(), std::move(patternsSet1),
                                     config))) {
      signalPassFailure();
    }

    RewritePatternSet patternsSet2(context);
    patternsSet2.add<SyrkOpLowering>(backend, blasIntWidth, context);
    patternsSet2.add<GemmOpLowering, TrsmOpLowering>(context);
    if (failed(applyPatternsGreedily(getOperation(), std::move(patternsSet2),
                                     config))) {
      signalPassFailure();
    }

    // Verify that all illegal ops have been lowered
    auto walkResult = getOperation()->walk([&](Operation *op) {
      if (isa<enzymexla::GemmOp, enzymexla::SymmOp, enzymexla::SyrkOp,
              enzymexla::TrsmOp>(op)) {
        op->emitError("Failed to lower enzymexla.blas operation");
        return WalkResult::interrupt();
      }
      return WalkResult::advance();
    });

    if (walkResult.wasInterrupted()) {
      signalPassFailure();
    }
  }
};
