#include "lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.h"

#include <cassert>
#include <cstdint>
#include <map>
#include <memory>
#include <numeric>
#include <optional>
#include <set>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include "lib/Dialect/Kernel/IR/KernelOps.h"
#include "lib/Dialect/ModuleAttributes.h"
#include "lib/Dialect/Secret/IR/SecretAttributes.h"
#include "lib/Dialect/Secret/IR/SecretDialect.h"
#include "lib/Dialect/Secret/IR/SecretOps.h"
#include "lib/Dialect/Secret/IR/SecretTypes.h"
#include "lib/Dialect/TensorExt/IR/TensorExtAttributes.h"
#include "lib/Dialect/TensorExt/IR/TensorExtDialect.h"
#include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
#include "lib/Kernel/AbstractValue.h"
#include "lib/Kernel/ArithmeticDag.h"
#include "lib/Kernel/IRMaterializingVisitor.h"
#include "lib/Kernel/KernelImplementation.h"
#include "lib/Kernel/KernelName.h"
#include "lib/Kernel/Utils.h"
#include "lib/Target/CompilationTarget/CompilationTarget.h"
#include "lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.h"
#include "lib/Transforms/ConvertToCiphertextSemantics/TypeConversion.h"
#include "lib/Transforms/DropUnitDims/DropUnitDims.h"
#include "lib/Transforms/LayoutPropagation/Utils.h"
#include "lib/Utils/AttributeUtils.h"
#include "lib/Utils/ContextAwareConversionUtils.h"
#include "lib/Utils/ContextAwareDialectConversion.h"
#include "lib/Utils/ContextAwareTypeConversion.h"
#include "lib/Utils/Layout/Codegen.h"
#include "lib/Utils/Layout/Convolution.h"
#include "lib/Utils/Layout/Utils.h"
#include "lib/Utils/MathUtils.h"
#include "lib/Utils/Utils.h"
#include "llvm/include/llvm/ADT/ArrayRef.h"            // from @llvm-project
#include "llvm/include/llvm/ADT/Hashing.h"             // from @llvm-project
#include "llvm/include/llvm/ADT/STLExtras.h"           // from @llvm-project
#include "llvm/include/llvm/ADT/SmallVector.h"         // from @llvm-project
#include "llvm/include/llvm/ADT/StringExtras.h"        // from @llvm-project
#include "llvm/include/llvm/Support/Debug.h"           // from @llvm-project
#include "llvm/include/llvm/Support/FormatVariadic.h"  // from @llvm-project
#include "llvm/include/llvm/Support/raw_ostream.h"     // from @llvm-project
#include "mlir/include/mlir/Analysis/Presburger/IntegerRelation.h"  // from @llvm-project
#include "mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"   // from @llvm-project
#include "mlir/include/mlir/Dialect/Linalg/IR/Linalg.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/SCF/IR/SCF.h"        // from @llvm-project
#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Utils/StaticValueUtils.h"  // from @llvm-project
#include "mlir/include/mlir/IR/AffineExpr.h"             // from @llvm-project
#include "mlir/include/mlir/IR/AffineMap.h"              // from @llvm-project
#include "mlir/include/mlir/IR/Attributes.h"             // from @llvm-project
#include "mlir/include/mlir/IR/Builders.h"               // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinOps.h"             // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
#include "mlir/include/mlir/IR/Matchers.h"               // from @llvm-project
#include "mlir/include/mlir/IR/OpDefinition.h"           // from @llvm-project
#include "mlir/include/mlir/IR/OperationSupport.h"       // from @llvm-project
#include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
#include "mlir/include/mlir/IR/TypeUtilities.h"          // from @llvm-project
#include "mlir/include/mlir/IR/Value.h"                  // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"              // from @llvm-project
#include "mlir/include/mlir/Support/LogicalResult.h"     // from @llvm-project
#include "mlir/include/mlir/Transforms/DialectConversion.h"  // from @llvm-project
#include "mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project

#define DEBUG_TYPE "convert-to-ciphertext-semantics"

namespace mlir {
namespace heir {

namespace {
using kernel::ArithmeticDagNode;
using kernel::implementHaleviShoup;
using kernel::IRMaterializingVisitor;
using kernel::SSAValue;
using ::mlir::heir::kernel::ArithmeticDagNode;
using ::mlir::heir::kernel::implementHaleviShoup;
using ::mlir::heir::kernel::implementRotateAndReduce;
using ::mlir::heir::kernel::IRMaterializingVisitor;
using ::mlir::heir::kernel::SSAValue;
using presburger::IntegerRelation;
using presburger::VarKind;
using tensor_ext::LayoutAttr;

auto& kLayoutAttrName = tensor_ext::TensorExtDialect::kLayoutAttrName;
auto& kMaterializedAttrName = "tensor_ext.layout_materialized";
auto& kOriginalTypeAttrName =
    tensor_ext::TensorExtDialect::kOriginalTypeAttrName;

NamedAttribute getArithFMF(OpBuilder& builder) {
  return mlir::NamedAttribute(
      mlir::StringAttr::get(builder.getContext(), "fastmath"),
      mlir::arith::FastMathFlagsAttr::get(
          builder.getContext(),
          mlir::arith::FastMathFlags::nsz | mlir::arith::FastMathFlags::nnan));
}

IntegerRelation restrictRelationToSlice(const IntegerRelation& relation,
                                        llvm::ArrayRef<int64_t> offsets,
                                        llvm::ArrayRef<int64_t> sizes) {
  auto result = relation.clone();
  // Set the domain vars to the range specified by the offset and size.
  for (auto [dim, offset] : llvm::enumerate(offsets)) {
    // d0 - offset >= 0
    SmallVector<int64_t> leConstraint(result->getNumCols(), 0);
    leConstraint[dim + result->getVarKindOffset(VarKind::Domain)] = 1;
    leConstraint.back() = -offset;
    result->addInequality(leConstraint);
    // offset + size - d0 - 1 >= 0
    SmallVector<int64_t> geConstraint(result->getNumCols(), 0);
    geConstraint[dim + result->getVarKindOffset(VarKind::Domain)] = -1;
    geConstraint.back() = offset + sizes[dim] - 1;
    result->addInequality(geConstraint);
  }
  return *result;
}

// tensor_ext.remap constrains its input and output types to be the same,
// i.e., remap occurs within one set of ciphertexts. The output of an
// extract_slice, however, may have a layout that has fewer ciphertexts
// in it. For example, extracting one row from a data-semantic matrix that
// is packed with one row per ciphertext would result in a single output
// ciphertext, and the expected layout of the result will reflect that.
// To bridge this gap, kernels must post-processes the remap's output to
// extract the subset ciphertexts relevant to the layout of the output
// slice.
Operation* remapAndExtractResult(ImplicitLocOpBuilder& builder, Value input,
                                 LayoutAttr resultLayout,
                                 RankedTensorType resultType) {
  assert(resultType.getRank() == 2 && "Expected 2D ciphertext semantic type");
  auto remapOp = tensor_ext::RemapOp::create(builder, input, resultLayout);

  SmallVector<OpFoldResult> strides(2, builder.getIndexAttr(1));
  SmallVector<OpFoldResult> offsets(2, builder.getIndexAttr(0));
  SmallVector<OpFoldResult> sizes;
  sizes.push_back(builder.getIndexAttr(resultType.getDimSize(0)));
  sizes.push_back(builder.getIndexAttr(resultType.getDimSize(1)));
  auto extractRemap = tensor::ExtractSliceOp::create(
      builder, resultType, remapOp.getResult(), offsets, sizes, strides);
  return extractRemap;
}

}  // namespace

// An unset value of a permutation as it's being built up.
static constexpr int kUnset = -1;

#define GEN_PASS_DEF_CONVERTTOCIPHERTEXTSEMANTICS
#include "lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.h.inc"

// This type converter converts types like tensor<NxMxi16> where the dimensions
// represent tensor-semantic data to tensor<ciphertext_count x num_slots x
// i16>, where the last dimension represents the ciphertext or plaintext slot
// count, and the other dimensions are determined by a layout attribute
// indexing.
//
// The presence of a layout attribute on the op define a value is required
// for this type converter to trigger. So patterns that use this and convert
// types must remove any layout attributes when they are done.
//
// TODO(#1450): Determine if we should support non-cyclic slot algebras here
// i.e., for the usual 2xN/2 case, how would we determine this situation
// at this stage of the compilation pipeline, and how would this pass update
// to convert to tensor<AxBx2xN/2xi16> where the last two dimensions now
// correspond to the slot algebra direct product?
struct LayoutMaterializationTypeConverter
    : public UniquelyNamedAttributeAwareTypeConverter {
 public:
  LayoutMaterializationTypeConverter(int minSlotCount)
      : UniquelyNamedAttributeAwareTypeConverter(kLayoutAttrName),
        minSlotCount(minSlotCount) {
    // For some reason, directly capturing minSlotCount here leads to memory
    // corruption on that int. Instead, pass the value to a member variable and
    // query it at call time. I have no idea why C++ does this. Debugging it
    // felt like having a stroke.
    addConversion([&](Type type, Attribute attr) { return std::nullopt; });
    addConversion([this](Type type, ArrayAttr attr) -> std::optional<Type> {
      auto composedAttr = LayoutAttr::composeLayouts(attr, type.getContext());
      auto converted = this->convertType(type, composedAttr);
      if (!converted) return std::nullopt;
      return converted;
    });
    addConversion([this](secret::SecretType type,
                         LayoutAttr attr) -> std::optional<Type> {
      auto innerType = type.getValueType();

      FailureOr<Type> convertedInnerType = materializeLayout(
          getElementTypeOrSelf(innerType), attr, getCiphertextSize());
      if (failed(convertedInnerType)) return std::nullopt;
      return secret::SecretType::get(convertedInnerType.value());
    });
    addConversion(
        [this](RankedTensorType type, LayoutAttr attr) -> std::optional<Type> {
          return materializeLayout(getElementTypeOrSelf(type), attr,
                                   getCiphertextSize());
        });
    addConversion(
        [this](IntegerType type, LayoutAttr attr) -> std::optional<Type> {
          return materializeLayout(type, attr, getCiphertextSize());
        });
    addConversion(
        [this](FloatType type, LayoutAttr attr) -> std::optional<Type> {
          return materializeLayout(type, attr, getCiphertextSize());
        });
    addConversion(
        [this](IndexType type, LayoutAttr attr) -> std::optional<Type> {
          return materializeLayout(type, attr, getCiphertextSize());
        });
    addConversion([this](secret::SecretType type,
                         DenseIntElementsAttr attr) -> std::optional<Type> {
      auto innerType = type.getValueType();

      FailureOr<Type> convertedInnerType =
          materializePermutationLayout(innerType, attr, getCiphertextSize());
      if (failed(convertedInnerType)) return std::nullopt;
      return secret::SecretType::get(convertedInnerType.value());
    });
    addConversion([this](RankedTensorType type,
                         DenseIntElementsAttr attr) -> std::optional<Type> {
      return materializePermutationLayout(type, attr, getCiphertextSize());
    });
  }

  int getCiphertextSize() const { return minSlotCount; }

 private:
  int minSlotCount;
};

bool hasMaterializedAttr(Operation* op) {
  return op->hasAttr(kMaterializedAttrName);
}

void setMaterializedAttr(Operation* op) {
  op->setAttr(kMaterializedAttrName, UnitAttr::get(op->getContext()));
}

void setMaterializedAttr(ArrayRef<Operation*> ops) {
  for (auto* op : ops) {
    setMaterializedAttr(op);
  }
}

Type maybeExtractSecretType(Type type) {
  if (auto secretType = dyn_cast<secret::SecretType>(type)) {
    return secretType.getValueType();
  }
  return type;
}

struct ConvertFunc : public ContextAwareFuncConversion {
 public:
  ConvertFunc(const ContextAwareTypeConverter& converter, MLIRContext* context)
      : ContextAwareFuncConversion(converter, context) {}

  LogicalResult finalizeFuncOpModification(
      func::FuncOp op, ArrayRef<Type> oldArgTypes,
      ArrayRef<Type> oldResultTypes, PatternRewriter& rewriter) const override {
    // Replace layout arg attrs with secret.original_type arg attrs This is
    // necessary so that later encoding/decoding functions can know what the
    // original type of the tensor was and how it was encoded.
    rewriter.modifyOpInPlace(op, [&] {
      setMaterializedAttr(op);
      for (int i = 0; i < op.getNumArguments(); ++i) {
        auto layoutAttr = op.getArgAttr(i, kLayoutAttrName);
        if (!layoutAttr || !isa<LayoutAttr, ArrayAttr>(layoutAttr)) {
          continue;
        }

        op.setArgAttr(i, kOriginalTypeAttrName,
                      tensor_ext::OriginalTypeAttr::get(
                          getContext(), maybeExtractSecretType(oldArgTypes[i]),
                          layoutAttr));
      }

      for (int i = 0; i < op.getNumResults(); ++i) {
        auto layoutAttr = op.getResultAttr(i, kLayoutAttrName);
        if (!layoutAttr || !isa<LayoutAttr, ArrayAttr>(layoutAttr)) {
          continue;
        }

        op.setResultAttr(
            i, kOriginalTypeAttrName,
            tensor_ext::OriginalTypeAttr::get(
                getContext(), maybeExtractSecretType(oldResultTypes[i]),
                layoutAttr));
      }
    });
    return success();
  };
};

struct ConvertSecretGeneric : public ConvertAnyContextAware<secret::GenericOp> {
 public:
  ConvertSecretGeneric(const ContextAwareTypeConverter& converter,
                       MLIRContext* context)
      : ConvertAnyContextAware(converter, context) {
    setDebugName("ConvertGeneric");
  }

  LogicalResult finalizeOpModification(
      secret::GenericOp op,
      ContextAwareConversionPatternRewriter& rewriter) const override {
    rewriter.modifyOpInPlace(op, [&] { setMaterializedAttr(op); });
    return success();
  };
};

// Convert an op generically, marking it as materialized. Lowest priority
// because it is only meant to handle ops that don't have special
// materialization rules.
struct ConvertAnyAddingMaterializedAttr : public ConvertAnyContextAware<> {
  ConvertAnyAddingMaterializedAttr(const ContextAwareTypeConverter& converter,
                                   MLIRContext* context)
      : ConvertAnyContextAware(converter, context, /*benefit=*/0) {
    setDebugName("ConvertAnyAddingMaterializedAttr");
  }

  LogicalResult finalizeOpModification(
      Operation* op,
      ContextAwareConversionPatternRewriter& rewriter) const override {
    rewriter.modifyOpInPlace(op, [&] { setMaterializedAttr(op); });
    return success();
  };
};

class ConvertAssignLayout
    : public ContextAwareOpConversionPattern<tensor_ext::AssignLayoutOp> {
 public:
  ConvertAssignLayout(const ContextAwareTypeConverter& typeConverter,
                      mlir::MLIRContext* context, int64_t minSlotCount,
                      CodegenStrategy strategy)
      : ContextAwareOpConversionPattern<tensor_ext::AssignLayoutOp>(
            typeConverter, context),
        minSlotCount(minSlotCount),
        strategy(strategy) {}

  LogicalResult matchAndRewrite(
      tensor_ext::AssignLayoutOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    Value input = adaptor.getValue();
    Attribute layout = op.getLayout();
    if (auto arrayAttr = dyn_cast<ArrayAttr>(layout)) {
      if (arrayAttr.empty() || !isa<LayoutAttr>(arrayAttr.getValue().front())) {
        return failure();
      }
    } else if (!isa<LayoutAttr>(layout) && !isa<DenseIntElementsAttr>(layout)) {
      return failure();
    }

    Type inputType = input.getType();
    Attribute finalLayout = layout;
    if (auto arrayAttr = dyn_cast<ArrayAttr>(layout)) {
      finalLayout = arrayAttr.getValue().back();
    }
    // The finalLayout can be used for type materialization since only the range
    // determines the materialized shape.
    Type resultType =
        getTypeConverter()->convertType(op.getType(), finalLayout);

    // The tensor_ext.layout attribute is still required to persist on any newly
    // created ops. This is distinct from finalLayout; the resultLayout is the
    // composition of all layouts applied in the layout assignment.
    auto resultLayout = findAttributeAssociatedWith(
        op.getResult(), tensor_ext::TensorExtDialect::kLayoutAttrName);
    if (failed(resultLayout)) return failure();
    Attribute resLayout = resultLayout.value();
    if (auto arrayAttr = dyn_cast<ArrayAttr>(resLayout)) {
      if (arrayAttr.empty() || !isa<LayoutAttr>(arrayAttr.getValue().front())) {
        return failure();
      }
    } else if (!isa<LayoutAttr>(resLayout) &&
               !isa<DenseIntElementsAttr>(resLayout)) {
      return failure();
    }

    FunctionKey key(layout, inputType, resultType);

    // Check cache for existing complex layout assignment function.
    func::FuncOp func;
    auto it = functionCache.find(key);
    if (it != functionCache.end()) {
      func = it->second;
      func::CallOp call = rewriter.replaceOpWithNewOp<func::CallOp>(
          op, func, ValueRange{adaptor.getValue()});
      setMaterializedAttr(call);
      call->setAttr(kLayoutAttrName, resultLayout.value());
      return success();
    }

    // Try implementing the layout assignment for this specific input value
    // in-place. If the implementation is simple, the block can be merged
    // directly inline. Otherwise, the block is moved to a function and we call
    // the layout assignment.
    Block* currentBlock = op->getBlock();
    Block* scratchBlock =
        rewriter.splitBlock(currentBlock, Block::iterator(op));
    Block* nextBlock = rewriter.splitBlock(scratchBlock, scratchBlock->begin());
    rewriter.setInsertionPointToEnd(scratchBlock);
    ImplicitLocOpBuilder b(op.getLoc(), rewriter);

    SmallVector<Operation*> createdOps;
    auto createdOpCallback = [&](Operation* createdOp) {
      setMaterializedAttr(createdOp);
      createdOps.push_back(createdOp);
    };

    auto res =
        implementAssignLayout(input, layout, minSlotCount, b, createdOpCallback,
                              op.getDomainSchedule(), strategy);
    if (failed(res)) {
      // Clean up split blocks if implementation failed
      rewriter.mergeBlocks(nextBlock, scratchBlock);
      rewriter.mergeBlocks(scratchBlock, currentBlock);
      return failure();
    }

    // Simple if it's only a few ops and no loops.
    bool isSimple = createdOps.size() <= 3;
    for (auto* createdOp : createdOps) {
      if (isa<scf::ForOp, scf::WhileOp>(createdOp)) {
        isSimple = false;
        break;
      }
    }

    if (isSimple) {
      // Keep implementation and merge blocks back.
      rewriter.mergeBlocks(nextBlock, scratchBlock);
      rewriter.mergeBlocks(scratchBlock, currentBlock);
      setAttributeAssociatedWith(res.value(), kLayoutAttrName,
                                 resultLayout.value());
      rewriter.replaceOp(op, res.value());
      return success();
    }

    // Create a new function for the layout assignment and move the block.
    func = outlineAssignLayoutFunction(op, input, res.value(), scratchBlock,
                                       resultType, rewriter);
    functionCache[key] = func;

    // scratchBlock is moved inside outlineAssignLayoutFunction, but we still
    // need to merge with currentBlock.
    rewriter.mergeBlocks(nextBlock, currentBlock);

    rewriter.setInsertionPointAfter(op);
    func::CallOp call = rewriter.replaceOpWithNewOp<func::CallOp>(
        op, func, ValueRange{adaptor.getValue()});
    setMaterializedAttr(call);
    setAttributeAssociatedWith(call.getResult(0), kLayoutAttrName,
                               resultLayout.value());
    return success();
  };

 private:
  int64_t minSlotCount;
  CodegenStrategy strategy;

  func::FuncOp outlineAssignLayoutFunction(
      tensor_ext::AssignLayoutOp op, Value originalInput, Value originalResult,
      Block* scratchBlock, Type resultType,
      ContextAwareConversionPatternRewriter& rewriter) const {
    func::FuncOp originalFunc = op->getParentOfType<func::FuncOp>();
    ModuleOp module = op->getParentOfType<ModuleOp>();
    auto loc = op.getLoc();

    Type inputType = originalInput.getType();
    std::string funcName = llvm::formatv(
        "_assign_layout_{0}",
        llvm::hash_combine(op.getLayout(), inputType, resultType));
    FunctionType funcType =
        FunctionType::get(rewriter.getContext(), {inputType}, {resultType});

    OpBuilder::InsertionGuard guard(rewriter);
    rewriter.setInsertionPointToStart(module.getBody());
    auto func = func::FuncOp::create(rewriter, loc, funcName, funcType);
    func->setAttr(kClientPackFuncAttrName,
                  rewriter.getDictionaryAttr({rewriter.getNamedAttr(
                      kClientHelperFuncName,
                      rewriter.getStringAttr(originalFunc.getSymName()))}));
    func.setPrivate();

    Block* funcBlock = func.addEntryBlock();

    // Move ops from scratchBlock to funcBlock and remap original input to the
    // funcBlock->getArgument(0).
    rewriter.replaceUsesWithIf(originalInput, funcBlock->getArgument(0),
                               [scratchBlock](OpOperand& use) {
                                 return scratchBlock->findAncestorOpInBlock(
                                            *use.getOwner()) != nullptr;
                               });
    rewriter.mergeBlocks(scratchBlock, funcBlock);

    rewriter.setInsertionPointToEnd(funcBlock);
    func::ReturnOp::create(rewriter, loc, originalResult);

    return func;
  }

  // Cache to store functions generated for a specific layout key is
  // (LayoutAttribute, InputType, ResultType). Layout is uniqued by the layout
  // string.
  using FunctionKey = std::tuple<Attribute, Type, Type>;
  mutable llvm::DenseMap<FunctionKey, func::FuncOp> functionCache;
};

class ConvertConvertLayout
    : public ContextAwareOpConversionPattern<tensor_ext::ConvertLayoutOp> {
 public:
  using ContextAwareOpConversionPattern<
      tensor_ext::ConvertLayoutOp>::ContextAwareOpConversionPattern;

  LogicalResult matchAndRewrite(
      tensor_ext::ConvertLayoutOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    LayoutAttr fromLayout = dyn_cast<LayoutAttr>(op.getFromLayout());
    LayoutAttr toLayout = dyn_cast<LayoutAttr>(op.getToLayout());
    if (!fromLayout || !toLayout) {
      return failure();
    }

    // This is persisted as an operation rather than lowered eagerly to a shift
    // network so as to allow VosVosErkinShiftNetworks to cache its work across
    // multiple instances of the same conversion.
    std::shared_ptr<IntegerRelation> composedLayout =
        fromLayout.getIntegerRelation().clone();
    composedLayout->inverse();
    composedLayout->compose(toLayout.getIntegerRelation());

    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
    LayoutAttr newLayoutAttr =
        LayoutAttr::getFromIntegerRelation(getContext(), *composedLayout);
    auto resultCiphertextSemanticType = cast<RankedTensorType>(
        getTypeConverter()->convertType(op.getResult().getType(), toLayout));
    auto remapAndExtract = remapAndExtractResult(
        b, adaptor.getValue(), newLayoutAttr, resultCiphertextSemanticType);

    setMaterializedAttr(remapAndExtract);
    setAttributeAssociatedWith(remapAndExtract->getResults()[0],
                               kLayoutAttrName, toLayout);
    rewriter.replaceOp(op, remapAndExtract->getResults()[0]);
    return success();
  };
};

// If the mapping is a partial rotation, return the rotation shift amount.
std::optional<int64_t> tryDetectPartialRotation(
    ::llvm::ArrayRef<int64_t> perm) {
  std::optional<int64_t> rotation = std::nullopt;
  for (int64_t i = 0; i < perm.size(); ++i) {
    int64_t input = i;
    int64_t output = perm[i];
    if (output == kUnset) continue;
    // We rotate left in this codebase, so invert normal output - input
    int64_t shiftAmount = -(output - input);
    if (!rotation.has_value()) {
      rotation = shiftAmount;
    } else if (shiftAmount != rotation.value()) {
      return std::nullopt;
    }
  }
  return rotation;
}

// Extend a partial permutation to a full permutation in an FHE-friendly way.
void extendPermutationGreedily(::llvm::MutableArrayRef<int64_t> perm) {
  std::set<int64_t> unmappedInputs;

  // Start with values 0..n-1 and remove when found in the permutation
  std::vector<int64_t> unmappedOutputsVector(perm.size());
  std::iota(unmappedOutputsVector.begin(), unmappedOutputsVector.end(), 0);
  std::set<int64_t> unmappedOutputs(unmappedOutputsVector.begin(),
                                    unmappedOutputsVector.end());

  for (int64_t i = 0; i < perm.size(); ++i) {
    if (perm[i] == kUnset) {
      unmappedInputs.insert(i);
    } else {
      unmappedOutputs.erase(perm[i]);
    }
  }

  // Set iteration is in sorted order, so we're mapping each unused input to
  // the first output index that hasn't been mapped to yet.
  for (const auto& [input, output] :
       llvm::zip(unmappedInputs, unmappedOutputs)) {
    perm[input] = output;
  }
}

// Extend a partial permutation to a full permutation in an FHE-friendly way.
//
// FHE-friendly means that the output permutation should lower to a small shift
// network. For example, if the permutation can be extended to a single
// rotation, it should be.
//
// The input partialPermutation must already be correctly sized (size n for a
// permutation on 1..n). Unset entries of the permutation are indicated by
// kUnset.
void extendPartialPermutation(MutableArrayRef<int64_t> partialPermutation) {
  // If the partially set entries correspond to a single rotation, extend it.
  std::optional<int64_t> rotation =
      tryDetectPartialRotation(partialPermutation);
  if (rotation.has_value()) {
    LLVM_DEBUG(llvm::dbgs() << "Detected partial rotation of offset "
                            << rotation.value() << "\n");
    for (int64_t i = 0; i < partialPermutation.size(); ++i) {
      if (partialPermutation[i] == kUnset) {
        int64_t target = i - rotation.value();
        if (target < 0) target += partialPermutation.size();
        partialPermutation[i] = target;
      }
    }
    return;
  }

  // Otherwise, try to fill in the unset entries greedily.
  extendPermutationGreedily(partialPermutation);
}

template <typename OpTy>
class ConversionBase : public ContextAwareOpConversionPattern<OpTy> {
 public:
  using ContextAwareOpConversionPattern<OpTy>::ContextAwareOpConversionPattern;

  LayoutAttr getLayoutAttr(Value value) const {
    auto layoutLookup = this->getTypeConverter()->getContextualAttr(value);
    if (failed(layoutLookup)) {
      return nullptr;
    }
    return dyn_cast<LayoutAttr>(layoutLookup.value());
  }

  Value materializeKernel(ContextAwareConversionPatternRewriter& rewriter,
                          Location loc,
                          std::shared_ptr<ArithmeticDagNode<SSAValue>> kernel,
                          Type convertedType, Attribute layoutAttr) const {
    IRMaterializingVisitor visitor(convertedType, [&](Operation* createdOp) {
      setMaterializedAttr(createdOp);
    });

    ImplicitLocOpBuilder b(loc, rewriter);
    Value finalOutput = visitor.process(kernel, b)[0];

    auto* finalOutputOp = finalOutput.getDefiningOp();
    finalOutputOp->setAttr(kLayoutAttrName, layoutAttr);
    setMaterializedAttr(finalOutputOp);
    return finalOutput;
  }

  void addBiasAndReplace(ContextAwareConversionPatternRewriter& rewriter,
                         Operation* op, Value finalOutput, Value acc,
                         Attribute layoutAttr) const {
    ImplicitLocOpBuilder b(op->getLoc(), rewriter);
    Operation* addBias =
        makeAppropriatelyTypedAddOp(b, op->getLoc(), finalOutput, acc);
    setAttributeAssociatedWith(addBias->getResults()[0], kLayoutAttrName,
                               layoutAttr);
    setMaterializedAttr(addBias);
    rewriter.replaceOp(op, addBias->getResults());
  }
};

class ConvertLinalgReduce : public ConversionBase<linalg::ReduceOp> {
 public:
  using ConversionBase<linalg::ReduceOp>::ConversionBase;

  void rotateAndReduceKernel(linalg::ReduceOp op, OpAdaptor adaptor,
                             ContextAwareConversionPatternRewriter& rewriter,
                             Operation* innerOp) const {
    auto input = op.getInputs()[0];

    auto originalShape = cast<RankedTensorType>(input.getType()).getShape();
    // Pre-conditions enforce that the reduction operation occurs along the
    // inputs single axis
    unsigned steps = originalShape[op.getDimensions()[0]];
    unsigned period = 1;

    SSAValue vectorLeaf(adaptor.getInputs()[0]);
    kernel::DagType dagType = kernel::mlirTypeToDagType(input.getType());

    auto implementedKernel =
        implementRotateAndReduce(vectorLeaf, {}, period, steps, dagType, {},
                                 innerOp->getName().getStringRef().str());
    rewriter.setInsertionPointAfter(op);

    auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));
    auto convertedType =
        getTypeConverter()->convertType(input.getType(), layoutAttr);

    Value finalOutput = materializeKernel(
        rewriter, op.getLoc(), implementedKernel, convertedType, layoutAttr);

    // Add the initial  value.
    Value result = adaptor.getInits()[0];
    addBiasAndReplace(rewriter, op, finalOutput, result, layoutAttr);
  }

  LogicalResult matchAndRewrite(
      linalg::ReduceOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    // Ensure the reduction op is single addition or multiplication, otherwise
    // there is no kernel.
    Block* body = op.getBlock();
    if (body->getOperations().size() != 2) {
      return op.emitError(
          "linalg.reduce only supported with a single reduction operation");
    }

    // TODO(#1543): support multi-dimension reductions
    if (op.getDimensions().size() != 1) {
      return op.emitError(
          "linalg.reduce only supported with a single reduction dimension");
    }

    if (!op.isSingleInputOutput()) {
      return op.emitError(
          "linalg.reduce only supported with a single reduction dimension");
    }

    Operation* innerOp = &body->getOperations().front();
    if (!isa<arith::AddFOp, arith::MulFOp, arith::AddIOp, arith::MulIOp>(
            innerOp)) {
      return op.emitError()
             << "linalg.reduce only supported with a single addition or "
                "multiplication operation, but found: "
             << innerOp->getName();
    }

    Value input = adaptor.getInputs()[0];

    LayoutAttr inputLayout = getLayoutAttr(input);
    if (!inputLayout)
      return rewriter.notifyMatchFailure(
          op, "missing new layout attribute for input");

    // Based on ImplementRotateAndReduce
    rotateAndReduceKernel(op, adaptor, rewriter, innerOp);
    return success();
  }
};

struct ConvertLinalgDot : public ConversionBase<linalg::DotOp> {
 public:
  using ConversionBase<linalg::DotOp>::ConversionBase;

  LogicalResult matchAndRewrite(
      linalg::DotOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    Value lhs = adaptor.getInputs()[0];
    Value rhs = adaptor.getInputs()[1];
    Value acc = adaptor.getOutputs()[0];

    LayoutAttr lhsLayout = getLayoutAttr(lhs);
    LayoutAttr rhsLayout = getLayoutAttr(rhs);
    if (!lhsLayout || !rhsLayout) {
      return rewriter.notifyMatchFailure(
          op, "missing new layout attribute for inputs");
    }

    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
    Type elementType = cast<RankedTensorType>(lhs.getType()).getElementType();
    std::string reduceOpName;
    if (isa<FloatType>(elementType)) {
      reduceOpName = "arith.addf";
    } else if (isa<IntegerType>(elementType)) {
      reduceOpName = "arith.addi";
    } else {
      return op.emitError("unsupported element type for linalg.dot");
    }

    auto originalShape =
        cast<RankedTensorType>(op.getInputs()[0].getType()).getShape();
    unsigned steps = originalShape[0];

    kernel::DagType dagType = kernel::mlirTypeToDagType(lhs.getType());
    auto implementedKernel =
        implementDot(SSAValue(lhs), SSAValue(rhs), steps, dagType);

    rewriter.setInsertionPointAfter(op);

    auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));

    auto convertedType =
        getTypeConverter()->convertType(lhs.getType(), layoutAttr);

    Value finalOutput = materializeKernel(
        rewriter, op.getLoc(), implementedKernel, convertedType, layoutAttr);

    addBiasAndReplace(rewriter, op, finalOutput, acc, layoutAttr);

    return success();
  }
};

struct ConvertLinalgMatvecLayout : public ConversionBase<linalg::MatvecOp> {
 public:
  using ConversionBase<linalg::MatvecOp>::ConversionBase;

  ConvertLinalgMatvecLayout(
      const ContextAwareTypeConverter& contextAwareTypeConverter,
      MLIRContext* context, bool unrollKernels = true)
      : ConversionBase<linalg::MatvecOp>(contextAwareTypeConverter, context,
                                         /*benefit=*/10),
        unrollKernels(unrollKernels) {}

  bool supportsHaleviShoup(linalg::MatvecOp op, OpAdaptor adaptor) const {
    auto kernelAttr = op->getAttrOfType<secret::KernelAttr>(
        secret::SecretDialect::kKernelAttrName);
    return kernelAttr && kernelAttr.getName() == KernelName::MatvecDiagonal;
  }

  void haleviShoupKernel(
      linalg::MatvecOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const {
    LLVM_DEBUG(llvm::dbgs()
               << "Converting linalg.matvec op with halevi shoup kernel: " << op
               << "\n");

    TypedValue<RankedTensorType> input =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[1]);
    SSAValue vectorLeaf(input);
    TypedValue<RankedTensorType> matrix =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[0]);
    SSAValue matrixLeaf(matrix);

    LayoutAttr matrixLayout = getLayoutAttr(matrix);
    auto matrixRelation = matrixLayout.getIntegerRelation();

    PointCollector collector;
    std::map<int, bool> zeroDiagonals;
    getCtComplementPoints(matrixRelation, collector, matrix.getType());
    for (const auto& point : collector.points) {
      zeroDiagonals[point[0]] = true;
    }
    LLVM_DEBUG(llvm::dbgs()
               << "Got " << zeroDiagonals.size()
               << " zero diagonals for filter: " << matrix << "\n");

    auto dagType = kernel::mlirTypeToDagType(input.getType());
    std::shared_ptr<ArithmeticDagNode<SSAValue>> implementedKernel =
        implementHaleviShoup(vectorLeaf, matrixLeaf,
                             cast<RankedTensorType>(op.getInputs()[0].getType())
                                 .getShape()
                                 .vec(),
                             dagType, zeroDiagonals,
                             /*unroll=*/unrollKernels);

    rewriter.setInsertionPointAfter(op);

    auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));
    Value finalOutput = materializeKernel(
        rewriter, op.getLoc(), implementedKernel, input.getType(), layoutAttr);

    // Add the initial accumulator value.
    Value result = adaptor.getOutputs()[0];
    addBiasAndReplace(rewriter, op, finalOutput, result, layoutAttr);
  }

  LogicalResult matchAndRewrite(
      linalg::MatvecOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    Value matrix = adaptor.getInputs()[0];
    Value vector = adaptor.getInputs()[1];
    LayoutAttr vectorLayout = getLayoutAttr(vector);
    LayoutAttr matrixLayout = getLayoutAttr(matrix);

    if (!matrixLayout || !vectorLayout)
      return rewriter.notifyMatchFailure(
          op, "missing new layout attribute for matrix and vector");

    if (supportsHaleviShoup(op, adaptor)) {
      haleviShoupKernel(op, adaptor, rewriter);
      return success();
    }

    // TODO(#1589): implement row-major naive matvec kernel
    return op.emitError() << "unsupported layout for matrix in matvec: "
                          << matrixLayout;
  }

 private:
  bool unrollKernels;
};

struct PreserveLinalgMatvecAsLinearTransform
    : public ConversionBase<linalg::MatvecOp> {
 public:
  using ConversionBase<linalg::MatvecOp>::ConversionBase;

  PreserveLinalgMatvecAsLinearTransform(
      const ContextAwareTypeConverter& typeConverter, MLIRContext* context)
      : ConversionBase<linalg::MatvecOp>(typeConverter, context,
                                         /*benefit=*/20) {}

  LogicalResult matchAndRewrite(
      linalg::MatvecOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    auto target = getTargetConfig(op->getParentOfType<ModuleOp>());
    if (failed(target) || !target->has_kernel_linear_transform) {
      return rewriter.notifyMatchFailure(op, "linear transform not enabled");
    }

    Value matrixOperand = op.getInputs()[0];
    LayoutAttr matrixLayout = getLayoutAttr(matrixOperand);
    if (!matrixLayout) {
      return rewriter.notifyMatchFailure(op, "missing layout for matrix");
    }

    Value matrix = matrixOperand;
    if (auto assignLayoutOp =
            matrix.getDefiningOp<tensor_ext::AssignLayoutOp>()) {
      matrix = assignLayoutOp.getValue();
    }
    auto constantMatrixOp = matrix.getDefiningOp<arith::ConstantOp>();
    if (!constantMatrixOp) {
      return rewriter.notifyMatchFailure(op, "matrix is not a constant");
    }
    auto denseAttr = dyn_cast<DenseElementsAttr>(constantMatrixOp.getValue());
    if (!denseAttr) {
      return rewriter.notifyMatchFailure(op,
                                         "matrix is not a DenseElementsAttr");
    }

    auto matrixType = cast<RankedTensorType>(matrix.getType());
    auto convertedMatrixType = cast<ShapedType>(
        getTypeConverter()->convertType(matrixType, matrixLayout));
    if (!convertedMatrixType) {
      return rewriter.notifyMatchFailure(op, "failed to convert matrix type");
    }

    int64_t numDiagonals = convertedMatrixType.getShape()[0];
    int64_t slots = convertedMatrixType.getShape()[1];
    auto elementType = matrixType.getElementType();

    Attribute zeroAttr = rewriter.getZeroAttr(elementType);
    std::vector<Attribute> diagonalValues(numDiagonals * slots, zeroAttr);

    auto matrixRelation = matrixLayout.getIntegerRelation();
    PointPairCollector collector(2, 2);
    enumeratePoints(matrixRelation, collector);

    int64_t numCols = matrixType.getDimSize(1);
    for (const auto& pointPair : collector.points) {
      int64_t row = pointPair.first[0];
      int64_t col = pointPair.first[1];
      int64_t d = pointPair.second[0];
      int64_t s = pointPair.second[1];

      int64_t flatIndex = row * numCols + col;
      Attribute val = denseAttr.getValues<Attribute>()[flatIndex];
      diagonalValues[d * slots + s] = val;
    }

    std::vector<int64_t> nonZeroDiagonalIndices;
    std::vector<Attribute> nonZeroDiagonalValues;
    for (int64_t d = 0; d < numDiagonals; ++d) {
      bool isZero = true;
      for (int64_t s = 0; s < slots; ++s) {
        if (diagonalValues[d * slots + s] != zeroAttr) {
          isZero = false;
          break;
        }
      }
      if (!isZero) {
        nonZeroDiagonalIndices.push_back(d);
        for (int64_t s = 0; s < slots; ++s) {
          nonZeroDiagonalValues.push_back(diagonalValues[d * slots + s]);
        }
      }
    }

    auto diagonalsType = RankedTensorType::get(
        {static_cast<int64_t>(nonZeroDiagonalIndices.size()), slots},
        elementType);
    auto diagonalsAttr =
        DenseElementsAttr::get(diagonalsType, nonZeroDiagonalValues);
    auto diagonalIndicesAttr =
        rewriter.getDenseI64ArrayAttr(nonZeroDiagonalIndices);

    auto resultLayout = findAttributeAssociatedWith(
        op.getResult(0), tensor_ext::TensorExtDialect::kLayoutAttrName);
    if (failed(resultLayout)) {
      return rewriter.notifyMatchFailure(op, "missing output layout");
    }

    auto outputType = op.getResult(0).getType();
    auto convertedOutputType =
        getTypeConverter()->convertType(outputType, resultLayout.value());

    rewriter.setInsertionPointAfter(op);
    auto linearTransformOp = rewriter.create<kernel::LinearTransformOp>(
        op.getLoc(), convertedOutputType, adaptor.getInputs()[1], diagonalsAttr,
        diagonalIndicesAttr, /*bsgs_ratio=*/nullptr);

    setMaterializedAttr(linearTransformOp);
    linearTransformOp->setAttr(kLayoutAttrName, resultLayout.value());

    addBiasAndReplace(rewriter, op, linearTransformOp.getResult(),
                      adaptor.getOutputs()[0], resultLayout.value());
    return success();
  }
};

struct ConvertLinalgConv1D : public ConversionBase<linalg::Conv1DOp> {
 public:
  using ConversionBase<linalg::Conv1DOp>::ConversionBase;

  ConvertLinalgConv1D(
      const ContextAwareTypeConverter& contextAwareTypeConverter,
      MLIRContext* context, bool unrollKernels = true)
      : ConversionBase<linalg::Conv1DOp>(contextAwareTypeConverter, context,
                                         /*benefit=*/10),
        unrollKernels(unrollKernels) {}

  bool supportsExpandedHaleviShoup(linalg::Conv1DOp op,
                                   OpAdaptor adaptor) const {
    Value filter = adaptor.getInputs().back();
    auto materializedFilterType = cast<RankedTensorType>(filter.getType());

    // If one of these dimensions is not a power of two, then we can't do
    // the Halevi-Shoup or Squat Packing Matrix Multiplication conversion.
    auto dimensions = materializedFilterType.getShape();
    int64_t numRows = dimensions[0];
    int64_t numCols = dimensions[1];
    bool isPowerOfTwoDims = isPowerOfTwo(numRows) && isPowerOfTwo(numCols);

    auto kernelAttr = op->getAttrOfType<secret::KernelAttr>(
        secret::SecretDialect::kKernelAttrName);
    bool isConv1dAsMatvec =
        kernelAttr && kernelAttr.getName() == KernelName::MatvecDiagonal;

    LLVM_DEBUG(llvm::dbgs()
               << "supports expanded conv1d as matvec with halevi-shoup: "
               << "isPowerOfTwoDims=" << isPowerOfTwoDims
               << " isConv1dAsMatvec=" << isConv1dAsMatvec << "\n");

    return isPowerOfTwoDims && isConv1dAsMatvec;
  }

  void haleviShoupKernel(
      linalg::Conv1DOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const {
    LLVM_DEBUG(llvm::dbgs()
               << "Converting linalg.conv1d op with halevi shoup kernel: " << op
               << "\n");

    TypedValue<RankedTensorType> data =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[0]);
    SSAValue vectorLeaf(data);
    TypedValue<RankedTensorType> filter =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[1]);
    SSAValue matrixLeaf(filter);

    // The original matrix shape is the shape of the expanded filter before
    // diagonalization.
    RankedTensorType expandedMatrixType = get1dConvFilterExpandedType(
        cast<RankedTensorType>(op.getInputs()[1].getType()),
        cast<RankedTensorType>(op.getInputs()[0].getType()), /*stride=*/1,
        /*padding=*/0);

    // Collect any zero diagonals of the filter matrix.
    LayoutAttr filterLayout = getLayoutAttr(adaptor.getInputs()[1]);
    auto filterRelation = filterLayout.getIntegerRelation();

    PointCollector collector;
    std::map<int, bool> zeroDiagonals;
    getCtComplementPoints(filterRelation, collector, filter.getType());
    for (const auto& point : collector.points) {
      zeroDiagonals[point[0]] = true;
    }
    LLVM_DEBUG(llvm::dbgs()
               << "Got " << zeroDiagonals.size()
               << " zero diagonals for filter: " << filter << "\n");

    auto dagType = kernel::mlirTypeToDagType(data.getType());
    std::shared_ptr<ArithmeticDagNode<SSAValue>> implementedKernel =
        implementHaleviShoup(vectorLeaf, matrixLeaf,
                             expandedMatrixType.getShape().vec(), dagType,
                             zeroDiagonals,
                             /*unroll=*/unrollKernels);

    rewriter.setInsertionPointAfter(op);
    auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));

    Value finalOutput = materializeKernel(
        rewriter, op.getLoc(), implementedKernel, data.getType(), layoutAttr);

    Value result = adaptor.getOutputs()[0];
    addBiasAndReplace(rewriter, op, finalOutput, result, layoutAttr);
  }

  LogicalResult matchAndRewrite(
      linalg::Conv1DOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    Value data = adaptor.getInputs().front();
    Value filter = adaptor.getInputs().back();
    LayoutAttr dataLayout = getLayoutAttr(data);
    LayoutAttr filterLayout = getLayoutAttr(filter);

    if (!dataLayout || !filterLayout)
      return rewriter.notifyMatchFailure(
          op, "missing new layout attribute for data and filter");

    if (supportsExpandedHaleviShoup(op, adaptor)) {
      haleviShoupKernel(op, adaptor, rewriter);
      return success();
    }

    return op.emitError() << "unsupported layout for 1d conv";
  }

 private:
  bool unrollKernels;
};

struct ConvertLinalgConv2D : public ConversionBase<linalg::Conv2DOp> {
 public:
  using ConversionBase<linalg::Conv2DOp>::ConversionBase;

  ConvertLinalgConv2D(
      const ContextAwareTypeConverter& contextAwareTypeConverter,
      MLIRContext* context, bool unrollKernels = true)
      : ConversionBase<linalg::Conv2DOp>(contextAwareTypeConverter, context,
                                         /*benefit=*/10),
        unrollKernels(unrollKernels) {}

  bool supportsExpandedHaleviShoup(linalg::Conv2DOp op,
                                   OpAdaptor adaptor) const {
    Value filter = adaptor.getInputs().back();
    auto materializedFilterType = cast<RankedTensorType>(filter.getType());

    // If one of these dimensions is not a power of two, then we can't do
    // the Halevi-Shoup or Squat Packing Matrix Multiplication conversion.
    auto dimensions = materializedFilterType.getShape();
    int64_t numRows = dimensions[0];
    int64_t numCols = dimensions[1];
    bool isPowerOfTwoDims = isPowerOfTwo(numRows) && isPowerOfTwo(numCols);

    auto kernelAttr = op->getAttrOfType<secret::KernelAttr>(
        secret::SecretDialect::kKernelAttrName);
    bool isConv2dAsMatvec =
        kernelAttr && kernelAttr.getName() == KernelName::MatvecDiagonal;

    LLVM_DEBUG(llvm::dbgs()
               << "supports expanded conv2d as matvec with halevi-shoup: "
               << "isPowerOfTwoDims=" << isPowerOfTwoDims
               << " isConv2dAsMatvec=" << isConv2dAsMatvec << "\n");

    return isPowerOfTwoDims && isConv2dAsMatvec;
  }

  void haleviShoupKernel(
      linalg::Conv2DOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const {
    LLVM_DEBUG(llvm::dbgs()
               << "Converting linalg.conv2d op with halevi shoup kernel: " << op
               << "\n");

    TypedValue<RankedTensorType> data =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[0]);
    SSAValue vectorLeaf(data);
    TypedValue<RankedTensorType> matrix =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[1]);
    SSAValue matrixLeaf(matrix);

    // The original matrix shape is the shape of the expanded filter before
    // diagonalization.
    RankedTensorType expandedMatrixType = get2dConvFilterExpandedType(
        cast<RankedTensorType>(op.getInputs()[1].getType()),
        cast<RankedTensorType>(op.getInputs()[0].getType()), /*padding=*/0);

    // Collect any zero diagonals of the filter matrix.
    LayoutAttr filterLayout = getLayoutAttr(adaptor.getInputs()[1]);
    auto filterRelation = filterLayout.getIntegerRelation();

    PointCollector collector;
    std::map<int, bool> zeroDiagonals;
    getCtComplementPoints(filterRelation, collector, matrix.getType());
    for (const auto& point : collector.points) {
      zeroDiagonals[point[0]] = true;
    }
    LLVM_DEBUG(llvm::dbgs()
               << "Got " << zeroDiagonals.size()
               << " zero diagonals for filter: " << matrix << "\n");

    auto dagType = kernel::mlirTypeToDagType(data.getType());
    std::shared_ptr<ArithmeticDagNode<SSAValue>> implementedKernel =
        implementHaleviShoup(vectorLeaf, matrixLeaf,
                             expandedMatrixType.getShape().vec(), dagType,
                             zeroDiagonals,
                             /*unroll=*/unrollKernels);

    rewriter.setInsertionPointAfter(op);

    auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));
    Value finalOutput = materializeKernel(
        rewriter, op.getLoc(), implementedKernel, data.getType(), layoutAttr);

    // Add the initial accumulator value.
    Value result = adaptor.getOutputs()[0];
    addBiasAndReplace(rewriter, op, finalOutput, result, layoutAttr);
  }

  LogicalResult matchAndRewrite(
      linalg::Conv2DOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    Value data = adaptor.getInputs().front();
    Value filter = adaptor.getInputs().back();
    LayoutAttr dataLayout = getLayoutAttr(data);
    LayoutAttr filterLayout = getLayoutAttr(filter);

    if (!dataLayout || !filterLayout)
      return rewriter.notifyMatchFailure(
          op, "missing new layout attribute for data and filter");

    if (supportsExpandedHaleviShoup(op, adaptor)) {
      haleviShoupKernel(op, adaptor, rewriter);
      return success();
    }

    return op.emitError() << "unsupported layout for 2d conv";
  }

 private:
  bool unrollKernels;
};

struct ConvertLinalgConv1DNcwFcw
    : public ConversionBase<linalg::Conv1DNcwFcwOp> {
 public:
  using ConversionBase<linalg::Conv1DNcwFcwOp>::ConversionBase;

  ConvertLinalgConv1DNcwFcw(
      const ContextAwareTypeConverter& contextAwareTypeConverter,
      MLIRContext* context, bool unrollKernels = true)
      : ConversionBase(contextAwareTypeConverter, context,
                       /*benefit=*/10),
        unrollKernels(unrollKernels) {}

  bool supportsExpandedHaleviShoup(linalg::Conv1DNcwFcwOp op,
                                   OpAdaptor adaptor) const {
    Value filter = adaptor.getInputs().back();
    auto materializedFilterType = cast<RankedTensorType>(filter.getType());

    // If one of these dimensions is not a power of two, then we can't do
    // the Halevi-Shoup or Squat Packing Matrix Multiplication conversion.
    auto dimensions = materializedFilterType.getShape();
    int64_t numRows = dimensions[0];
    int64_t numCols = dimensions[1];
    bool isPowerOfTwoDims = isPowerOfTwo(numRows) && isPowerOfTwo(numCols);

    auto kernelAttr = op->getAttrOfType<secret::KernelAttr>(
        secret::SecretDialect::kKernelAttrName);
    bool isConv1dAsMatvec =
        kernelAttr && kernelAttr.getName() == KernelName::MatvecDiagonal;

    LLVM_DEBUG(llvm::dbgs()
               << "supports expanded conv1d as matvec with halevi-shoup: "
               << "isPowerOfTwoDims=" << isPowerOfTwoDims
               << " isConv1dAsMatvec=" << isConv1dAsMatvec << "\n");

    return isPowerOfTwoDims && isConv1dAsMatvec;
  }

  // Rebuild the operand shape LayoutPropagation diagonalized the filter
  // against, honoring any zero tensor.pad it folded into the conv's `padding`
  // parameter.
  FailureOr<RankedTensorType> expandedFilterShape(
      linalg::Conv1DNcwFcwOp op) const {
    auto filterType = cast<RankedTensorType>(op.getInputs()[1].getType());
    auto dataType = cast<RankedTensorType>(op.getInputs()[0].getType());
    int64_t stride =
        llvm::to_vector(op.getStrides().getValues<int64_t>()).front();

    int64_t padding = getConvFoldedPadding(op);
    std::optional<ConvMatrixOperand> matrixOperand =
        foldConvWidthPadding(dataType, padding);
    if (!matrixOperand) {
      return op.emitError()
             << kConvFoldedPaddingAttrName << " of " << padding
             << " does not fit this conv's data operand " << dataType;
    }
    return get1dConvCwFcwFilterExpandedType(filterType, matrixOperand->dataType,
                                            stride, matrixOperand->padding);
  }

  LogicalResult haleviShoupKernel(
      linalg::Conv1DNcwFcwOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const {
    LLVM_DEBUG(
        llvm::dbgs()
        << "Converting linalg.conv_1d_ncw_fcw op with halevi shoup kernel: "
        << op << "\n");

    TypedValue<RankedTensorType> data =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[0]);
    SSAValue vectorLeaf(data);
    TypedValue<RankedTensorType> matrix =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[1]);
    SSAValue matrixLeaf(matrix);

    FailureOr<RankedTensorType> expandedMatrixType = expandedFilterShape(op);
    if (failed(expandedMatrixType)) return failure();
    // Collect any zero diagonals of the filter matrix.
    LayoutAttr filterLayout = getLayoutAttr(adaptor.getInputs()[1]);
    auto filterRelation = filterLayout.getIntegerRelation();

    PointCollector collector;
    std::map<int, bool> zeroDiagonals;
    getCtComplementPoints(filterRelation, collector, matrix.getType());
    for (const auto& point : collector.points) {
      zeroDiagonals[point[0]] = true;
    }
    LLVM_DEBUG(llvm::dbgs()
               << "Got " << zeroDiagonals.size()
               << " zero diagonals for filter: " << matrix << "\n");

    auto dagType = kernel::mlirTypeToDagType(data.getType(),
                                             data.getType().getShape().back());
    std::shared_ptr<ArithmeticDagNode<SSAValue>> implementedKernel =
        implementHaleviShoup(vectorLeaf, matrixLeaf,
                             expandedMatrixType->getShape(), dagType,
                             zeroDiagonals,
                             /*unroll=*/unrollKernels);

    rewriter.setInsertionPointAfter(op);

    auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));

    Value finalOutput = materializeKernel(
        rewriter, op.getLoc(), implementedKernel, data.getType(), layoutAttr);

    // Add the initial accumulator value.
    Value result = adaptor.getOutputs()[0];
    addBiasAndReplace(rewriter, op, finalOutput, result, layoutAttr);
    return success();
  }

  LogicalResult matchAndRewrite(
      linalg::Conv1DNcwFcwOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    Value data = adaptor.getInputs().front();
    Value filter = adaptor.getInputs().back();
    LayoutAttr dataLayout = getLayoutAttr(data);
    LayoutAttr filterLayout = getLayoutAttr(filter);

    if (!dataLayout || !filterLayout) {
      return rewriter.notifyMatchFailure(
          op, "missing new layout attribute for data and filter");
    }

    if (supportsExpandedHaleviShoup(op, adaptor)) {
      return haleviShoupKernel(op, adaptor, rewriter);
    }

    return op.emitError() << "unsupported layout for 1d conv";
  }

 private:
  bool unrollKernels;
};

struct ConvertLinalgConv2DNchwFchw
    : public ConversionBase<linalg::Conv2DNchwFchwOp> {
 public:
  using ConversionBase<linalg::Conv2DNchwFchwOp>::ConversionBase;

  ConvertLinalgConv2DNchwFchw(
      const ContextAwareTypeConverter& contextAwareTypeConverter,
      MLIRContext* context, bool unrollKernels = true)
      : ConversionBase<linalg::Conv2DNchwFchwOp>(contextAwareTypeConverter,
                                                 context,
                                                 /*benefit=*/10),
        unrollKernels(unrollKernels) {}

  bool supportsExpandedHaleviShoup(linalg::Conv2DNchwFchwOp op,
                                   OpAdaptor adaptor) const {
    Value filter = adaptor.getInputs().back();
    auto materializedFilterType = cast<RankedTensorType>(filter.getType());

    // If one of these dimensions is not a power of two, then we can't do
    // the Halevi-Shoup or Squat Packing Matrix Multiplication conversion.
    auto dimensions = materializedFilterType.getShape();
    int64_t numRows = dimensions[0];
    int64_t numCols = dimensions[1];
    bool isPowerOfTwoDims = isPowerOfTwo(numRows) && isPowerOfTwo(numCols);

    auto kernelAttr = op->getAttrOfType<secret::KernelAttr>(
        secret::SecretDialect::kKernelAttrName);
    bool isConv2dAsMatvec =
        kernelAttr && kernelAttr.getName() == KernelName::MatvecDiagonal;

    LLVM_DEBUG(llvm::dbgs()
               << "supports expanded conv2d as matvec with halevi-shoup: "
               << "isPowerOfTwoDims=" << isPowerOfTwoDims
               << " isConv2dAsMatvec=" << isConv2dAsMatvec << "\n");

    return isPowerOfTwoDims && isConv2dAsMatvec;
  }

  void haleviShoupKernel(
      linalg::Conv2DNchwFchwOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const {
    LLVM_DEBUG(
        llvm::dbgs()
        << "Converting linalg.conv_2d_nchw_fchw op with halevi shoup kernel: "
        << op << "\n");

    TypedValue<RankedTensorType> data =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[0]);
    SSAValue vectorLeaf(data);
    TypedValue<RankedTensorType> matrix =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[1]);
    SSAValue matrixLeaf(matrix);

    RankedTensorType dataType =
        cast<RankedTensorType>(op.getInputs()[0].getType());
    auto infoAttr = op->getAttr(kKernelInfoAttrName);
    if (auto info = getKernelInfo(infoAttr)) {
      // Use the kernel info attribute's input shape for the kernel. This
      // accounts for any padding on the data-semantic tensor to allow for
      // channel packing.
      dataType =
          RankedTensorType::get(info->inputShape, dataType.getElementType());
    }

    // The original matrix shape is the shape of the expanded filter before
    // diagonalization.
    // NOTE: `padding=0` is only correct because nothing folds a `tensor.pad`
    // into a 2-D conv's padding parameter. If that changes, go through
    // foldConvWidthPadding the way ConvertLinalgConv1DNcwFcw does; see
    // ConvMatrixOperand.
    RankedTensorType expandedMatrixType = get2dConvChwFchwFilterExpandedType(
        cast<RankedTensorType>(op.getInputs()[1].getType()), dataType,
        /*padding=*/0, llvm::to_vector(op.getStrides().getValues<int64_t>()));
    // Collect any zero diagonals of the filter matrix.
    LayoutAttr filterLayout = getLayoutAttr(adaptor.getInputs()[1]);
    auto filterRelation = filterLayout.getIntegerRelation();

    PointCollector collector;
    std::map<int, bool> zeroDiagonals;
    getCtComplementPoints(filterRelation, collector, matrix.getType());
    for (const auto& point : collector.points) {
      zeroDiagonals[point[0]] = true;
    }
    LLVM_DEBUG(llvm::dbgs()
               << "Got " << zeroDiagonals.size()
               << " zero diagonals for filter: " << matrix << "\n");

    auto dagType = kernel::mlirTypeToDagType(data.getType(),
                                             data.getType().getShape().back());
    std::shared_ptr<ArithmeticDagNode<SSAValue>> implementedKernel =
        implementHaleviShoup(vectorLeaf, matrixLeaf,
                             expandedMatrixType.getShape(), dagType,
                             zeroDiagonals,
                             /*unroll=*/unrollKernels);

    rewriter.setInsertionPointAfter(op);

    auto layoutAttr = op->getAttr(kLayoutAttrName);
    if (auto arrayAttr = dyn_cast<ArrayAttr>(layoutAttr)) {
      layoutAttr = LayoutAttr::composeLayouts(arrayAttr, op.getContext());
    }
    Value finalOutput = materializeKernel(
        rewriter, op.getLoc(), implementedKernel, data.getType(), layoutAttr);

    // Add the initial accumulator value.
    Value result = adaptor.getOutputs()[0];
    addBiasAndReplace(rewriter, op, finalOutput, result, layoutAttr);
  }

  LogicalResult matchAndRewrite(
      linalg::Conv2DNchwFchwOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    Value data = adaptor.getInputs().front();
    Value filter = adaptor.getInputs().back();
    LayoutAttr dataLayout = getLayoutAttr(data);
    LayoutAttr filterLayout = getLayoutAttr(filter);

    if (!dataLayout || !filterLayout) {
      return rewriter.notifyMatchFailure(
          op, "missing new layout attribute for data and filter");
    }

    if (supportsExpandedHaleviShoup(op, adaptor)) {
      haleviShoupKernel(op, adaptor, rewriter);
      return success();
    }

    return op.emitError() << "unsupported layout for 2d conv";
  }

 private:
  bool unrollKernels;
};

Value makeMask(ContextAwareConversionPatternRewriter& rewriter, Location loc,
               Value index, RankedTensorType ciphertextSemanticType) {
  // The ciphertext tensor is a 1D tensor, so the applyOp's result is a
  // single value we can use to build a mask.
  // A tensor of zeros
  auto maskHolder =
      arith::ConstantOp::create(rewriter, loc, ciphertextSemanticType,
                                rewriter.getZeroAttr(ciphertextSemanticType));
  // A scalar 1
  auto one = arith::ConstantOp::create(
      rewriter, loc, ciphertextSemanticType.getElementType(),
      rewriter.getOneAttr(ciphertextSemanticType.getElementType()));
  // insert 1 into the right index
  auto mask = tensor::InsertOp::create(rewriter, loc, one, maskHolder, index);
  setMaterializedAttr({maskHolder, one, mask});
  return mask.getResult();
}

Value makeInverseMask(ContextAwareConversionPatternRewriter& rewriter,
                      Location loc, Value index,
                      RankedTensorType ciphertextSemanticType) {
  // The ciphertext tensor is a 1D tensor, so the applyOp's result is a
  // single value we can use to build a mask.
  // A tensor of ones
  auto maskHolder =
      arith::ConstantOp::create(rewriter, loc, ciphertextSemanticType,
                                rewriter.getOneAttr(ciphertextSemanticType));
  // A scalar 0
  auto one = arith::ConstantOp::create(
      rewriter, loc, ciphertextSemanticType.getElementType(),
      rewriter.getZeroAttr(ciphertextSemanticType.getElementType()));
  // insert 0 into the right index
  auto mask = tensor::InsertOp::create(rewriter, loc, one, maskHolder, index);
  setMaterializedAttr({maskHolder, one, mask});
  return mask.getResult();
}

// Generate IR that loops over the (ciphertext, slot) pairs of the integer
// relation, and at each iteration checks if the corresponding domain index
// tuple matches the dynamicIndices given. If so, it calls thenBuilder
// (forwarded to the then block of an scf.if) with the calculated slot indices
// and iter args, and the return values from thenBuilder are yielded as iter
// args to the next loop iteration.
//
// This is used in some kernels in this file to convert insertion accesses of
// data-semantic tensors into plaintext masks on the corresponding slots of
// ciphertext-semantic tensors.
static FailureOr<SmallVector<Value>> generateLoopWithDynamicIndexCheck(
    ImplicitLocOpBuilder& b, const IntegerRelation& relation,
    ValueRange iterInits, ValueRange dynamicIndices,
    function_ref<SmallVector<Value>(OpBuilder&, Location, ValueRange,
                                    ValueRange)>
        thenBuilder) {
  MLIRLoopNestGenerator generator(b);
  auto loop = generator.generateForLoop(
      relation, iterInits,
      [&](OpBuilder& builder, Location loc, ValueRange exprs,
          ValueRange iterArgs) {
        SmallVector<Value> indices;
        for (int i = relation.getVarKindOffset(VarKind::Range);
             i < relation.getVarKindOffset(VarKind::Range) +
                     relation.getNumRangeVars();
             ++i) {
          indices.push_back(arith::IndexCastOp::create(
              builder, loc, builder.getIndexType(), exprs[i]));
        }

        // If statement to check that the current set of exprs matches
        // the dynamic insertion indices.
        Value exprsMatchIndices;

        // conjunction over all exprs[i] == indices[i] where i is a domain
        // variable.
        int domainStart = relation.getVarKindOffset(VarKind::Domain);
        for (int i = domainStart; i < domainStart + relation.getNumDomainVars();
             ++i) {
          auto eqOp = arith::CmpIOp::create(
              b, arith::CmpIPredicate::eq,
              arith::IndexCastOp::create(builder, loc, builder.getIndexType(),
                                         exprs[i]),
              dynamicIndices[i - domainStart]);
          setMaterializedAttr(eqOp);

          if (i == domainStart) {
            exprsMatchIndices = eqOp.getResult();
          } else {
            auto andOp =
                arith::AndIOp::create(b, exprsMatchIndices, eqOp.getResult());
            exprsMatchIndices = andOp.getResult();
            setMaterializedAttr(andOp);
          }
        }

        auto ifOp = scf::IfOp::create(
            b, exprsMatchIndices,
            /*thenBuilder=*/
            [&](OpBuilder& builder, Location loc) {
              SmallVector<Value> results =
                  thenBuilder(builder, loc, indices, iterArgs);
              scf::YieldOp::create(b, loc, results);
            },
            /*elseBuilder=*/
            [&](OpBuilder& builder, Location loc) {
              // No-op, just scf.yield the current
              // iter args
              scf::YieldOp::create(b, loc, iterArgs);
            });
        setMaterializedAttr(ifOp);

        // Here the scalarMaskValue can be an actual scalar (in the case of
        // a cleartext scalar) or a 1 (in the case of a secret scalar, so
        // that the resulting mask requires an extra mul).
        return scf::ValueVector(ifOp.getResults());
      });
  if (failed(loop)) {
    return failure();
  }

  SmallVector<Value> results = loop.value().getResults();
  return results;
}

class ConvertTensorExtractLayout
    : public ContextAwareOpConversionPattern<tensor::ExtractOp> {
 public:
  using ContextAwareOpConversionPattern<
      tensor::ExtractOp>::ContextAwareOpConversionPattern;

  struct MaskResult {
    Value mask;
    IntegerRelation maskLayout;
  };

  // A helper to create a mask for the case of tensor.extract from a ciphertext
  // when the extraction indices are statically known.
  MaskResult createMaskFromStaticIndices(LayoutAttr tensorLayout,
                                         ArrayRef<int64_t> staticIndices,
                                         OpAdaptor adaptor,
                                         ImplicitLocOpBuilder& b) const {
    IntegerRelation tensorRel = tensorLayout.getIntegerRelation();
    IntegerRelation relWithFixedDomain =
        fixDomainVars(tensorRel, staticIndices);
    relWithFixedDomain.projectOut(tensorRel.getVarKindOffset(VarKind::Domain),
                                  tensorRel.getNumDomainVars());

    RankedTensorType ciphertextSemanticType =
        cast<RankedTensorType>(adaptor.getTensor().getType());
    arith::ConstantOp zeroTensorOp = arith::ConstantOp::create(
        b, ciphertextSemanticType, b.getZeroAttr(ciphertextSemanticType));
    arith::ConstantOp oneScalarOp = arith::ConstantOp::create(
        b, ciphertextSemanticType.getElementType(),
        b.getOneAttr(ciphertextSemanticType.getElementType()));

    // TODO(#2216): use a smarter mask via the desired scalar result layout
    std::vector<int64_t> targetSlot = anyRangePoint(relWithFixedDomain);
    SmallVector<Value> indexCsts;
    for (int64_t v : targetSlot) {
      auto cst = arith::ConstantIndexOp::create(b, v);
      indexCsts.push_back(cst);
      setMaterializedAttr(cst);
    }

    auto insertIntoMask =
        tensor::InsertOp::create(b, oneScalarOp, zeroTensorOp, indexCsts);
    setMaterializedAttr({insertIntoMask, zeroTensorOp, oneScalarOp});

    // Since we only insert a 1 in a single range index, the actual layout
    // here requires the additional constraint that we fix the range variables
    // to that chosen point.
    IntegerRelation resultRelation =
        fixRangeVars(relWithFixedDomain, targetSlot);
    return MaskResult{insertIntoMask.getResult(), resultRelation};
  }

  LogicalResult matchAndRewrite(
      tensor::ExtractOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    // Mask the tensor at the given index, rotate it to the first position, and
    // then ensure its layout matches the expected output layout with a
    // layout_conversion op.
    FailureOr<Attribute> tensorLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getTensor());
    if (failed(tensorLayoutResult)) {
      // If the tensor has no layout, it is a cleartext operation and
      // can be skipped.
      setMaterializedAttr(op);
      return success();
    }
    FailureOr<Attribute> resultLayoutResult =
        getTypeConverter()->getContextualAttr(op.getResult());
    if (failed(resultLayoutResult)) {
      return op.emitError() << "failed to fetch layout attribute for input";
    }

    LayoutAttr tensorLayout = cast<LayoutAttr>(tensorLayoutResult.value());
    LayoutAttr resultLayout = cast<LayoutAttr>(resultLayoutResult.value());
    RankedTensorType ciphertextSemanticType =
        cast<RankedTensorType>(adaptor.getTensor().getType());

    SmallVector<Value> indices(op.getIndices());
    auto staticIndicesResult = getConstantIntValues(getAsOpFoldResult(indices));
    ImplicitLocOpBuilder b(op.getLoc(), rewriter);

    if (!staticIndicesResult.has_value()) {
      // mask = createMaskFromDynamicIndices(tensorLayout, adaptor, b);
      //
      // The problem is that, although we could construct the right mask
      // dynamically, then we need to figure out how to convert the masked
      // ciphertext from its natural layout (zeros in most slots) to the
      // appropriate scalar layout (perhaps replicated everywhere, depending on
      // what layout-optimization picks for the result of this op). This
      // requires, at the very least, finding the index of the first nonzero
      // slot (which we can determine from the mask using a loop, yikes), and
      // then rotating the masked ciphertext to put that first nonzero slot in
      // slot 0, and masking out the rest of the slots so we statically know the
      // layout.
      //
      // As a consequence, we need to perform a rotation by a dynamic index. And
      // it's not clear how we could determine which rotation keys to generate
      // to support that.
      //
      // Punting until we have a strong need for this use case.
      //
      // TODO(#2257): Support dynamic indices in tensor.extract
      return op.emitError() << "tensor.extract with dynamic cleartext indices "
                               "not supported yet";
    }

    MaskResult maskResult = createMaskFromStaticIndices(
        tensorLayout, staticIndicesResult.value(), adaptor, b);
    Value mask = maskResult.mask;

    StringRef mulOpName =
        isa<IntegerType>(ciphertextSemanticType.getElementType())
            ? "arith.muli"
            : "arith.mulf";
    Operation* mulOp = rewriter.create(
        OperationState(op->getLoc(), mulOpName, {mask, adaptor.getTensor()},
                       {ciphertextSemanticType}));

    auto convertLayoutOp = tensor_ext::ConvertLayoutOp::create(
        b, mulOp->getResult(0),
        LayoutAttr::getFromIntegerRelation(op.getContext(),
                                           maskResult.maskLayout),
        resultLayout);

    // Intentionally don't set materialized attr on convert_layout so it will be
    // processed by a subsequent pattern.
    setMaterializedAttr(mulOp);
    // The layout conversion may be folded away, so the mul op also needs an
    // attribute
    setAttributeAssociatedWith(mulOp->getResults()[0], kLayoutAttrName,
                               resultLayout);
    setAttributeAssociatedWith(convertLayoutOp.getResult(), kLayoutAttrName,
                               resultLayout);
    rewriter.replaceOp(op, convertLayoutOp);
    return success();
  }
};

class ConvertTensorPad : public ContextAwareOpConversionPattern<tensor::PadOp> {
 public:
  using ContextAwareOpConversionPattern<
      tensor::PadOp>::ContextAwareOpConversionPattern;

  LogicalResult secretSourceSecretResult(
      tensor::PadOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const {
    auto sourceTy = dyn_cast<RankedTensorType>(
        maybeExtractSecretType(op.getSource().getType()));
    auto resultTy = dyn_cast<RankedTensorType>(
        maybeExtractSecretType(op.getResult().getType()));
    if (!sourceTy || !resultTy || !sourceTy.hasStaticShape() ||
        !resultTy.hasStaticShape()) {
      return op.emitError("Only static shapes are supported for secret pad");
    }
    for (int64_t val : op.getStaticLow()) {
      if (ShapedType::isDynamic(val)) {
        return op.emitError(
            "Only static low padding is supported for secret pad");
      }
    }
    for (int64_t val : op.getStaticHigh()) {
      if (ShapedType::isDynamic(val)) {
        return op.emitError(
            "Only static high padding is supported for secret pad");
      }
    }

    FailureOr<Attribute> sourceLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getSource());
    FailureOr<Attribute> resultLayoutResult =
        getTypeConverter()->getContextualAttr(op.getResult());

    LayoutAttr sourceLayout = cast<LayoutAttr>(sourceLayoutResult.value());
    LayoutAttr resultLayout = cast<LayoutAttr>(resultLayoutResult.value());

    // Check if padded value is constant 0.
    Block& body = op.getRegion().front();
    auto yieldOp = cast<tensor::YieldOp>(body.getTerminator());
    Value yieldedValue = yieldOp.getValue();

    if (!matchPattern(yieldedValue, m_AnyZeroFloat()) &&
        !matchPattern(yieldedValue, m_Zero())) {
      return op.emitError("Only zero padding is supported");
    }

    // Check if sourceCTType == targetCTType
    auto sourceCTType = cast<RankedTensorType>(adaptor.getSource().getType());
    auto targetCTType =
        getTypeConverter()->convertType(op.getResultType(), resultLayout);
    if (!targetCTType) {
      return op.emitError("failed to convert target type");
    }
    if (sourceCTType != targetCTType) {
      return op.emitError("Unsupported pad that changes number of ciphertexts");
    }

    SmallVector<int64_t> lowPadding = llvm::to_vector(op.getStaticLow());

    IntegerRelation paddingRel =
        getPaddingRelation(op.getResultType(), op.getSourceType(), lowPadding);

    IntegerRelation sourceRel = sourceLayout.getIntegerRelation();
    IntegerRelation resultRel = resultLayout.getIntegerRelation();

    paddingRel.inverse();
    sourceRel.inverse();
    sourceRel.compose(paddingRel);
    sourceRel.compose(resultRel);

    LayoutAttr remapLayoutAttr =
        LayoutAttr::getFromIntegerRelation(op.getContext(), sourceRel);

    auto resultCiphertextSemanticType = cast<RankedTensorType>(targetCTType);
    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
    auto remapAndExtract = remapAndExtractResult(
        b, adaptor.getSource(), remapLayoutAttr, resultCiphertextSemanticType);

    setMaterializedAttr(remapAndExtract);
    setAttributeAssociatedWith(remapAndExtract->getResult(0), kLayoutAttrName,
                               resultLayout);
    rewriter.replaceOp(op, remapAndExtract->getResult(0));
    return success();
  }

  LogicalResult matchAndRewrite(
      tensor::PadOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    if (hasMaterializedAttr(op)) return failure();

    FailureOr<Attribute> sourceLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getSource());
    FailureOr<Attribute> resultLayoutResult =
        getTypeConverter()->getContextualAttr(op.getResult());

    bool isSecretSource = succeeded(sourceLayoutResult);
    bool isSecretResult = succeeded(resultLayoutResult);

    if (isSecretSource && isSecretResult) {
      return secretSourceSecretResult(op, adaptor, rewriter);
    }

    if (isSecretSource && !isSecretResult) {
      return op.emitError()
             << "result tensor should have been assigned a layout "
                "by layout-propagation";
    }

    if (!isSecretSource && isSecretResult) {
      return op.emitError()
             << "source tensor should have been assigned a layout "
                "by layout-propagation";
    }

    // Cleartext pad can be elided.
    setMaterializedAttr(op);
    return success();
  }
};

class ConvertTensorInsertSlice
    : public ContextAwareOpConversionPattern<tensor::InsertSliceOp> {
 public:
  using ContextAwareOpConversionPattern<
      tensor::InsertSliceOp>::ContextAwareOpConversionPattern;

  // A helper to create two masks for the case of tensor.insert_slice a
  // cleartext into a ciphertext when the insertion indices are statically
  // known. The scalarMaskValue is the value inserted into the plaintext mask.
  //
  // Returns a pair of scalarMask and destMask.
  std::pair<std::vector<Value>, std::vector<Value>>
  createMasksFromStaticIndicesAtCt(const PointCollector& collector,
                                   RankedTensorType destCiphertextSemanticType,
                                   int64_t scalarMaskValue,
                                   ImplicitLocOpBuilder& b) const {
    // Now create pre-masked packed plainexts that have zeros in all but those
    // slots of dest's ciphertext-semantic tensor that the scalar should be
    // "inserted" into. The nonzero slots contain repeated copies of the scalar
    // value.
    auto numCiphertexts = destCiphertextSemanticType.getDimSize(0);
    Type elementType = destCiphertextSemanticType.getElementType();
    RankedTensorType singleCiphertextType = RankedTensorType::get(
        {1, destCiphertextSemanticType.getDimSize(1)}, elementType);

    // Create dense constants for the scalar mask and dest mask.
    std::vector<SmallVector<Attribute>> scalarMaskAttrs(
        numCiphertexts,
        SmallVector<Attribute>(singleCiphertextType.getNumElements(),
                               b.getZeroAttr(elementType)));
    std::vector<SmallVector<Attribute>> destMaskAttrs(
        numCiphertexts,
        SmallVector<Attribute>(singleCiphertextType.getNumElements(),
                               b.getOneAttr(elementType)));
    Attribute scalarMaskAttr =
        isa<IntegerType>(elementType)
            ? cast<Attribute>(b.getIntegerAttr(elementType, scalarMaskValue))
            : b.getFloatAttr(elementType, scalarMaskValue);

    // Iterate over all these points for each ciphertext, and for the plaintext
    // mask insert a copy of the scalar, while for the dest mask insert a zero.
    for (auto& point : collector.points) {
      auto ctIndex = point[0];
      auto slotIndex = point[1];

      // Here the scalarMaskValue can be an actual scalar (in the case of a
      // cleartext scalar) or a 1 (in the case of a secret scalar, so that the
      // resulting mask requires an extra mul).
      scalarMaskAttrs[ctIndex][slotIndex] = scalarMaskAttr;
      destMaskAttrs[ctIndex][slotIndex] = b.getZeroAttr(elementType);
    }

    std::vector<Value> scalarMasks;
    std::vector<Value> destMasks;
    for (auto ct = 0; ct < numCiphertexts; ++ct) {
      arith::ConstantOp scalarMask = arith::ConstantOp::create(
          b, singleCiphertextType,
          DenseElementsAttr::get(singleCiphertextType, scalarMaskAttrs[ct]));
      arith::ConstantOp destMask = arith::ConstantOp::create(
          b, singleCiphertextType,
          DenseElementsAttr::get(singleCiphertextType, destMaskAttrs[ct]));

      setMaterializedAttr({scalarMask, destMask});
      scalarMasks.push_back(scalarMask.getResult());
      destMasks.push_back(destMask.getResult());
    }

    return {scalarMasks, destMasks};
  }

  LogicalResult secretScalarSecretTensor(
      tensor::InsertSliceOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const {
    MLIRContext* ctx = op.getContext();

    FailureOr<Attribute> scalarLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getSource());
    FailureOr<Attribute> destLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getDest());

    LayoutAttr scalarLayout = cast<LayoutAttr>(scalarLayoutResult.value());
    LayoutAttr destLayout = cast<LayoutAttr>(destLayoutResult.value());

    IntegerRelation scalarRel = scalarLayout.getIntegerRelation();
    IntegerRelation destRel = destLayout.getIntegerRelation();

    // Initial support for this kernel requires that the scalar source slice
    // changes layout to match the destination tensor layout at the subspace
    // defined by the slicing.
    if (!scalarRel.getRangeSet().isSubsetOf(destRel.getRangeSet())) {
      return op.emitError()
             << "tensor.insert requires slice layout to be a subset of the "
                "dest layout, but "
                "got slice layout "
             << scalarLayout << " and tensor layout " << destLayout << "\n";
    }

    SmallVector<int64_t> staticStrides =
        SmallVector<int64_t>(op.getStaticStrides());
    if (!llvm::all_of(staticStrides, [](int64_t s) { return s == 1; })) {
      return op.emitError()
             << "only unit strides of a tensor.insert_slice are supported";
    }

    // Convert the scalar layout to the destination layout subspace.
    auto maybeSliceLayout = getSliceInsertionRelation(
        op.getSourceType(), op.getResultType(),
        SmallVector<int64_t>(op.getStaticOffsets()),
        SmallVector<int64_t>(op.getStaticSizes()), staticStrides);
    if (failed(maybeSliceLayout)) {
      return op.emitError() << "failed to get layout for insert slice";
    }
    auto sliceInsertionLayout = maybeSliceLayout.value();
    sliceInsertionLayout.compose(destRel);

    // If we insert at an offset other than the zero, we may need to reindex
    // the resulting ciphertext so that it is indexed at 0. Shift the range of
    // the ciphertext output var back to index at 0.
    auto ctLowerBound = sliceInsertionLayout.getConstantBound64(
        presburger::BoundType::LB,
        sliceInsertionLayout.getVarKindOffset(presburger::VarKind::Range));
    if (!ctLowerBound) {
      return op.emitError()
             << "failed to get constant bound on ciphertext index";
    }
    // Shift range ct var by -ctLowerBound.
    auto ctVar =
        sliceInsertionLayout.getVarKindOffset(presburger::VarKind::Range);
    auto shiftedSliceInsertionLayout =
        shiftVar(sliceInsertionLayout, ctVar, -ctLowerBound.value());

    // In case the slice is already in a layout compatible with the destination
    // layout, we don't need to insert a conversion.
    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
    Value convertedSource = adaptor.getSource();
    if (!isRelationEqual(scalarRel, shiftedSliceInsertionLayout)) {
      LayoutAttr newScalarLayout =
          LayoutAttr::getFromIntegerRelation(ctx, shiftedSliceInsertionLayout);
      LLVM_DEBUG(llvm::dbgs()
                 << "converting insert_slice source from " << scalarLayout
                 << " to " << newScalarLayout << "\n");
      convertedSource = tensor_ext::ConvertLayoutOp::create(
          b, op.getLoc(), adaptor.getSource(), scalarLayout, newScalarLayout);
    }

    // For each ciphertext in the destination, mask the points originating from
    // the slice with ones and the remaining points from the destination with
    // ones. To do this, we collect all points of the slice in the destination
    // layout.
    IntegerRelation sliceOfDestRelation = restrictRelationToSlice(
        destRel, op.getStaticOffsets(), op.getStaticSizes());
    // collector.points contains all (ciphertext, slot) pairs that correspond
    // to the locations the scalar should be inserted into.
    RankedTensorType ciphertextSemanticType =
        cast<RankedTensorType>(adaptor.getDest().getType());
    PointCollector collector;
    getRangePoints(sliceOfDestRelation, collector);
    std::pair<std::vector<Value>, std::vector<Value>> results =
        createMasksFromStaticIndicesAtCt(collector, ciphertextSemanticType, 1,
                                         b);

    Value result =
        tensor::EmptyOp::create(b, ciphertextSemanticType.getShape(),
                                ciphertextSemanticType.getElementType());
    auto numCt = ciphertextSemanticType.getDimSize(0);
    auto slots = ciphertextSemanticType.getDimSize(1);
    for (auto ct = 0; ct < numCt; ++ct) {
      Value scalarMask = results.first[ct];
      Value destMask = results.second[ct];

      // Mask the scalarMask with the packed scalar. Mask the dest with the
      // destMask, then add the two results.
      SmallVector<OpFoldResult> ctOffsets;
      ctOffsets.push_back(rewriter.getIndexAttr(ct));
      ctOffsets.push_back(rewriter.getIndexAttr(0));
      SmallVector<OpFoldResult> sizes;
      sizes.push_back(b.getIndexAttr(1));
      sizes.push_back(b.getIndexAttr(slots));
      SmallVector<OpFoldResult> strides(2, b.getIndexAttr(1));
      Operation* extractedDest = tensor::ExtractSliceOp::create(
          b, op.getLoc(), cast<RankedTensorType>(convertedSource.getType()),
          adaptor.getDest(), ctOffsets, sizes, strides);
      Operation* scalarMul = makeAppropriatelyTypedMulOp(
          b, op.getLoc(), scalarMask, convertedSource, {getArithFMF(b)});
      Operation* destMul = makeAppropriatelyTypedMulOp(
          b, op.getLoc(), destMask, extractedDest->getResult(0),
          {getArithFMF(b)});
      Operation* finalAdd =
          makeAppropriatelyTypedAddOp(b, op.getLoc(), scalarMul->getResult(0),
                                      destMul->getResult(0), {getArithFMF(b)});

      // Insert the final result into the ciphertext at position ct.
      Operation* insertOp = tensor::InsertSliceOp::create(
          b, finalAdd->getResult(0), result, ctOffsets, sizes, strides);
      setMaterializedAttr(
          {extractedDest, scalarMul, destMul, finalAdd, insertOp});
      result = insertOp->getResult(0);
    }

    setAttributeAssociatedWith(result, kLayoutAttrName, destLayout);
    rewriter.replaceOp(op, result);
    return success();
  }

  LogicalResult matchAndRewrite(
      tensor::InsertSliceOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    FailureOr<Attribute> sourceLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getSource());
    FailureOr<Attribute> destLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getDest());

    bool isSecretScalar = succeeded(sourceLayoutResult);
    bool isSecretDest = succeeded(destLayoutResult);

    if (isSecretScalar && isSecretDest) {
      return secretScalarSecretTensor(op, adaptor, rewriter);
    }

    if (!isSecretScalar && isSecretDest) {
      return op.emitError()
             << "unsupported case of inserting cleartext into secret tensor";
    }

    if (isSecretScalar && !isSecretDest) {
      return op.emitError() << "dest tensor should have been assigned a layout "
                               "by layout-propagation";
    }

    // cleartext scalar and cleartext tensor means this is a cleartext op
    // that can be elided.
    setMaterializedAttr(op);
    return success();
  }
};

class ConvertTensorInsertLayout
    : public ContextAwareOpConversionPattern<tensor::InsertOp> {
 public:
  using ContextAwareOpConversionPattern<
      tensor::InsertOp>::ContextAwareOpConversionPattern;

  // A helper to create two masks for the case of tensor.insert a cleartext
  // into a ciphertext when the insertion indices are statically known. The
  // scalarMaskValue is the value inserted into the plaintext mask.
  //
  // Returns a pair of scalarMask and destMask.
  std::pair<Value, Value> createMasksFromStaticIndices(
      LayoutAttr destLayout, ArrayRef<int64_t> staticIndices,
      tensor::InsertOp op, OpAdaptor adaptor, Value scalarMaskValue,
      ImplicitLocOpBuilder& b) const {
    IntegerRelation destRel = destLayout.getIntegerRelation();
    IntegerRelation fixedRel = fixDomainVars(destRel, staticIndices);
    fixedRel.projectOut(destRel.getVarKindOffset(VarKind::Domain),
                        destRel.getNumDomainVars());

    // Now create a pre-masked packed plainext that has zeros in all but those
    // slots of dest's ciphertext-semantic tensor that the scalar should be
    // "inserted" into. The nonzero slots contain repeated copies of the scalar
    // value.
    RankedTensorType ciphertextSemanticType =
        cast<RankedTensorType>(adaptor.getDest().getType());

    arith::ConstantOp zeroTensorOp = arith::ConstantOp::create(
        b, ciphertextSemanticType, b.getZeroAttr(ciphertextSemanticType));
    arith::ConstantOp zeroScalarOp = arith::ConstantOp::create(
        b, ciphertextSemanticType.getElementType(),
        b.getZeroAttr(ciphertextSemanticType.getElementType()));
    arith::ConstantOp oneOp = arith::ConstantOp::create(
        b, ciphertextSemanticType, b.getOneAttr(ciphertextSemanticType));

    // collector.points contains all (ciphertext, slot) pairs that correspond
    // to the locations the scalar should be inserted into.
    PointCollector collector;
    getRangePoints(fixedRel, collector);

    // Iterate over all these points, and for the plaintext mask insert a copy
    // of the scalar, while for the dest mask insert a zero.
    Value scalarMask = zeroTensorOp.getResult();
    Value destMask = oneOp.getResult();
    for (const auto& point : collector.points) {
      SmallVector<Value> indexCsts;
      for (int64_t v : point) {
        auto cst = arith::ConstantIndexOp::create(b, v);
        indexCsts.push_back(cst);
        setMaterializedAttr(cst);
      }

      // Here the scalarMaskValue can be an actual scalar (in the case of a
      // cleartext scalar) or a 1 (in the case of a secret scalar, so that the
      // resulting mask requires an extra mul).
      auto insertIntoScalarMask =
          tensor::InsertOp::create(b, scalarMaskValue, scalarMask, indexCsts);
      auto insertIntoDestMask =
          tensor::InsertOp::create(b, zeroScalarOp, destMask, indexCsts);
      scalarMask = insertIntoScalarMask.getResult();
      destMask = insertIntoDestMask.getResult();
      setMaterializedAttr({insertIntoScalarMask, insertIntoDestMask});
    }

    setMaterializedAttr({zeroTensorOp, zeroScalarOp, oneOp});
    return {scalarMask, destMask};
  }

  // A helper to create two masks for the case of tensor.insert a cleartext
  // into a ciphertext when the insertion indices are not known statically.
  // The scalarMaskValue is the value inserted into the plaintext mask.
  //
  // Returns a pair of scalarMask and destMask.
  std::pair<Value, Value> createMasksFromDynamicIndices(
      LayoutAttr destLayout, tensor::InsertOp op, OpAdaptor adaptor,
      Value scalarMaskValue, ImplicitLocOpBuilder& b) const {
    IntegerRelation destRel = destLayout.getIntegerRelation();
    // Here we need to loop over the relation in the same manner as we do
    // codegen for assign_layout, but insert mask entries instead of copying
    // from a data tensor.
    RankedTensorType ciphertextSemanticType =
        cast<RankedTensorType>(adaptor.getDest().getType());

    arith::ConstantOp zeroTensorOp = arith::ConstantOp::create(
        b, ciphertextSemanticType, b.getZeroAttr(ciphertextSemanticType));
    arith::ConstantOp zeroScalarOp = arith::ConstantOp::create(
        b, ciphertextSemanticType.getElementType(),
        b.getZeroAttr(ciphertextSemanticType.getElementType()));
    arith::ConstantOp oneOp = arith::ConstantOp::create(
        b, ciphertextSemanticType, b.getOneAttr(ciphertextSemanticType));

    Value scalarMask = zeroTensorOp.getResult();
    Value destMask = oneOp.getResult();
    SmallVector<Value> indices(op.getIndices());

    auto loop = generateLoopWithDynamicIndexCheck(
        b, destRel, {scalarMask, destMask}, indices,
        [&](OpBuilder& builder, Location loc, ValueRange slotIndices,
            ValueRange iterArgs) {
          Value curScalarMask = iterArgs[0];
          Value curDestMask = iterArgs[1];
          SmallVector<Value> slotIndicesAsIndex;
          for (Value index : slotIndices) {
            slotIndicesAsIndex.push_back(arith::IndexCastOp::create(
                builder, loc, builder.getIndexType(), index));
          }
          auto insertIntoScalarMask = tensor::InsertOp::create(
              b, scalarMaskValue, curScalarMask, slotIndicesAsIndex);
          auto insertIntoDestMask = tensor::InsertOp::create(
              b, zeroScalarOp, curDestMask, slotIndicesAsIndex);
          setMaterializedAttr({insertIntoScalarMask, insertIntoDestMask});
          SmallVector<Value> results = {insertIntoScalarMask.getResult(),
                                        insertIntoDestMask.getResult()};
          return results;
        });
    if (failed(loop)) {
      op.emitError() << "Failed to generate loop nest for layout "
                     << destLayout;
      return {Value(), Value()};
    }

    setMaterializedAttr({zeroTensorOp, zeroScalarOp, oneOp});
    return {loop.value()[0], loop.value()[1]};
  }

  LogicalResult cleartextScalarSecretTensor(
      tensor::InsertOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const {
    LayoutAttr destLayout = cast<LayoutAttr>(
        getTypeConverter()->getContextualAttr(adaptor.getDest()).value());

    Value scalarMask;
    Value destMask;

    SmallVector<Value> indices(op.getIndices());
    auto staticIndicesResult = getConstantIntValues(getAsOpFoldResult(indices));
    RankedTensorType ciphertextSemanticType =
        cast<RankedTensorType>(adaptor.getDest().getType());
    ImplicitLocOpBuilder b(op.getLoc(), rewriter);

    if (staticIndicesResult.has_value()) {
      SmallVector<int64_t> staticIndices = staticIndicesResult.value();
      std::pair<Value, Value> results = createMasksFromStaticIndices(
          destLayout, staticIndices, op, adaptor, adaptor.getScalar(), b);
      scalarMask = results.first;
      destMask = results.second;
    } else {
      std::pair<Value, Value> results = createMasksFromDynamicIndices(
          destLayout, op, adaptor, adaptor.getScalar(), b);
      scalarMask = results.first;
      destMask = results.second;
    }

    // Mask the dest with the destMask, then add the result to the
    // scalarMask
    StringRef mulOpName =
        isa<IntegerType>(ciphertextSemanticType.getElementType())
            ? "arith.muli"
            : "arith.mulf";
    StringRef addOpName =
        isa<IntegerType>(ciphertextSemanticType.getElementType())
            ? "arith.addi"
            : "arith.addf";

    Operation* destMul = b.create(OperationState(op.getLoc(), mulOpName,
                                                 {destMask, adaptor.getDest()},
                                                 {ciphertextSemanticType}));
    Operation* finalAdd = b.create(OperationState(
        op.getLoc(), addOpName, {scalarMask, destMul->getResult(0)},
        {ciphertextSemanticType}));
    Value result = finalAdd->getResult(0);
    setMaterializedAttr({destMul, finalAdd});
    setAttributeAssociatedWith(result, kLayoutAttrName, destLayout);
    rewriter.replaceOp(op, result);
    return success();
  }

  LogicalResult secretScalarSecretTensor(
      tensor::InsertOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const {
    FailureOr<Attribute> scalarLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getScalar());
    FailureOr<Attribute> destLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getDest());

    LayoutAttr scalarLayout = dyn_cast<LayoutAttr>(scalarLayoutResult.value());
    LayoutAttr destLayout = dyn_cast<LayoutAttr>(destLayoutResult.value());

    // Initial support for this kernel requires the sizes of the range to
    // match. Ideally we could be smarter here, for example to take a scalar
    // layout which is a single slot in a single ciphertext, and insert it into
    // the correct slot of the single ciphertext of the dest tensor via a mask
    // and rotate. But this kernel is not expected to be used, so we can instead
    // incur the cost of a layout conversion before the insert.
    IntegerRelation scalarRel = scalarLayout.getIntegerRelation();
    IntegerRelation destRel = destLayout.getIntegerRelation();
    if (!isRelationEqual(scalarRel.getRangeSet(), destRel.getRangeSet())) {
      return op.emitError()
             << "tensor.insert requires scalar and tensor layout to match, but "
                "got scalar layout "
             << scalarLayout << " and tensor layout " << destLayout << "\n";
    }

    Value scalarMask;
    Value destMask;

    SmallVector<Value> indices(op.getIndices());
    auto staticIndicesResult = getConstantIntValues(getAsOpFoldResult(indices));
    RankedTensorType ciphertextSemanticType =
        cast<RankedTensorType>(adaptor.getDest().getType());
    ImplicitLocOpBuilder b(op.getLoc(), rewriter);

    // In the cleartextScalar case, we inserted the scalar value directly in
    // the slots of the plaintext mask. Here we insert a 1, and then multiply
    // the resulting mask by the packed scalar ciphertext.
    Value scalarOne = arith::ConstantOp::create(
        b, ciphertextSemanticType.getElementType(),
        b.getOneAttr(ciphertextSemanticType.getElementType()));

    if (staticIndicesResult.has_value()) {
      SmallVector<int64_t> staticIndices = staticIndicesResult.value();
      std::pair<Value, Value> results = createMasksFromStaticIndices(
          destLayout, staticIndices, op, adaptor, scalarOne, b);
      scalarMask = results.first;
      destMask = results.second;
    } else {
      std::pair<Value, Value> results =
          createMasksFromDynamicIndices(destLayout, op, adaptor, scalarOne, b);
      scalarMask = results.first;
      destMask = results.second;
    }

    // Mask the scalarMask with the packed scalar. Mask the dest with the
    // destMask, then add the two results.
    StringRef mulOpName =
        isa<IntegerType>(ciphertextSemanticType.getElementType())
            ? "arith.muli"
            : "arith.mulf";
    StringRef addOpName =
        isa<IntegerType>(ciphertextSemanticType.getElementType())
            ? "arith.addi"
            : "arith.addf";

    // Mul the masks to the respective ciphertexts
    Operation* scalarMul = b.create(OperationState(
        op.getLoc(), mulOpName, {scalarMask, adaptor.getScalar()},
        {ciphertextSemanticType}));
    Operation* destMul = b.create(OperationState(op.getLoc(), mulOpName,
                                                 {destMask, adaptor.getDest()},
                                                 {ciphertextSemanticType}));
    Operation* finalAdd = b.create(
        OperationState(op.getLoc(), addOpName,
                       {scalarMul->getResult(0), destMul->getResult(0)},
                       {ciphertextSemanticType}));
    Value result = finalAdd->getResult(0);
    setMaterializedAttr({scalarMul, destMul, finalAdd});
    setAttributeAssociatedWith(result, kLayoutAttrName, destLayout);
    rewriter.replaceOp(op, result);
    return success();
  }

  LogicalResult matchAndRewrite(
      tensor::InsertOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    FailureOr<Attribute> scalarLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getScalar());
    FailureOr<Attribute> destLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getDest());

    bool isSecretScalar = succeeded(scalarLayoutResult);
    bool isSecretDest = succeeded(destLayoutResult);

    if (isSecretScalar && isSecretDest) {
      return secretScalarSecretTensor(op, adaptor, rewriter);
    }

    if (!isSecretScalar && isSecretDest) {
      return cleartextScalarSecretTensor(op, adaptor, rewriter);
    }

    if (isSecretScalar && !isSecretDest) {
      return op.emitError() << "dest tensor should have been assigned a layout "
                               "by layout-propagation";
    }

    // cleartext scalar and cleartext tensor means this is a cleartext op
    // that can be elided.
    setMaterializedAttr(op);
    return success();
  }
};

class ConvertTensorExtractSlice
    : public ContextAwareOpConversionPattern<tensor::ExtractSliceOp> {
 public:
  using ContextAwareOpConversionPattern<
      tensor::ExtractSliceOp>::ContextAwareOpConversionPattern;

  LogicalResult secretSourceSecretResult(
      tensor::ExtractSliceOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const {
    MLIRContext* ctx = op.getContext();

    FailureOr<Attribute> sourceLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getSource());
    FailureOr<Attribute> resultLayoutResult =
        getTypeConverter()->getContextualAttr(op.getResult());
    LayoutAttr resultLayout = cast<LayoutAttr>(resultLayoutResult.value());
    IntegerRelation sourceRel =
        cast<LayoutAttr>(sourceLayoutResult.value()).getIntegerRelation();
    IntegerRelation resultRel = resultLayout.getIntegerRelation();

    // Compute the layout relation of the extract_slice operation.
    auto extractSliceLayout =
        getSliceExtractionRelation(op.getSourceType(), op.getResultType(),
                                   SmallVector<int64_t>(op.getStaticOffsets()),
                                   SmallVector<int64_t>(op.getStaticSizes()),
                                   SmallVector<int64_t>(op.getStaticStrides()));
    if (failed(extractSliceLayout)) {
      return op.emitError() << "failed to get layout for extract slice";
    }

    // Remap the source ciphertext semantic tensor to the result ciphertext
    // semantic tensor layout. To do this, we compose the relations to traverse
    // the following diagram, starting from the source ciphertext tensor to the
    // extracted slice ciphertext tensor.
    //  Source tensor ─────────> Slice tensor
    //     \                        │
    //    /│\                       │
    //     │                       \│ /
    //     │                        \/
    // Source ciphertext        Slice ciphertext
    // (ct, slot)               (ct, slot)
    sourceRel.inverse();
    sourceRel.compose(extractSliceLayout.value());
    sourceRel.compose(resultRel);

    LayoutAttr sliceLayoutAttr =
        LayoutAttr::getFromIntegerRelation(ctx, sourceRel);
    auto resultCiphertextSemanticType = cast<RankedTensorType>(
        getTypeConverter()->convertType(op.getResultType(), resultLayout));
    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
    auto remapAndExtract = remapAndExtractResult(
        b, adaptor.getSource(), sliceLayoutAttr, resultCiphertextSemanticType);

    setMaterializedAttr(remapAndExtract);
    setAttributeAssociatedWith(remapAndExtract->getResult(0), kLayoutAttrName,
                               sliceLayoutAttr);
    rewriter.replaceOp(op, remapAndExtract->getResult(0));
    return success();
  }

  LogicalResult matchAndRewrite(
      tensor::ExtractSliceOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    // If this is already materialized, then we can skip it.
    if (hasMaterializedAttr(op)) return failure();

    // Extract a secret slice from a secret tensor.
    FailureOr<Attribute> sourceLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getSource());
    FailureOr<Attribute> resultLayoutResult =
        getTypeConverter()->getContextualAttr(op.getResult());

    bool isSecretSource = succeeded(sourceLayoutResult);
    bool isSecretResult = succeeded(resultLayoutResult);

    if (isSecretSource && isSecretResult) {
      return secretSourceSecretResult(op, adaptor, rewriter);
    }

    if (isSecretSource && !isSecretResult) {
      return op.emitError()
             << "result tensor should have been assigned a layout "
                "by layout-propagation";
    }

    if (!isSecretSource && isSecretResult) {
      return op.emitError()
             << "source tensor should have been assigned a layout "
                "by layout-propagation";
    }

    // cleartext scalar and cleartext tensor means this is a cleartext op
    // that can be elided.
    setMaterializedAttr(op);
    return success();
  }
};

class ConvertTensorCollapseShape
    : public ContextAwareOpConversionPattern<tensor::CollapseShapeOp> {
 public:
  using ContextAwareOpConversionPattern<
      tensor::CollapseShapeOp>::ContextAwareOpConversionPattern;

  LogicalResult matchAndRewrite(
      tensor::CollapseShapeOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    // if the layouts are equal (modulo the rank reduction) then this can be a
    // no-op. check if the layouts are equal after collapsing the dimensions.
    // then if the input and output types are the same, just replace with input.
    SliceVerificationResult res =
        isRankReducedType(op.getSrcType(), op.getResultType());
    if (res != SliceVerificationResult::Success)
      return rewriter.notifyMatchFailure(
          op, "Only rank-reduced types are supported for CollapseShapeOp");

    auto srcType = adaptor.getSrc().getType();
    FailureOr<Attribute> tensorLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getSrc());
    if (failed(tensorLayoutResult)) {
      // If the tensor has no layout, it is a cleartext operation and
      // can be skipped.
      setMaterializedAttr(op);
      return success();
    }
    auto tensorLayout = dyn_cast<LayoutAttr>(tensorLayoutResult.value());
    if (!tensorLayout) {
      return op.emitError() << "failed to fetch new layout attribute for input";
    }
    FailureOr<Attribute> resultLayoutResult =
        getTypeConverter()->getContextualAttr(op.getResult());
    if (failed(resultLayoutResult)) {
      return op.emitError() << "failed to fetch layout attribute for input";
    }
    LayoutAttr resultLayout = dyn_cast<LayoutAttr>(resultLayoutResult.value());
    if (!resultLayout) {
      return op.emitError() << "failed to fetch new layout attribute for input";
    }
    Type resultType =
        getTypeConverter()->convertType(op.getResultType(), resultLayout);

    if (resultType != srcType) {
      return rewriter.notifyMatchFailure(
          op, "result type does not match input type");
    }

    auto srcRelation = tensorLayout.getIntegerRelation();
    auto collapsedRelation = collapseDimensions(srcRelation, op.getSrcType(),
                                                op.getReassociationIndices());
    if (!isRelationEqual(collapsedRelation,
                         resultLayout.getIntegerRelation())) {
      return rewriter.notifyMatchFailure(
          op, "result layout is not equal to input layout");
    }

    // Put in a no-op unrealized conversion cast operation to persist the new
    // attribute for downstream ops.
    auto castOp = UnrealizedConversionCastOp::create(
        rewriter, op.getLoc(), resultType, adaptor.getSrc());
    setMaterializedAttr(castOp);
    setAttributeAssociatedWith(castOp.getResult(0), kLayoutAttrName,
                               resultLayout);

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

class ConvertTensorExpandShape
    : public ContextAwareOpConversionPattern<tensor::ExpandShapeOp> {
 public:
  using ContextAwareOpConversionPattern<
      tensor::ExpandShapeOp>::ContextAwareOpConversionPattern;

  LogicalResult matchAndRewrite(
      tensor::ExpandShapeOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    // if the layouts are equal (modulo the rank reduction) then this can be a
    // no-op. check if the layouts are equal after collapsing the dimensions.
    // then if the input and output types are the same, just replace with input.
    SliceVerificationResult res =
        isRankReducedType(op.getResultType(), op.getSrcType());
    if (res != SliceVerificationResult::Success)
      return rewriter.notifyMatchFailure(
          op, "Only rank-reduced types are supported for CollapseShapeOp");

    auto srcType = adaptor.getSrc().getType();
    FailureOr<Attribute> sourceLayoutResult =
        getTypeConverter()->getContextualAttr(adaptor.getSrc());
    if (failed(sourceLayoutResult)) {
      // If the tensor has no layout, it is a cleartext operation and
      // can be skipped.
      setMaterializedAttr(op);
      return success();
    }
    auto sourceLayout = dyn_cast<LayoutAttr>(sourceLayoutResult.value());
    if (!sourceLayout) {
      return op.emitError() << "failed to fetch new layout attribute for input";
    }

    FailureOr<Attribute> resultLayoutResult =
        getTypeConverter()->getContextualAttr(op.getResult());
    if (failed(resultLayoutResult)) {
      // The result may not have a layout if the source was a plaintext. Since
      // that is checked above, this should not happen.
      return op.emitError() << "failed to fetch layout attribute for result";
    }
    LayoutAttr resultLayout = dyn_cast<LayoutAttr>(resultLayoutResult.value());
    if (!resultLayout) {
      return op.emitError()
             << "failed to fetch new layout attribute for result";
    }
    Type resultType =
        getTypeConverter()->convertType(op.getResultType(), resultLayout);

    if (resultType != srcType) {
      return rewriter.notifyMatchFailure(
          op, "result type does not match input type");
    }

    auto srcRelation = sourceLayout.getIntegerRelation();
    auto expandedRelation = expandDimensions(srcRelation, op.getResultType(),
                                             op.getReassociationIndices());
    if (!isRelationEqual(expandedRelation, resultLayout.getIntegerRelation())) {
      return rewriter.notifyMatchFailure(
          op, "result layout is not equal to input layout");
    }

    // Put in a no-op unrealized conversion cast operation to persist the new
    // attribute for downstream ops.
    auto castOp = UnrealizedConversionCastOp::create(
        rewriter, op.getLoc(), resultType, adaptor.getSrc());
    setMaterializedAttr(castOp);
    setAttributeAssociatedWith(castOp.getResult(0), kLayoutAttrName,
                               resultLayout);

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

struct DropRotateUnitDims : OpRewritePattern<tensor_ext::RotateOp> {
  using OpRewritePattern<tensor_ext::RotateOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(tensor_ext::RotateOp rotateOp,
                                PatternRewriter& rewriter) const override {
    SmallVector<int64_t> operandUnitDims =
        getUnitDims(rotateOp.getTensor().getType());
    if (operandUnitDims.empty()) {
      LLVM_DEBUG(llvm::dbgs() << "no unit dims to drop\n");
      return failure();
    }

    SmallVector<Value> collapsedOperands =
        collapseOperands(rewriter, {rotateOp.getTensor()}, operandUnitDims);

    tensor_ext::RotateOp collapsedOp = tensor_ext::RotateOp::create(
        rewriter, rotateOp.getLoc(), collapsedOperands[0], rotateOp.getShift());
    rewriter.replaceOp(rotateOp, expandResult(rewriter, collapsedOp.getResult(),
                                              rotateOp.getOutput().getType(),
                                              operandUnitDims));
    return success();
  }
};

struct DropRotateAndReduceUnitDims
    : OpRewritePattern<tensor_ext::RotateAndReduceOp> {
  using OpRewritePattern<tensor_ext::RotateAndReduceOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(tensor_ext::RotateAndReduceOp rotateOp,
                                PatternRewriter& rewriter) const override {
    SmallVector<int64_t> operandUnitDims =
        getUnitDims(rotateOp.getTensor().getType());
    if (operandUnitDims.empty()) {
      LLVM_DEBUG(llvm::dbgs() << "no unit dims to drop\n");
      return failure();
    }

    SmallVector<Value> collapsedOperands =
        collapseOperands(rewriter, {rotateOp.getTensor()}, operandUnitDims);

    auto collapsedOp = tensor_ext::RotateAndReduceOp::create(
        rewriter, rotateOp.getLoc(), collapsedOperands[0],
        rotateOp.getPlaintexts(), rotateOp.getPeriod().getZExtValue(),
        rotateOp.getSteps().getZExtValue(), rotateOp.getReduceOp());
    rewriter.replaceOp(rotateOp, expandResult(rewriter, collapsedOp.getResult(),
                                              rotateOp.getOutput().getType(),
                                              operandUnitDims));
    return success();
  }
};

struct ConvertLinalgMatmul
    : public ContextAwareOpConversionPattern<linalg::MatmulOp> {
 public:
  using ContextAwareOpConversionPattern<
      linalg::MatmulOp>::ContextAwareOpConversionPattern;

  ConvertLinalgMatmul(
      const ContextAwareTypeConverter& contextAwareTypeConverter,
      MLIRContext* context)
      : ContextAwareOpConversionPattern(contextAwareTypeConverter, context,
                                        /*benefit=*/10) {}

  LayoutAttr getLayoutAttr(Value value) const {
    auto layoutLookup = getTypeConverter()->getContextualAttr(value);
    if (failed(layoutLookup)) {
      return nullptr;
    }
    return dyn_cast<LayoutAttr>(layoutLookup.value());
  }

  bool supportsBicyclicDiagonal(linalg::MatmulOp op) const {
    auto kernelAttr = op->getAttrOfType<secret::KernelAttr>(
        secret::SecretDialect::kKernelAttrName);
    return kernelAttr &&
           kernelAttr.getName() == KernelName::MatmulBicyclicDiagonal;
  }

  // Bicyclic matmul against a cleartext operand: the cleartext operand arrives
  // packed as the generalized diagonal matrix (one row per summation step) and
  // the kernel is a BSGS rotate-and-reduce of the secret operand over n steps,
  // with rotation period m (the secret lhs's packed row count) or p (the secret
  // rhs's packed column count).
  void bicyclicDiagonalKernel(
      linalg::MatmulOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const {
    LLVM_DEBUG(llvm::dbgs()
               << "Converting linalg.matmul op with bicyclic diagonal "
                  "kernel: "
               << op << "\n");

    // Determine if the lhs or rhs is the secret operand.
    bool secretLhs = false;
    auto genericOp = op->getParentOfType<secret::GenericOp>();
    if (genericOp) {
      if (auto blockArg = dyn_cast<BlockArgument>(op.getInputs()[0])) {
        if (blockArg.getArgNumber() < genericOp.getNumOperands()) {
          auto opArg = genericOp.getOperand(blockArg.getArgNumber());
          if (isa<secret::SecretType>(opArg.getType())) {
            secretLhs = true;
          }
        }
      }
    }

    TypedValue<RankedTensorType> ct = cast<TypedValue<RankedTensorType>>(
        adaptor.getInputs()[secretLhs ? 0 : 1]);
    SSAValue ctLeaf(ct);
    TypedValue<RankedTensorType> pt = cast<TypedValue<RankedTensorType>>(
        adaptor.getInputs()[secretLhs ? 1 : 0]);
    SSAValue ptLeaf(pt);

    auto lhsType = cast<RankedTensorType>(op.getInputs()[0].getType());
    auto rhsType = cast<RankedTensorType>(op.getInputs()[1].getType());
    auto secretType = ct.getType();

    auto dagType = kernel::mlirTypeToDagType(secretType);
    int64_t period = secretLhs ? lhsType.getDimSize(0) : rhsType.getDimSize(1);
    int64_t steps = secretLhs ? lhsType.getDimSize(1) : rhsType.getDimSize(0);
    std::string reduceOp = isa<FloatType>(secretType.getElementType())
                               ? "arith.addf"
                               : "arith.addi";

    std::shared_ptr<ArithmeticDagNode<SSAValue>> implementedKernel =
        implementRotateAndReduce(ctLeaf, std::optional<SSAValue>(ptLeaf),
                                 period, steps, dagType, /*zeroDiagonals=*/{},
                                 reduceOp);

    rewriter.setInsertionPointAfter(op);
    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
    IRMaterializingVisitor visitor(ct.getType(), [&](Operation* createdOp) {
      setMaterializedAttr(createdOp);
    });
    Value finalOutput = visitor.process(implementedKernel, b)[0];

    auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));
    auto* finalOutputOp = finalOutput.getDefiningOp();
    finalOutputOp->setAttr(kLayoutAttrName, layoutAttr);
    setMaterializedAttr(finalOutputOp);

    // Add the initial accumulator value.
    Value result = adaptor.getOutputs()[0];

    Operation* addBias =
        makeAppropriatelyTypedAddOp(b, op->getLoc(), finalOutput, result);
    addBias->setAttr(kLayoutAttrName, layoutAttr);
    setMaterializedAttr(addBias);
    rewriter.replaceOp(op, addBias);
  }

  bool supportsBicyclic(linalg::MatmulOp op, OpAdaptor adaptor) const {
    auto kernelAttr = op->getAttrOfType<secret::KernelAttr>(
        secret::SecretDialect::kKernelAttrName);
    return kernelAttr && kernelAttr.getName() == KernelName::MatmulBicyclic;
  }

  void bicyclicKernel(linalg::MatmulOp op, OpAdaptor adaptor,
                      ContextAwareConversionPatternRewriter& rewriter) const {
    LLVM_DEBUG(llvm::dbgs()
               << "Converting linalg.matmul op with bicyclic kernel: " << op
               << "\n");

    TypedValue<RankedTensorType> lhs =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[0]);
    SSAValue lhsLeaf(lhs);
    TypedValue<RankedTensorType> rhs =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[1]);
    SSAValue rhsLeaf(rhs);

    auto lhsType = cast<RankedTensorType>(op.getInputs()[0].getType());
    auto rhsType = cast<RankedTensorType>(op.getInputs()[1].getType());

    auto dagType = kernel::mlirTypeToDagType(lhsType);
    // TODO(#2386): zero-pad inputs to ensure coprime dimensions
    std::shared_ptr<ArithmeticDagNode<SSAValue>> implementedKernel =
        kernel::implementBicyclicMatmul(lhsLeaf, rhsLeaf, lhsType.getShape()[0],
                                        lhsType.getShape()[1],
                                        rhsType.getShape()[1], dagType);

    rewriter.setInsertionPointAfter(op);
    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
    IRMaterializingVisitor visitor(
        lhs.getType(), [&](Operation* createdOp) { setMaterializedAttr(op); });
    Value finalOutput = visitor.process(implementedKernel, b)[0];

    auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));
    auto* finalOutputOp = finalOutput.getDefiningOp();
    finalOutputOp->setAttr(kLayoutAttrName, layoutAttr);
    setMaterializedAttr(finalOutputOp);

    // Add the initial accumulator value.
    Value result = adaptor.getOutputs()[0];
    Operation* addBias =
        makeAppropriatelyTypedAddOp(b, op->getLoc(), finalOutput, result);
    addBias->setAttr(kLayoutAttrName, layoutAttr);
    setMaterializedAttr(addBias);
    rewriter.replaceOp(op, addBias);
  }

  LogicalResult matchAndRewrite(
      linalg::MatmulOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    if (supportsBicyclicDiagonal(op)) {
      bicyclicDiagonalKernel(op, adaptor, rewriter);
      return success();
    }
    if (supportsBicyclic(op, adaptor)) {
      bicyclicKernel(op, adaptor, rewriter);
      return success();
    }
    return failure();
  }
};

struct ConvertLinalgBatchMatmul
    : public ContextAwareOpConversionPattern<linalg::BatchMatmulOp> {
 public:
  using ContextAwareOpConversionPattern<
      linalg::BatchMatmulOp>::ContextAwareOpConversionPattern;

  LayoutAttr getLayoutAttr(Value value) const {
    auto layoutLookup = getTypeConverter()->getContextualAttr(value);
    if (failed(layoutLookup)) {
      return nullptr;
    }
    return dyn_cast<LayoutAttr>(layoutLookup.value());
  }

  bool supportsTricyclic(linalg::BatchMatmulOp op, OpAdaptor adaptor) const {
    auto kernelAttr = op->getAttrOfType<secret::KernelAttr>(
        secret::SecretDialect::kKernelAttrName);
    return kernelAttr &&
           kernelAttr.getName() == KernelName::BatchMatmulTricyclic;
  }

  void tricyclicKernel(linalg::BatchMatmulOp op, OpAdaptor adaptor,
                       ContextAwareConversionPatternRewriter& rewriter) const {
    LLVM_DEBUG(llvm::dbgs()
               << "Converting linalg.batch_matmul op with tricyclic kernel: "
               << op << "\n");

    TypedValue<RankedTensorType> lhs =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[0]);
    SSAValue lhsLeaf(lhs);
    TypedValue<RankedTensorType> rhs =
        cast<TypedValue<RankedTensorType>>(adaptor.getInputs()[1]);
    SSAValue rhsLeaf(rhs);

    auto lhsType = cast<RankedTensorType>(op.getInputs()[0].getType());
    auto rhsType = cast<RankedTensorType>(op.getInputs()[1].getType());

    auto dagType = kernel::mlirTypeToDagType(lhsType);
    std::shared_ptr<ArithmeticDagNode<SSAValue>> implementedKernel =
        kernel::implementTricyclicBatchMatmul(
            lhsLeaf, rhsLeaf, lhsType.getShape()[0], lhsType.getShape()[1],
            lhsType.getShape()[2], rhsType.getShape()[2], dagType);

    rewriter.setInsertionPointAfter(op);
    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
    IRMaterializingVisitor visitor(
        lhs.getType(), [&](Operation* createdOp) { setMaterializedAttr(op); });
    Value finalOutput = visitor.process(implementedKernel, b)[0];

    auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));
    auto* finalOutputOp = finalOutput.getDefiningOp();
    finalOutputOp->setAttr(kLayoutAttrName, layoutAttr);
    setMaterializedAttr(finalOutputOp);

    // Add the initial accumulator value.
    Value result = adaptor.getOutputs()[0];
    Operation* addBias =
        makeAppropriatelyTypedAddOp(b, op->getLoc(), finalOutput, result);
    addBias->setAttr(kLayoutAttrName, layoutAttr);
    setMaterializedAttr(addBias);
    rewriter.replaceOp(op, addBias);
  }

  LogicalResult matchAndRewrite(
      linalg::BatchMatmulOp op, OpAdaptor adaptor,
      ContextAwareConversionPatternRewriter& rewriter) const final {
    if (supportsTricyclic(op, adaptor)) {
      tricyclicKernel(op, adaptor, rewriter);
      return success();
    }
    return failure();
  }
};

struct ConvertToCiphertextSemantics
    : impl::ConvertToCiphertextSemanticsBase<ConvertToCiphertextSemantics> {
  using ConvertToCiphertextSemanticsBase::ConvertToCiphertextSemanticsBase;

  void runOnOperation() override {
    MLIRContext* context = &getContext();
    auto* module = getOperation();

    int64_t ctSize = minSlotCount;
    LayoutMaterializationTypeConverter typeConverter =
        LayoutMaterializationTypeConverter(ctSize);

    RewritePatternSet patterns(context);
    ConversionTarget target(*context);
    target.markUnknownOpDynamicallyLegal([&](Operation* op) {
      return isa<ModuleOp>(op) || hasMaterializedAttr(op);
    });

    patterns.add<
        ConvertAnyAddingMaterializedAttr, ConvertConvertLayout, ConvertFunc,
        ConvertLinalgMatmul, ConvertLinalgBatchMatmul, ConvertLinalgReduce,
        ConvertLinalgDot, ConvertSecretGeneric, ConvertTensorCollapseShape,
        ConvertTensorExpandShape, ConvertTensorExtractLayout,
        ConvertTensorExtractSlice, ConvertTensorPad, ConvertTensorInsertLayout,
        ConvertTensorInsertSlice, PreserveLinalgMatvecAsLinearTransform>(
        typeConverter, context);
    patterns.add<ConvertLinalgMatvecLayout, ConvertLinalgConv1D,
                 ConvertLinalgConv2D, ConvertLinalgConv2DNchwFchw,
                 ConvertLinalgConv1DNcwFcw>(typeConverter, context,
                                            unrollKernels);
    patterns.add<ConvertAssignLayout>(typeConverter, context, minSlotCount,
                                      codegenStrategy);

    ConversionConfig config;
    config.buildMaterializations = false;
    if (failed(applyContextAwarePartialConversion(
            module, target, std::move(patterns), config))) {
      return signalPassFailure();
    }

    // Decompose tensor.concat into repeated tensor.insert_slice ops.
    // Note ConvertAssignLayout generates tensor.concat
    RewritePatternSet cleanupPatterns2(context);
    tensor::populateDecomposeTensorConcatPatterns(cleanupPatterns2);

    // Folding here will remove any unrealized conversion cast ops that were
    // inserted to persist new layouts.
    if (failed(applyPatternsGreedily(module, std::move(cleanupPatterns2)))) {
      return signalPassFailure();
    }

    // Walk the IR to validate that there are no remaining unrealized conversion
    // cast ops.
    module->walk([&](UnrealizedConversionCastOp op) {
      op->emitError() << "unexpected unrealized conversion cast op found";
      signalPassFailure();
    });

    clearAttrs(module, kLayoutAttrName);
    clearAttrs(module, kMaterializedAttrName);
  }
};

}  // namespace heir
}  // namespace mlir
