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

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <numeric>
#include <optional>
#include <string>
#include <utility>

#include "lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h"
#include "lib/Dialect/HEIRInterfaces.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/Dialect/TensorExt/Transforms/Patterns.h"
#include "lib/Kernel/KernelName.h"
#include "lib/Transforms/LayoutPropagation/Utils.h"
#include "lib/Utils/AttributeUtils.h"
#include "lib/Utils/Layout/Convolution.h"
#include "lib/Utils/Layout/Hoisting.h"
#include "lib/Utils/Layout/IslConversion.h"
#include "lib/Utils/Layout/Utils.h"
#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/SmallVectorExtras.h"       // from @llvm-project
#include "llvm/include/llvm/ADT/TypeSwitch.h"              // from @llvm-project
#include "llvm/include/llvm/Support/Debug.h"               // from @llvm-project
#include "llvm/include/llvm/Support/ErrorHandling.h"       // from @llvm-project
#include "llvm/include/llvm/Support/FormatVariadic.h"      // from @llvm-project
#include "mlir/include/mlir/Analysis/DataFlow/Utils.h"     // from @llvm-project
#include "mlir/include/mlir/Analysis/DataFlowFramework.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/AsmParser/AsmParser.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.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/Linalg/IR/LinalgInterfaces.h"  // from @llvm-project
#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.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/Builders.h"               // from @llvm-project
#include "mlir/include/mlir/IR/BuiltinAttributes.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/Diagnostics.h"            // from @llvm-project
#include "mlir/include/mlir/IR/Matchers.h"               // from @llvm-project
#include "mlir/include/mlir/IR/Operation.h"              // from @llvm-project
#include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
#include "mlir/include/mlir/IR/Types.h"                  // from @llvm-project
#include "mlir/include/mlir/IR/Value.h"                  // from @llvm-project
#include "mlir/include/mlir/IR/Visitors.h"               // from @llvm-project
#include "mlir/include/mlir/Support/LLVM.h"              // from @llvm-project
#include "mlir/include/mlir/Support/WalkResult.h"        // from @llvm-project
#include "mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project

#define DEBUG_TYPE "layout-propagation"

namespace mlir {
namespace heir {

using linalg::BatchMatmulOp;
using linalg::Conv1DNcwFcwOp;
using linalg::Conv1DOp;
using linalg::Conv2DNchwFchwOp;
using linalg::Conv2DOp;
using linalg::DotOp;
using linalg::MatmulOp;
using linalg::MatvecOp;
using linalg::ReduceOp;
using linalg::VecmatOp;
using presburger::IntegerRelation;
using secret::GenericOp;
using secret::SecretType;
using secret::YieldOp;
using tensor::CollapseShapeOp;
using tensor::ExpandShapeOp;
using tensor::InsertOp;
using tensor_ext::AssignLayoutOp;
using tensor_ext::ConvertLayoutOp;
using tensor_ext::LayoutAttr;

#define GEN_PASS_DEF_LAYOUTPROPAGATION
#include "lib/Transforms/LayoutPropagation/LayoutPropagation.h.inc"

namespace {

// The result of a compatibility check for the layouts of an op's operands (cf.
// hasCompatibleArgumentLayouts). If the check fails, the presence of a
// diagnostic signals that the failure is unrecoverable and should cause the
// pass to fail. If the diagnostic is nullopt, then the failure can be
// recovered by rectifyIncompatibleOperandLayouts.
struct CompatibilityResult {
  bool compatible;
  std::optional<InFlightDiagnostic> diag;
};

void visitDebugInfo(Operation* op, DataFlowSolver* solver) {
  LLVM_DEBUG({
    llvm::dbgs() << "Visiting: " << op->getName() << "\n";
    llvm::dbgs() << "Operand secretness:\n";
    for (OpOperand& opOperand : op->getOpOperands()) {
      bool secretness = isSecret(opOperand.get(), solver);
      llvm::dbgs() << "- Operand " << opOperand.getOperandNumber() << ": "
                   << secretness << "\n";
    }
  });
}

void debugAssignLayout(Value value, Attribute layout) {
  LLVM_DEBUG(llvm::dbgs() << "Assigning layout " << layout << " to value "
                          << value << "\n");
}

std::pair<Value, LayoutAttr> convertToLayout(
    MLIRContext* ctx, mlir::IRRewriter& builder, Operation* op, Value value,
    LayoutAttr oldLayout, const IntegerRelation& newRelation,
    ArrayRef<int64_t> domainSchedule = {}) {
  builder.setInsertionPoint(op);
  LayoutAttr layoutAttr = LayoutAttr::getFromIntegerRelation(ctx, newRelation);
  ConvertLayoutOp convertLayoutOp = ConvertLayoutOp::create(
      builder, op->getLoc(), value, oldLayout, layoutAttr,
      builder.getDenseI64ArrayAttr(domainSchedule));
  convertLayoutOp->setAttr(tensor_ext::TensorExtDialect::kLayoutAttrName,
                           layoutAttr);
  Value toReplace = convertLayoutOp.getResult();
  builder.replaceUsesWithIf(value, toReplace, [&](OpOperand& operand) {
    return operand.getOwner() == op;
  });
  return std::make_pair(toReplace, layoutAttr);
}

// The outcome of folding a zero `tensor.pad` on the width dim into a 1-D
// conv's own `padding` parameter.
struct FoldedConvPadding {
  // What the Toeplitz matrix must be built against.
  ConvMatrixOperand matrixOperand;
  // The layout the padded operand must already carry for the fold to be valid.
  IntegerRelation targetRelation;
};

// Try to fold a zero `tensor.pad` on the width dim of a conv's `data` operand
// into the conv's own `padding` parameter.
//
// Returns nullopt when the pattern does not apply, in which case the caller
// keeps the unfolded path. `dataType` must be rank 3 with N=1.
std::optional<FoldedConvPadding> tryFoldPadIntoConvPadding(
    Value data, RankedTensorType dataType, LayoutAttr dataLayout,
    int64_t ciphertextSize) {
  assert(dataType.getRank() == 3 && dataType.getDimSize(0) == 1 &&
         "expected a rank-3 N=1 conv data operand");

  // DropUnitDims rewrites a rank-3 pad into
  // collapse_shape -> tensor.pad (rank 2) -> expand_shape, so peel any
  // reshape/cast chain to find the pad.
  Value cursor = data;
  tensor::PadOp padOp;
  while (Operation* def = cursor.getDefiningOp()) {
    if (auto p = dyn_cast<tensor::PadOp>(def)) {
      padOp = p;
      break;
    }
    if (isa<tensor::ExpandShapeOp, tensor::CollapseShapeOp, tensor::CastOp>(
            def)) {
      cursor = def->getOperand(0);
      continue;
    }
    break;
  }
  if (!padOp) return std::nullopt;

  Value padValue = padOp.getConstantPaddingValue();
  bool zeroPad = padValue && (matchPattern(padValue, m_AnyZeroFloat()) ||
                              matchPattern(padValue, m_Zero()));
  if (!zeroPad) return std::nullopt;

  // Only a symmetric pad on the trailing (width) dim is expressible as the
  // conv's `padding` parameter, and only when the bounds are static. The pad
  // may be rank 2 (unit batch dim dropped) or rank 3.
  if (!padOp.getLow().empty() || !padOp.getHigh().empty()) return std::nullopt;
  ArrayRef<int64_t> low = padOp.getStaticLow();
  ArrayRef<int64_t> high = padOp.getStaticHigh();
  bool widthOnly = !low.empty() && low.size() == high.size() &&
                   low.back() == high.back() && low.back() > 0;
  for (size_t i = 0; widthOnly && i + 1 < low.size(); ++i) {
    widthOnly &= low[i] == 0 && high[i] == 0;
  }
  if (!widthOnly) return std::nullopt;

  // The reshape chain above means the pad's width dim is not guaranteed to be
  // this conv operand's width dim, so validate before building a type from it.
  int64_t p = low.back();
  std::optional<ConvMatrixOperand> matrixOperand =
      foldConvWidthPadding(dataType, p);
  if (!matrixOperand) return std::nullopt;

  // The layout we expect on the padded value: the unpadded row-major layout
  // with the width index shifted by `p`. If the actual layout is anything else
  // (a conversion intervened, a non-row-major producer, reshapes that did not
  // cancel) do not fold, rather than silently mis-indexing the matrix.
  IntegerRelation expected =
      getRowMajorLayoutRelation(matrixOperand->dataType, ciphertextSize);
  expected = shiftVar(
      expected, expected.getVarKindOffset(presburger::VarKind::Domain) + 2, p);
  if (!dataLayout.getIntegerRelation().isEqual(expected)) {
    LLVM_DEBUG(llvm::dbgs()
               << "conv_1d found a pad of " << p
               << " but the operand layout does not match the shifted "
                  "unpadded row-major layout; not folding\n");
    return std::nullopt;
  }

  LLVM_DEBUG(llvm::dbgs() << "conv_1d folding tensor.pad of " << p
                          << " into the conv padding parameter\n");
  return FoldedConvPadding{*matrixOperand, expected};
}

// Return a copy of the kernel info associated with the value and update the
// result shape to the new result shape. If the value does not have a kernel
// info, return an empty Attribute.

Attribute cloneKernelInfoWithResultShape(Value value,
                                         ArrayRef<int64_t> resultShape) {
  auto kernelInfo = findAttributeAssociatedWith(value, kKernelInfoAttrName);
  if (succeeded(kernelInfo)) {
    auto info = getKernelInfo(kernelInfo.value()).value();
    info.resultShape = llvm::to_vector(resultShape);
    return makeKernelInfoAttr(value.getContext(), info);
  }
  return Attribute{};
}

}  // namespace

struct LayoutPropagation : impl::LayoutPropagationBase<LayoutPropagation> {
  using LayoutPropagationBase::LayoutPropagationBase;

  // Top level visit method handling common logic and dispatching to specific
  // visitOperation overloads.
  LogicalResult visitOperation(Operation* op);

  // Op-specific transfer functions
  LogicalResult visitOperation(CollapseShapeOp op);
  LogicalResult visitOperation(ExpandShapeOp op);
  LogicalResult visitOperation(GenericOp op);
  LogicalResult visitOperation(ReduceOp op);
  LogicalResult visitOperation(Conv1DOp op);
  LogicalResult visitOperation(Conv1DNcwFcwOp op);
  LogicalResult visitOperation(Conv2DOp op);
  LogicalResult visitOperation(Conv2DNchwFchwOp op);
  LogicalResult visitOperation(VecmatOp op);
  LogicalResult visitOperation(MatvecOp op);
  LogicalResult visitOperation(MatmulOp op);
  LogicalResult visitOperation(BatchMatmulOp op);
  LogicalResult visitOperation(DotOp op);
  LogicalResult visitOperation(YieldOp op);
  LogicalResult visitOperation(affine::AffineForOp op);
  LogicalResult visitOperation(func::FuncOp op);
  LogicalResult visitOperation(func::ReturnOp op);
  LogicalResult visitOperation(tensor::ExtractOp op);
  LogicalResult visitOperation(tensor::InsertOp op);
  LogicalResult visitOperation(tensor::InsertSliceOp op);
  LogicalResult visitOperation(tensor::ExtractSliceOp op);
  LogicalResult visitOperation(tensor::PadOp op);

  // Determine if the operation arguments have compatible layouts for the
  // given op. If the check fails, the CompatibilityResult::compatible field
  // is false. If there is also a diagnostic populated in the result, the
  // failure is unrecoverable.
  CompatibilityResult hasCompatibleArgumentLayouts(Operation* op);

  // Op-specific compatibility functions
  CompatibilityResult hasCompatibleArgumentLayouts(DotOp op);
  CompatibilityResult hasCompatibleArgumentLayouts(Conv1DOp op);
  CompatibilityResult hasCompatibleArgumentLayouts(Conv1DNcwFcwOp op);
  CompatibilityResult hasCompatibleArgumentLayouts(Conv2DOp op);
  CompatibilityResult hasCompatibleArgumentLayouts(Conv2DNchwFchwOp op);
  CompatibilityResult hasCompatibleArgumentLayouts(ReduceOp op);
  CompatibilityResult hasCompatibleArgumentLayouts(VecmatOp op);
  CompatibilityResult hasCompatibleArgumentLayouts(MatvecOp op);
  CompatibilityResult hasCompatibleArgumentLayouts(MatmulOp op);
  CompatibilityResult hasCompatibleArgumentLayouts(tensor::InsertSliceOp op);

  // Insert conversion ops to rectify incompatible operand layouts
  void rectifyIncompatibleOperandLayouts(Operation* op);

  // Op-specific overrides
  void rectifyIncompatibleOperandLayouts(DotOp op);
  void rectifyIncompatibleOperandLayouts(ReduceOp op);
  void rectifyIncompatibleOperandLayouts(tensor::InsertOp op);
  void rectifyIncompatibleOperandLayouts(tensor::InsertSliceOp op);

  // Return the default layout for a given type
  FailureOr<LayoutAttr> defaultLayoutForType(Type type);
  FailureOr<LayoutAttr> defaultLayoutForScalarType(Type type);

  // Create an assign_layout op for the given value, and return the resulting
  // op. The given builder should have its insertion point set before calling.
  FailureOr<AssignLayoutOp> assignDefaultLayoutForOpOperand(
      Operation* op, Value operand, IRRewriter& builder);

  // The conv/matvec kernels add the init (bias) operand directly to the
  // kernel output, which is packed per the op's result layout — for strided
  // convs a pixel-shuffled "gap" layout, NOT the row-major layout the init
  // gets by default. Force the init operand's layout to match the result
  // layout (a free re-packing for the usual plaintext-constant bias).
  LogicalResult alignInitWithResultLayout(Operation* op, Value init,
                                          Attribute resultLayoutAttr);

  // Helper to pass layouts through generic ops
  void passLayoutThroughOp(Operation* op);

  // Add an op attribute denoting the layouts of the op results. Assumes the
  // assignedLayouts map contains the layout for the result SSA values already.
  void setResultLayoutAttr(Operation* op, Attribute kernelInfoAttr = {});

  void runOnOperation() override;

  LayoutAttr getComposedLayoutAttr(Value v) {
    Attribute attr = assignedLayouts.at(v);
    if (auto layoutAttr = dyn_cast<LayoutAttr>(attr)) {
      return layoutAttr;
    }
    if (auto arrayAttr = dyn_cast<ArrayAttr>(attr)) {
      auto layoutAttrs = arrayAttr.getAsRange<LayoutAttr>();
      auto it = layoutAttrs.begin();
      presburger::IntegerRelation composedRel = (*it).getIntegerRelation();
      ++it;
      for (; it != layoutAttrs.end(); ++it) {
        composedRel.compose((*it).getIntegerRelation());
      }
      return LayoutAttr::getFromIntegerRelation(&getContext(), composedRel);
    }
    llvm_unreachable("unexpected attribute type in assignedLayouts");
  }

  DenseMap<Value, Attribute> assignedLayouts;
  DataFlowSolver* solver;
};

FailureOr<AssignLayoutOp> LayoutPropagation::assignDefaultLayoutForOpOperand(
    Operation* op, Value operand, IRRewriter& builder) {
  FailureOr<LayoutAttr> layout = defaultLayoutForType(operand.getType());
  if (failed(layout)) {
    return failure();
  }
  LayoutAttr layoutAttr = layout.value();
  AssignLayoutOp assignLayoutOp =
      AssignLayoutOp::create(builder, op->getLoc(), operand, layoutAttr);
  setAttributeAssociatedWith(assignLayoutOp.getResult(),
                             tensor_ext::TensorExtDialect::kLayoutAttrName,
                             layoutAttr);
  Value toReplace = assignLayoutOp.getResult();
  // This may create duplicate layout assignment ops, and we expect CSE
  // to later clean them up. Otherwise we risk replacing a use of the
  // cleartext value in some other context.
  builder.replaceUsesWithIf(operand, toReplace, [&](OpOperand& otherOperand) {
    return otherOperand.getOwner() == op;
  });
  debugAssignLayout(toReplace, layoutAttr);
  assignedLayouts.insert({toReplace, layoutAttr});

  return assignLayoutOp;
}

LogicalResult LayoutPropagation::visitOperation(Operation* op) {
  visitDebugInfo(op, solver);

  if (!isa<func::FuncOp, func::ReturnOp, GenericOp, YieldOp, AssignLayoutOp>(
          op) &&
      !isSecret(op->getOperands(), solver) &&
      !isSecret(op->getResults(), solver)) {
    LLVM_DEBUG(llvm::dbgs()
               << "Skipping op " << op->getName()
               << " because its operands and results are non-secret, or it is "
                  "in a special allowlist of ops to ignore\n");
    return success();
  }

  // If an operand has no layout, it may for example be produced as a plaintext
  // constant, such as a zero-valued tensor for the initializer of a reduction.
  // In this case, we insert a layout assignment.
  for (OpOperand& opOperand : op->getOpOperands()) {
    Value operand = opOperand.get();
    if (!assignedLayouts.contains(operand)) {
      // Some operations may function properly with operands that have no
      // layout, e.g., a tensor_ext.rotate op doesn't need a layout for the
      // shift operand, and a tensor.insert op has a kernel that can work with a
      // cleartext scalar operand (and is, in fact, more efficient than when the
      // scalar operand is packed).
      if (auto layoutReqIface =
              dyn_cast<OperandLayoutRequirementOpInterface>(op)) {
        bool secretness = isSecret(opOperand.get(), solver);
        if (!layoutReqIface.operandRequiresLayout(opOperand.getOperandNumber(),
                                                  secretness)) {
          LLVM_DEBUG(
              llvm::dbgs()
              << "OperandLayoutRequirementOpInterface ensures us that operand "
              << opOperand.getOperandNumber() << " of op " << op->getName()
              << " does not need a layout.\n");
          continue;
        }
      }

      LLVM_DEBUG(llvm::dbgs() << "operand has no layout: " << operand << "\n");
      mlir::IRRewriter builder(&getContext());
      builder.setInsertionPoint(op);
      if (failed(assignDefaultLayoutForOpOperand(op, operand, builder)))
        return op->emitError()
               << "Failed to assign default layout to operand " << operand;
    }
  }

  auto [compatible, diag] = hasCompatibleArgumentLayouts(op);
  if (!compatible) {
    LLVM_DEBUG(llvm::dbgs() << "Argument layouts for " << op->getName()
                            << " are incompatible\n");
    if (diag.has_value()) {
      // An InFlightDiagnostic casts to a failure()
      return diag.value();
    }
    rectifyIncompatibleOperandLayouts(op);
  }

  return TypeSwitch<Operation*, LogicalResult>(op)
      // func ops
      .Case<func::FuncOp, func::ReturnOp>(
          [&](auto op) { return visitOperation(op); })
      // secret ops
      .Case<GenericOp, YieldOp>([&](auto op) { return visitOperation(op); })
      // linalg ops
      .Case<DotOp, MatvecOp, VecmatOp, ReduceOp, MatmulOp, BatchMatmulOp,
            Conv1DOp, Conv1DNcwFcwOp, Conv2DOp, Conv2DNchwFchwOp>(
          [&](auto op) { return visitOperation(op); })
      // affine ops
      .Case<affine::AffineForOp>([&](auto op) { return visitOperation(op); })
      // tensor ops
      .Case<tensor::ExtractOp, tensor::InsertOp, tensor::InsertSliceOp,
            tensor::ExtractSliceOp, tensor::PadOp, CollapseShapeOp,
            ExpandShapeOp>([&](auto op) { return visitOperation(op); })
      // AddI, AddF, mgmt.* all pass the layout through unchanged.
      .Default([&](Operation* op) {
        passLayoutThroughOp(op);
        return success();
      });
}

LogicalResult LayoutPropagation::visitOperation(func::FuncOp op) {
  // Set a default value for each secret argument
  for (Value arg : op.getArguments()) {
    if (!isSecret(arg, solver)) {
      // Cleartext arguments don't get layouts, they are later given
      // assign_layout ops and materialized to plaintexts server-side.
      continue;
    }
    FailureOr<LayoutAttr> layout = defaultLayoutForType(arg.getType());
    if (failed(layout)) {
      return op->emitOpError()
             << "Failed to assign default layout to func argument " << arg;
    }
    debugAssignLayout(arg, layout.value());
    assignedLayouts.insert({arg, layout.value()});
    setAttributeAssociatedWith(
        arg, tensor_ext::TensorExtDialect::kLayoutAttrName, layout.value());

    // Set a default kernel info for each secret argument.
    if (auto tensorType = dyn_cast<RankedTensorType>(
            cast<SecretType>(arg.getType()).getValueType())) {
      setAttributeAssociatedWith(
          arg, kKernelInfoAttrName,
          makeKernelInfoAttr(
              arg.getContext(),
              KernelInfo{.resultShape = llvm::to_vector(tensorType.getShape()),
                         .gapFactor = 1}));
    }
  }

  // Func result attrs are handled by the ReturnOp
  return success();
}

LogicalResult LayoutPropagation::visitOperation(func::ReturnOp op) {
  func::FuncOp func = op->getParentOfType<func::FuncOp>();
  for (OpOperand& operand : op->getOpOperands()) {
    if (!assignedLayouts.contains(operand.get())) {
      if (isSecret(operand.get(), solver)) {
        return op->emitError("secret return value has no assigned layout");
      }

      // It needs no layout.
      continue;
    }
    Attribute layout = assignedLayouts.at(operand.get());
    func.setResultAttr(operand.getOperandNumber(),
                       tensor_ext::TensorExtDialect::kLayoutAttrName, layout);
  }
  return success();
}

LogicalResult LayoutPropagation::visitOperation(GenericOp op) {
  // Every block argument has the same layout as its corresponding operand.
  for (OpOperand& operand : op->getOpOperands()) {
    if (!assignedLayouts.contains(operand.get())) {
      // Assume it is not a tensor type and doesn't need a layout.
      continue;
    }
    Attribute layout = assignedLayouts.at(operand.get());
    BlockArgument blockArg =
        op.getRegion().getArgument(operand.getOperandNumber());
    assignedLayouts.insert({blockArg, layout});
    setAttributeAssociatedWith(
        blockArg, tensor_ext::TensorExtDialect::kLayoutAttrName, layout);
    debugAssignLayout(blockArg, layout);
    // Pass through kernel info
    auto kernelInfo =
        findAttributeAssociatedWith(operand.get(), kKernelInfoAttrName);
    if (succeeded(kernelInfo)) {
      setAttributeAssociatedWith(blockArg, kKernelInfoAttrName,
                                 kernelInfo.value());
    }
  }
  // The layout of the result of the generic op is handled when the YieldOp is
  // visited.
  return success();
}

LogicalResult LayoutPropagation::visitOperation(YieldOp op) {
  // The results of the generic op has the same layouts as the yielded values
  GenericOp generic = op->getParentOfType<GenericOp>();
  for (OpOperand& operand : op->getOpOperands()) {
    Type operandType = operand.get().getType();
    if (!assignedLayouts.contains(operand.get())) {
      // If it's a tensor type, it may be something like a tensor.empty()
      // that would not be assigned a layout earlier in the walk, because
      // it does not depend on any secret information. In this case, use the
      // default layout.
      LLVM_DEBUG(llvm::dbgs() << "No layout assigned to operand "
                              << operand.get() << ", using default layout\n");
      if (isa<RankedTensorType>(operandType)) {
        FailureOr<LayoutAttr> layout = defaultLayoutForType(operandType);
        if (failed(layout)) {
          return failure();
        }
        LayoutAttr layoutAttr = layout.value();
        debugAssignLayout(operand.get(), layoutAttr);
        assignedLayouts.insert({operand.get(), layoutAttr});
      } else {
        // Assume it is not a tensor type and doesn't need a layout.
        continue;
      }
    }
    Attribute layout = assignedLayouts.at(operand.get());
    Value result = generic.getResult(operand.getOperandNumber());
    assignedLayouts.insert({result, layout});
    debugAssignLayout(result, layout);
    setAttributeAssociatedWith(
        result, tensor_ext::TensorExtDialect::kLayoutAttrName, layout);
  }
  return success();
}

LogicalResult LayoutPropagation::visitOperation(CollapseShapeOp op) {
  // Reassociate the dimensions to collapse. For e.g. if we collapse d0 x d1 x
  // d2 to d' then we need to construct the row-major relation from (d0, d1, d2)
  // -> d'. Then the relation on the collapsed tensor will be the composition
  // from d' -> (d0, d1, d2) -> (ct, slot).
  auto rowMajorRelation = getCollapsedRelation(
      op.getSrcType(), op.getResultType(), op.getReassociationIndices());

  auto tensor = op.getSrc();
  LayoutAttr inputLayout = getComposedLayoutAttr(tensor);
  IntegerRelation relation = inputLayout.getIntegerRelation();

  relation.applyDomain(rowMajorRelation);
  MLIRContext* ctx = &getContext();
  LayoutAttr resultLayoutAttr =
      LayoutAttr::getFromIntegerRelation(ctx, relation);
  Attribute kernelInfoAttr =
      cloneKernelInfoWithResultShape(tensor, op.getResultType().getShape());

  assignedLayouts.insert({op.getResult(), resultLayoutAttr});
  setResultLayoutAttr(op, kernelInfoAttr);
  debugAssignLayout(op.getResult(), resultLayoutAttr);
  return success();
}

LogicalResult LayoutPropagation::visitOperation(ExpandShapeOp op) {
  // Only support rank-reduced types for now, i.e., where the expanded shape
  // only adds static dimensions of size 1.
  SliceVerificationResult res =
      isRankReducedType(op.getResultType(), op.getSrcType());
  if (res != SliceVerificationResult::Success)
    return op->emitError(
        "Only rank-reduced types are supported for ExpandShapeOp");

  auto tensor = op.getSrc();
  LayoutAttr inputLayout = getComposedLayoutAttr(tensor);
  IntegerRelation relation = inputLayout.getIntegerRelation();
  IntegerRelation expandedRelation = expandDimensions(
      relation, op.getResultType(), op.getReassociationIndices());

  MLIRContext* ctx = &getContext();
  LayoutAttr resultLayoutAttr =
      LayoutAttr::getFromIntegerRelation(ctx, expandedRelation);
  Attribute kernelInfoAttr =
      cloneKernelInfoWithResultShape(tensor, op.getResultType().getShape());

  assignedLayouts.insert({op.getResult(), resultLayoutAttr});
  setResultLayoutAttr(op, kernelInfoAttr);
  debugAssignLayout(op.getResult(), resultLayoutAttr);
  return success();
}

LogicalResult LayoutPropagation::visitOperation(VecmatOp op) {
  auto vecmatOp = cast<linalg::ContractionOpInterface>(*op);
  auto vec = vecmatOp.lhs();
  auto vecType = cast<RankedTensorType>(vec.getType());
  auto matrix = vecmatOp.rhs();
  if (!isSecret(vec, solver) || isSecret(matrix, solver)) {
    return failure();
  }

  MLIRContext* ctx = &getContext();
  mlir::IRRewriter builder(ctx);

  if (vecType.getDimSize(0) > minSlotCount) {
    return op->emitError() << "Vector must fit into a single ciphertext";
  }

  // Assign a row major layout to the input vec.
  LayoutAttr vecLayout = getComposedLayoutAttr(vec);
  if (!isRelationRowMajor(vecType, minSlotCount,
                          vecLayout.getIntegerRelation())) {
    // Insert a layout conversion op to make the vec layout per-row
    auto [toReplace, newVecLayoutAttr] =
        convertToLayout(ctx, builder, op, vec, vecLayout,
                        getPerRowLayoutRelation(vecType, minSlotCount));
    debugAssignLayout(toReplace, newVecLayoutAttr);
    assignedLayouts.insert({toReplace, newVecLayoutAttr});
    vec = toReplace;
  }

  // Assign or change the filter matrix to a diagonal layout. For a vecmat
  // operation, we need the diagonals to be extracted from the transpose.
  // So we get the regular diagonal layout for the transpose matrix type, and
  // then swap the dimensions.
  auto matrixType = cast<RankedTensorType>(matrix.getType());
  auto matrixTransposeType = RankedTensorType::get(
      {matrixType.getDimSize(1), matrixType.getDimSize(0)},
      matrixType.getElementType());
  LayoutAttr matrixLayout = getComposedLayoutAttr(matrix);
  IntegerRelation matrixLayoutRelation = matrixLayout.getIntegerRelation();
  auto domainOffset = matrixLayout.getIntegerRelation().getVarKindOffset(
      presburger::VarKind::Domain);
  auto clonedMatrixRelation = matrixLayoutRelation.clone();
  clonedMatrixRelation->swapVar(domainOffset, domainOffset + 1);
  if (!isRelationSquatDiagonal(matrixTransposeType, minSlotCount,
                               *clonedMatrixRelation)) {
    // Insert a layout conversion op to make the matrix layout squat diagonal
    auto [toReplace, newMatrixLayoutAttr] = convertToLayout(
        ctx, builder, op, matrix, matrixLayout,
        getDiagonalLayoutRelation(matrixTransposeType, minSlotCount));
    debugAssignLayout(toReplace, newMatrixLayoutAttr);
    assignedLayouts.insert({toReplace, newMatrixLayoutAttr});
    matrix = toReplace;
  }

  // The output has the same per-row layout as the input matrix.
  auto result = vecmatOp->getResult(0);
  RankedTensorType outputType = cast<RankedTensorType>(result.getType());
  IntegerRelation outputLayoutResult =
      getRowMajorLayoutRelation(outputType, minSlotCount);
  LayoutAttr outputLayoutAttr =
      LayoutAttr::getFromIntegerRelation(ctx, outputLayoutResult);
  Attribute kernelInfoAttr =
      cloneKernelInfoWithResultShape(vec, outputType.getShape());

  assignedLayouts.insert({result, outputLayoutAttr});
  setResultLayoutAttr(op, kernelInfoAttr);
  debugAssignLayout(result, outputLayoutAttr);

  auto kernelAttr =
      secret::KernelAttr::get(ctx, KernelName::VecmatDiagonal, /*force=*/false);
  op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr);

  return success();
}

LogicalResult LayoutPropagation::visitOperation(MatvecOp op) {
  auto matvecOp = cast<linalg::ContractionOpInterface>(*op);
  auto matrix = matvecOp.lhs();
  auto matrixType = cast<RankedTensorType>(matrix.getType());

  // TODO(#1597): a layout optimizer should really be selecting the diagonal
  // layout instead of this pass.

  LayoutAttr matrixLayout = getComposedLayoutAttr(matrix);
  // The Halevi-Shoup kernel (all we support at this time) requires one
  // ciphertext per matrix row.
  if (!isRelationSquatDiagonal(matrixType, minSlotCount,
                               matrixLayout.getIntegerRelation())) {
    // Insert a layout conversion op to make the matrix layout squat diagonal
    MLIRContext* ctx = &getContext();
    mlir::IRRewriter builder(ctx);
    auto [toReplace, newMatrixLayoutAttr] =
        convertToLayout(ctx, builder, op, matrix, matrixLayout,
                        getDiagonalLayoutRelation(matrixType, minSlotCount));
    debugAssignLayout(toReplace, newMatrixLayoutAttr);
    assignedLayouts.insert({toReplace, newMatrixLayoutAttr});
    matrix = toReplace;
  }

  // We also require the input to be row-major.
  auto vector = matvecOp.rhs();
  auto vectorType = cast<RankedTensorType>(vector.getType());
  LayoutAttr vectorLayout = getComposedLayoutAttr(vector);
  if (!isRelationRowMajor(vectorType, minSlotCount,
                          vectorLayout.getIntegerRelation())) {
    // Insert a layout conversion op to make the matrix layout squat diagonal
    MLIRContext* ctx = &getContext();
    mlir::IRRewriter builder(ctx);
    auto [toReplace, newVectorLayoutAttr] =
        convertToLayout(ctx, builder, op, vector, vectorLayout,
                        getRowMajorLayoutRelation(vectorType, minSlotCount));
    debugAssignLayout(toReplace, newVectorLayoutAttr);
    assignedLayouts.insert({toReplace, newVectorLayoutAttr});
    vector = toReplace;
  }

  // Always one result, and for the kernels we have right now it's always a
  // row-major replicated vector. Since the matrix may be rectangular, the
  // output layout may have different alignment from the input layout.
  auto result = matvecOp->getResult(0);
  RankedTensorType outputType = cast<RankedTensorType>(result.getType());
  FailureOr<LayoutAttr> outputLayoutResult = defaultLayoutForType(outputType);
  if (failed(outputLayoutResult)) {
    return failure();
  }
  LayoutAttr resultLayout = outputLayoutResult.value();
  Attribute kernelInfoAttr =
      cloneKernelInfoWithResultShape(matrix, outputType.getShape());

  assignedLayouts.insert({result, resultLayout});
  setResultLayoutAttr(op, kernelInfoAttr);
  debugAssignLayout(result, resultLayout);

  // Add secret.kernel attribute for MatvecDiagonal
  MLIRContext* ctx = &getContext();
  auto kernelAttr =
      secret::KernelAttr::get(ctx, KernelName::MatvecDiagonal, /*force=*/false);
  op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr);

  return success();
}

LogicalResult LayoutPropagation::visitOperation(DotOp op) {
  LLVM_DEBUG(llvm::dbgs() << "visitOperation(DotOp): " << op << "\n");
  Value result = op.getResult(0);
  FailureOr<LayoutAttr> layout = defaultLayoutForType(result.getType());
  if (failed(layout)) {
    return failure();
  }
  Attribute kernelInfoAttr = cloneKernelInfoWithResultShape(
      op.getInputs().front(),
      cast<RankedTensorType>(result.getType()).getShape());
  assignedLayouts.insert({result, layout.value()});
  setResultLayoutAttr(op, kernelInfoAttr);
  debugAssignLayout(result, layout.value());

  // Add secret.kernel attribute for Dot
  MLIRContext* ctx = &getContext();
  auto kernelAttr =
      secret::KernelAttr::get(ctx, KernelName::Dot, /*force=*/false);
  op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr);

  return success();
}

LogicalResult LayoutPropagation::alignInitWithResultLayout(
    Operation* op, Value init, Attribute resultLayoutAttr) {
  MLIRContext* ctx = &getContext();
  LayoutAttr targetLayout;
  if (auto arrayAttr = dyn_cast<ArrayAttr>(resultLayoutAttr)) {
    targetLayout = LayoutAttr::composeLayouts(arrayAttr, ctx);
  } else {
    targetLayout = dyn_cast<LayoutAttr>(resultLayoutAttr);
  }
  if (!targetLayout) {
    return op->emitError() << "expected a layout for the op result";
  }

  if (!assignedLayouts.contains(init)) {
    return op->emitError() << "init operand has no assigned layout";
  }
  LayoutAttr initLayout = getComposedLayoutAttr(init);
  if (initLayout == targetLayout ||
      isRelationEqual(initLayout.getIntegerRelation(),
                      targetLayout.getIntegerRelation())) {
    return success();
  }

  LLVM_DEBUG(llvm::dbgs() << "init layout does not match the kernel's result "
                             "layout, inserting layout conversion.\n");
  mlir::IRRewriter builder(ctx);
  auto [toReplace, newInitLayoutAttr] = convertToLayout(
      ctx, builder, op, init, initLayout, targetLayout.getIntegerRelation());
  debugAssignLayout(toReplace, newInitLayoutAttr);
  assignedLayouts.insert({toReplace, newInitLayoutAttr});
  return success();
}

LogicalResult LayoutPropagation::visitOperation(Conv1DOp op) {
  LLVM_DEBUG(llvm::dbgs() << "Specializing visitor on Conv1DOp\n");
  Value data = op.getInputs().front();
  Value filter = op.getInputs().back();
  auto dataType = cast<RankedTensorType>(data.getType());
  auto filterType = cast<RankedTensorType>(filter.getType());

  // Flattened data must fit into the ciphertext size.
  if (dataType.getNumElements() > minSlotCount) {
    return op->emitOpError()
           << "Flattened data must fit into a single ciphertext, but got "
           << dataType.getNumElements() << " elements and ciphertext size is "
           << minSlotCount;
  }

  MLIRContext* ctx = &getContext();
  mlir::IRRewriter builder(ctx);

  // TODO(#1597): a layout optimizer should really be selecting the
  // layout instead of this pass.
  LayoutAttr dataLayout = getComposedLayoutAttr(data);
  if (!isRelationRowMajor(dataType, minSlotCount,
                          dataLayout.getIntegerRelation())) {
    LLVM_DEBUG(llvm::dbgs() << "conv_1d data input is not row major, inserting "
                               "layout conversion.\n");
    auto [toReplace, newDataLayoutAttr] =
        convertToLayout(ctx, builder, op, data, dataLayout,
                        getRowMajorLayoutRelation(dataType, minSlotCount));
    debugAssignLayout(toReplace, newDataLayoutAttr);
    assignedLayouts.insert({toReplace, newDataLayoutAttr});
  }

  // The kernel for this operation requires expanding the conv filter matrix
  // into a larger matrix and then diagonalizing.
  LayoutAttr filterLayout = getComposedLayoutAttr(filter);
  if (!isRelationConvFilterDiagonalized(filterType, dataType, /*padding=*/0,
                                        minSlotCount,
                                        filterLayout.getIntegerRelation())) {
    LLVM_DEBUG(llvm::dbgs() << "conv_1d filter input is not diagonalized, "
                               "inserting layout conversion.\n");
    // Insert a layout conversion op to make the matrix layout expanded and
    // squat diagonal
    auto convRelation = getConvFilterDiagonalizedRelation(
        filterType, dataType, /*padding=*/0, minSlotCount);
    if (failed(convRelation)) {
      return failure();
    }
    auto [toReplace, newFilterLayoutAttr] = convertToLayout(
        ctx, builder, op, filter, filterLayout, convRelation.value());
    debugAssignLayout(toReplace, newFilterLayoutAttr);
    assignedLayouts.insert({toReplace, newFilterLayoutAttr});
  }
  // Always one result, and for the kernels we have right now it's always a
  // row-major replicated vector. Since the
  // output matrix will have different shape than the input, assign the new
  // layout.
  auto result = op->getResult(0);
  RankedTensorType outputType = cast<RankedTensorType>(result.getType());
  FailureOr<LayoutAttr> outputLayoutResult = defaultLayoutForType(outputType);
  if (failed(outputLayoutResult)) {
    return failure();
  }
  LayoutAttr resultLayout = outputLayoutResult.value();
  Attribute kernelInfoAttr =
      cloneKernelInfoWithResultShape(data, outputType.getShape());

  assignedLayouts.insert({result, resultLayout});
  setResultLayoutAttr(op, kernelInfoAttr);
  auto kernelAttr =
      secret::KernelAttr::get(ctx, KernelName::MatvecDiagonal, /*force=*/false);
  op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr);

  return alignInitWithResultLayout(op, op.getOutputs().front(), resultLayout);
}

LogicalResult LayoutPropagation::visitOperation(Conv2DOp op) {
  LLVM_DEBUG(llvm::dbgs() << "Specializing visitor on Conv2DOp\n");
  Value data = op.getInputs().front();
  Value filter = op.getInputs().back();
  auto dataType = cast<RankedTensorType>(data.getType());
  auto filterType = cast<RankedTensorType>(filter.getType());

  // Flattened data must fit into the ciphertext size.
  if (dataType.getNumElements() > minSlotCount) {
    return op->emitOpError()
           << "Flattened data must fit into a single ciphertext, but got "
           << dataType.getNumElements() << " elements and ciphertext size is "
           << minSlotCount;
  }

  MLIRContext* ctx = &getContext();
  mlir::IRRewriter builder(ctx);

  // TODO(#1597): a layout optimizer should really be selecting the
  // layout instead of this pass.
  LayoutAttr dataLayout = getComposedLayoutAttr(data);
  if (!isRelationRowMajor(dataType, minSlotCount,
                          dataLayout.getIntegerRelation())) {
    LLVM_DEBUG(llvm::dbgs() << "conv_2d data input is not row major, inserting "
                               "layout conversion.\n");
    auto [toReplace, newDataLayoutAttr] =
        convertToLayout(ctx, builder, op, data, dataLayout,
                        getRowMajorLayoutRelation(dataType, minSlotCount));
    debugAssignLayout(toReplace, newDataLayoutAttr);
    assignedLayouts.insert({toReplace, newDataLayoutAttr});
  }

  // The kernel for this operation requires expanding the conv filter matrix
  // into a larger matrix and then diagonalizing.
  LayoutAttr filterLayout = getComposedLayoutAttr(filter);
  if (!isRelationConvFilterDiagonalized(filterType, dataType, /*padding=*/0,
                                        minSlotCount,
                                        filterLayout.getIntegerRelation())) {
    LLVM_DEBUG(llvm::dbgs() << "conv_2d filter input is not diagonalized, "
                               "inserting layout conversion.\n");
    // Insert a layout conversion op to make the matrix layout expanded and
    // squat diagonal
    auto convRelation = getConvFilterDiagonalizedRelation(
        filterType, dataType, /*padding=*/0, minSlotCount);
    if (failed(convRelation)) {
      return failure();
    }
    auto [toReplace, newFilterLayoutAttr] = convertToLayout(
        ctx, builder, op, filter, filterLayout, convRelation.value());
    debugAssignLayout(toReplace, newFilterLayoutAttr);
    assignedLayouts.insert({toReplace, newFilterLayoutAttr});
  }
  // Always one result, and for the kernels we have right now it's always a
  // row-major replicated vector. Since the
  // output matrix will have different shape than the input, assign the new
  // layout.
  auto result = op->getResult(0);
  RankedTensorType outputType = cast<RankedTensorType>(result.getType());
  FailureOr<LayoutAttr> outputLayoutResult = defaultLayoutForType(outputType);
  if (failed(outputLayoutResult)) {
    return failure();
  }
  LayoutAttr resultLayout = outputLayoutResult.value();
  Attribute kernelInfoAttr =
      cloneKernelInfoWithResultShape(data, outputType.getShape());

  assignedLayouts.insert({result, resultLayout});
  setResultLayoutAttr(op, kernelInfoAttr);
  auto kernelAttr =
      secret::KernelAttr::get(ctx, KernelName::MatvecDiagonal, /*force=*/false);
  op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr);

  return alignInitWithResultLayout(op, op.getOutputs().front(), resultLayout);
}

LogicalResult LayoutPropagation::visitOperation(Conv1DNcwFcwOp op) {
  LLVM_DEBUG(llvm::dbgs() << "Specializing visitor on Conv1DNcwFcwOp\n");
  Value data = op.getInputs().front();
  Value filter = op.getInputs().back();
  auto dataType = cast<RankedTensorType>(data.getType());
  auto filterType = cast<RankedTensorType>(filter.getType());
  RankedTensorType outputType =
      cast<RankedTensorType>(op.getResult(0).getType());

  bool interchangeRows = true;

  int64_t stride = op.getStrides().getValues<int64_t>().begin()[0];

  MLIRContext* ctx = &getContext();
  mlir::IRRewriter builder(ctx);

  // Ensure data is in gapped row-major layout with current inputGap.
  // We expect 3-D tensor (N, C, W) but only support N=1.
  if (dataType.getRank() != 3 || dataType.getDimSize(0) != 1) {
    return op->emitOpError() << "Expected 3-D data tensor (N=1, C, W)";
  }

  // Since the stride will be used as the gap factor, the layout requires that
  // the number of output channels is divisible by gap.
  // TODO(#2883): handle padding the output channels to be divisible by gap.
  if (outputType.getDimSize(1) % stride != 0) {
    return op->emitOpError()
           << "Expected number of output channels to be divisible by gap";
  }

  LayoutAttr dataLayout = getComposedLayoutAttr(data);

  ConvMatrixOperand matrixOperand{dataType};
  IntegerRelation targetDataRelation =
      getRowMajorLayoutRelation(dataType, minSlotCount);
  if (auto folded =
          tryFoldPadIntoConvPadding(data, dataType, dataLayout, minSlotCount)) {
    matrixOperand = folded->matrixOperand;
    targetDataRelation = folded->targetRelation;
  }

  if (!isRelationEqual(dataLayout.getIntegerRelation(), targetDataRelation)) {
    LLVM_DEBUG(llvm::dbgs() << "conv_1d data input is not row major, "
                               "inserting layout conversion.\n");
    auto [toReplace, newDataLayoutAttr] =
        convertToLayout(ctx, builder, op, data, dataLayout, targetDataRelation);
    debugAssignLayout(toReplace, newDataLayoutAttr);
    assignedLayouts.insert({toReplace, newDataLayoutAttr});
  }

  // The kernel for this operation requires expanding the conv filter matrix
  // into a larger matrix and then diagonalizing.
  LayoutAttr filterLayout = getComposedLayoutAttr(filter);
  auto convRelation = get1dConvCwFcwFilterDiagonalizedRelation(
      filterType, matrixOperand.dataType, stride, matrixOperand.padding,
      minSlotCount, /*interchangeRows=*/interchangeRows);
  if (failed(convRelation)) {
    return failure();
  }
  if (!isRelationEqual(filterLayout.getIntegerRelation(),
                       convRelation.value())) {
    LLVM_DEBUG(llvm::dbgs() << "conv_1d filter input is not diagonalized, "
                               "inserting layout conversion.\n");

    // Insert a layout conversion op to make the matrix layout expanded and
    // squat diagonal. The added domain schedule ensures ISL can efficiently
    // generate a loop nest implementing the layout. However, the choice of
    // which domain indices to include is arbitrary (so long as ISL remains
    // fast).
    auto [toReplace, newFilterLayoutAttr] = convertToLayout(
        ctx, builder, op, filter, filterLayout, convRelation.value(),
        /*domainSchedule=*/{0, 1});
    debugAssignLayout(toReplace, newFilterLayoutAttr);
    assignedLayouts.insert({toReplace, newFilterLayoutAttr});
  }

  // Always one result. For a gapped output, the result will also have a
  // row-shuffled gap. Future users may need to insert a layout conversion.
  auto result = op->getResult(0);
  presburger::IntegerRelation resultRelation =
      get1dConvResultRelation(outputType, stride, /*padding=*/0, minSlotCount,
                              /*interchangeRows=*/interchangeRows);
  LayoutAttr resultLayoutAttr =
      LayoutAttr::getFromIntegerRelation(ctx, resultRelation);
  Attribute kernelInfoAttr =
      cloneKernelInfoWithResultShape(data, outputType.getShape());

  // Record what the filter was diagonalized against, so that
  // ConvertToCiphertextSemantics rebuilds the same expanded matrix shape.
  setConvFoldedPadding(op, matrixOperand.padding);

  assignedLayouts.insert({result, resultLayoutAttr});
  setResultLayoutAttr(op, kernelInfoAttr);
  auto kernelAttr =
      secret::KernelAttr::get(ctx, KernelName::MatvecDiagonal, /*force=*/false);
  op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr);

  return alignInitWithResultLayout(op, op.getOutputs().front(),
                                   resultLayoutAttr);
}

LogicalResult LayoutPropagation::visitOperation(Conv2DNchwFchwOp op) {
  MLIRContext* ctx = &getContext();

  LLVM_DEBUG(llvm::dbgs() << "Specializing visitor on Conv2DNchwFchwOp\n");
  Value data = op.getInputs().front();
  Value filter = op.getInputs().back();
  auto dataType = cast<RankedTensorType>(data.getType());
  auto filterType = cast<RankedTensorType>(filter.getType());
  RankedTensorType outputType =
      cast<RankedTensorType>(op.getResult(0).getType());

  SmallVector<int64_t> strides(op.getStrides().getValues<int64_t>().begin(),
                               op.getStrides().getValues<int64_t>().end());
  if (!llvm::all_equal(strides)) {
    return op->emitOpError() << "Expected equal strides for Conv2DNchwFchwOp";
  }

  // Interchange rows only if the stride is greater than 1. Otherwise, we can
  // densely pack the rows and outputs without channel gapping.
  bool interchangeRows = strides[0] > 1;

  auto dataParentInfo = findAttributeAssociatedWith(data, kKernelInfoAttrName);
  if (failed(dataParentInfo)) {
    return op->emitOpError() << "Failed to find kernel info for data input";
  }
  auto dataKernelInfo = getKernelInfo(dataParentInfo.value());
  if (!dataKernelInfo) {
    return op->emitOpError() << "Failed to get kernel info for data input";
  }
  auto gapFactor = strides[0] * dataKernelInfo->gapFactor;
  RankedTensorType fheInputType = RankedTensorType::get(
      dataKernelInfo->resultShape, outputType.getElementType());
  RankedTensorType fheOutputType = outputType;

  if (interchangeRows) {
    // If interchangeRows is on, then the output shape may include reshaping the
    // striding and gapping.
    int64_t hFhe = std::max(dataKernelInfo->resultShape[2],
                            outputType.getDimSize(2) * gapFactor);
    int64_t wFhe = std::max(dataKernelInfo->resultShape[3],
                            outputType.getDimSize(3) * gapFactor);
    int64_t cFhe = outputType.getDimSize(1) / (gapFactor * gapFactor);
    fheOutputType =
        RankedTensorType::get({outputType.getDimSize(0), cFhe, hFhe, wFhe},
                              outputType.getElementType());
  }
  KernelInfo kernelInfo = {
      .inputShape = llvm::to_vector(fheInputType.getShape()),
      .resultShape = llvm::to_vector(fheOutputType.getShape()),
      .gapFactor = gapFactor};
  Attribute kernelInfoAttr = makeKernelInfoAttr(ctx, kernelInfo);

  // Ensure data is in gapped row-major layout with current inputGap.
  // We expect 4-D tensor (N, C, H, W) but only support N=1.
  if (dataType.getRank() != 4 || dataType.getDimSize(0) != 1) {
    return op->emitOpError() << "Expected 4-D data tensor (N=1, C, H, W)";
  }

  // Since the stride will be used as the gap factor, the layout requires that
  // the number of output channels is divisible by gap^2.
  // TODO(#2883): handle padding the output channels to be divisible by gap^2.
  if (outputType.getDimSize(1) % (strides[0] * strides[0]) != 0) {
    return op->emitOpError()
           << "Expected number of output channels to be divisible by gap^2";
  }

  LayoutAttr dataLayout = getComposedLayoutAttr(data);
  IntegerRelation targetDataRelation =
      getRowMajorLayoutRelation(fheInputType, minSlotCount);

  mlir::IRRewriter builder(ctx);
  if (!isRelationEqual(dataLayout.getIntegerRelation(), targetDataRelation)) {
    LLVM_DEBUG(llvm::dbgs() << "conv_2d data input is not row major, "
                               "inserting layout conversion.\n");
    auto [toReplace, newDataLayoutAttr] =
        convertToLayout(ctx, builder, op, data, dataLayout, targetDataRelation);
    debugAssignLayout(toReplace, newDataLayoutAttr);
    assignedLayouts.insert({toReplace, newDataLayoutAttr});
  }

  // The kernel for this operation requires expanding the conv filter matrix
  // into a larger matrix and then diagonalizing.
  Value originalFilter = filter;
  if (auto assignOp = filter.getDefiningOp<AssignLayoutOp>()) {
    originalFilter = assignOp.getValue();
  }
  auto maybeRels = get2dConvChwFchwFilterAsSequence(
      filterType, fheInputType, strides, /*padding=*/0, minSlotCount,
      interchangeRows);
  if (failed(maybeRels)) {
    return failure();
  }
  auto rels = maybeRels.value();
  SmallVector<Attribute> layoutAttrs;
  for (const auto& rel : rels) {
    layoutAttrs.push_back(LayoutAttr::getFromIntegerRelation(ctx, rel));
  }
  ArrayAttr layoutList = ArrayAttr::get(ctx, layoutAttrs);

  builder.setInsertionPoint(op);
  auto assignLayoutOp =
      AssignLayoutOp::create(builder, op->getLoc(), originalFilter, layoutList);

  Value toReplace = assignLayoutOp.getResult();
  builder.replaceUsesWithIf(filter, toReplace, [&](OpOperand& otherOperand) {
    return otherOperand.getOwner() == op;
  });
  LayoutAttr composedLayout = LayoutAttr::composeLayouts(layoutList, ctx);

  debugAssignLayout(toReplace, composedLayout);
  assignedLayouts.insert({toReplace, composedLayout});
  setAttributeAssociatedWith(
      toReplace, tensor_ext::TensorExtDialect::kLayoutAttrName, composedLayout);

  // Always one result. For a gapped output, the result will also have a
  // pixel-shuffled gap. Future users may need to insert a layout conversion.
  auto result = op->getResult(0);
  Attribute resultLayoutAttr;
  presburger::IntegerRelation rel1 =
      get2dConvResultRelation(outputType, strides, /*padding=*/0, minSlotCount);

  if (interchangeRows) {
    presburger::IntegerRelation rel2 = get2dConvRowInterchangeLayoutRelation(
        outputType, strides, minSlotCount);
    LayoutAttr layout1 = LayoutAttr::getFromIntegerRelation(ctx, rel1);
    LayoutAttr layout2 = LayoutAttr::getFromIntegerRelation(ctx, rel2);
    resultLayoutAttr = ArrayAttr::get(ctx, {layout1, layout2});
  } else {
    resultLayoutAttr = LayoutAttr::getFromIntegerRelation(ctx, rel1);
  }

  assignedLayouts.insert({result, resultLayoutAttr});
  auto resultKernelInfo =
      cloneKernelInfoWithResultShape(data, outputType.getShape());

  setResultLayoutAttr(op, resultKernelInfo);
  auto kernelAttr =
      secret::KernelAttr::get(ctx, KernelName::MatvecDiagonal, /*force=*/false);
  op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr);

  setAttributeAssociatedWith(op.getResult(0), kKernelInfoAttrName,
                             kernelInfoAttr);

  return alignInitWithResultLayout(op, op.getOutputs().front(),
                                   resultLayoutAttr);
}

LogicalResult LayoutPropagation::visitOperation(BatchMatmulOp op) {
  MLIRContext* ctx = &getContext();
  mlir::IRRewriter builder(ctx);

  Value lhs = op.getOperand(0);
  Value rhs = op.getOperand(1);
  auto lhsType = cast<RankedTensorType>(lhs.getType());
  auto rhsType = cast<RankedTensorType>(rhs.getType());
  auto result = op->getResult(0);

  bool inputSecret = isSecret(lhs, solver);
  bool filterSecret = isSecret(rhs, solver);

  LLVM_DEBUG(llvm::dbgs() << "lhs=" << lhs << ";\nrhs=" << rhs << "\n");

  // TODO(#3173): support layout propagation for pt-ct / ct-pt batch matrix
  // multiplication.
  if (!inputSecret) {
    return builder.notifyMatchFailure(
        op, "pt-ct batch matrix multiplication is not supported");
  }
  if (!filterSecret) {
    return builder.notifyMatchFailure(
        op, "ct-pt batch matrix multiplication is not supported");
  }

  int64_t hLhs = lhsType.getDimSize(0);
  int64_t mDim = lhsType.getDimSize(1);
  int64_t nDim = lhsType.getDimSize(2);
  int64_t hRhs = rhsType.getDimSize(0);
  int64_t pDim = rhsType.getDimSize(2);

  bool batchEqual = (hLhs == hRhs);
  bool lhsCoprime = std::gcd(hLhs, mDim) == 1 && std::gcd(hLhs, nDim) == 1 &&
                    std::gcd(mDim, nDim) == 1;
  bool rhsCoprime = std::gcd(hRhs, nDim) == 1 && std::gcd(hRhs, pDim) == 1 &&
                    std::gcd(nDim, pDim) == 1;
  bool outputCoprime = std::gcd(hLhs, mDim) == 1 && std::gcd(hLhs, pDim) == 1 &&
                       std::gcd(mDim, pDim) == 1;

  // Tricyclic ct-ct batch matmul.
  if (inputSecret && filterSecret && batchEqual && lhsCoprime && rhsCoprime &&
      outputCoprime) {
    LayoutAttr lhsLayout = getComposedLayoutAttr(lhs);
    if (!isRelationTricyclic(lhsType, minSlotCount,
                             lhsLayout.getIntegerRelation())) {
      auto [toReplace, newInputMatrixLayoutAttr] =
          convertToLayout(ctx, builder, op, lhs, lhsLayout,
                          getTricyclicLayoutRelation(lhsType, minSlotCount));
      debugAssignLayout(toReplace, newInputMatrixLayoutAttr);
      assignedLayouts.insert({toReplace, newInputMatrixLayoutAttr});
    }

    LayoutAttr rhsLayout = getComposedLayoutAttr(rhs);
    if (!isRelationTricyclic(rhsType, minSlotCount,
                             rhsLayout.getIntegerRelation())) {
      auto [toReplace, newFilterMatrixLayoutAttr] =
          convertToLayout(ctx, builder, op, rhs, rhsLayout,
                          getTricyclicLayoutRelation(rhsType, minSlotCount));
      debugAssignLayout(toReplace, newFilterMatrixLayoutAttr);
      assignedLayouts.insert({toReplace, newFilterMatrixLayoutAttr});
    }

    RankedTensorType outputType = cast<RankedTensorType>(result.getType());
    IntegerRelation outputLayoutResult =
        getTricyclicLayoutRelation(outputType, minSlotCount);
    LayoutAttr outputLayoutAttr =
        LayoutAttr::getFromIntegerRelation(ctx, outputLayoutResult);

    assignedLayouts.insert({result, outputLayoutAttr});
    setResultLayoutAttr(op);
    debugAssignLayout(result, outputLayoutAttr);

    auto kernelAttr = secret::KernelAttr::get(
        ctx, KernelName::BatchMatmulTricyclic, /*force=*/false);
    op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr);
    return success();
  }

  return builder.notifyMatchFailure(
      op, "batch matmul shape requirements for tricyclic layout are not met");
}

LogicalResult LayoutPropagation::visitOperation(MatmulOp op) {
  MLIRContext* ctx = &getContext();
  mlir::IRRewriter builder(ctx);

  Value lhs = op.getOperand(0);
  Value rhs = op.getOperand(1);
  auto lhsType = cast<RankedTensorType>(lhs.getType());
  auto rhsType = cast<RankedTensorType>(rhs.getType());
  auto result = op->getResult(0);

  bool inputSecret = isSecret(lhs, solver);
  bool filterSecret = isSecret(rhs, solver);
  int64_t mDim = lhsType.getDimSize(0);
  int64_t nDim = lhsType.getDimSize(1);
  int64_t pDim = rhsType.getDimSize(1);
  bool lhsCoprime = std::gcd(mDim, nDim) == 1;
  bool rhsCoprime = std::gcd(nDim, pDim) == 1;
  bool outputCoprime = std::gcd(mDim, pDim) == 1;

  LLVM_DEBUG(llvm::dbgs() << "lhs=" << lhs << ";\nrhs=" << rhs << "\n");

  // Bicyclic ct-ct matmul.
  if (inputSecret && filterSecret && lhsCoprime && rhsCoprime &&
      outputCoprime) {
    LayoutAttr lhsLayout = getComposedLayoutAttr(lhs);
    if (!isRelationBicyclic(lhsType, minSlotCount,
                            lhsLayout.getIntegerRelation())) {
      auto [toReplace, newInputMatrixLayoutAttr] =
          convertToLayout(ctx, builder, op, lhs, lhsLayout,
                          getBicyclicLayoutRelation(lhsType, minSlotCount));
      debugAssignLayout(toReplace, newInputMatrixLayoutAttr);
      assignedLayouts.insert({toReplace, newInputMatrixLayoutAttr});
    }

    LayoutAttr rhsLayout = getComposedLayoutAttr(rhs);
    if (!isRelationBicyclic(rhsType, minSlotCount,
                            rhsLayout.getIntegerRelation())) {
      auto [toReplace, newFilterMatrixLayoutAttr] =
          convertToLayout(ctx, builder, op, rhs, rhsLayout,
                          getBicyclicLayoutRelation(rhsType, minSlotCount));
      debugAssignLayout(toReplace, newFilterMatrixLayoutAttr);
      assignedLayouts.insert({toReplace, newFilterMatrixLayoutAttr});
    }

    RankedTensorType outputType = cast<RankedTensorType>(result.getType());
    IntegerRelation outputLayoutResult =
        getBicyclicLayoutRelation(outputType, minSlotCount);
    LayoutAttr outputLayoutAttr =
        LayoutAttr::getFromIntegerRelation(ctx, outputLayoutResult);
    auto kernelInfoAttr =
        cloneKernelInfoWithResultShape(lhs, outputType.getShape());

    assignedLayouts.insert({result, outputLayoutAttr});
    setResultLayoutAttr(op, kernelInfoAttr);
    debugAssignLayout(result, outputLayoutAttr);

    auto kernelAttr = secret::KernelAttr::get(ctx, KernelName::MatmulBicyclic,
                                              /*force=*/false);
    op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr);
    return success();
  }

  // Bicyclic ct-pt matmul with generalized-diagonal-layout plaintext.
  if (outputCoprime && ((inputSecret && !filterSecret && lhsCoprime) ||
                        (!inputSecret && filterSecret && rhsCoprime))) {
    Value secretOperand = inputSecret ? lhs : rhs;
    RankedTensorType secretType = inputSecret ? lhsType : rhsType;
    Value weight = inputSecret ? rhs : lhs;
    RankedTensorType weightType = inputSecret ? rhsType : lhsType;

    LayoutAttr secretLayout = getComposedLayoutAttr(secretOperand);
    if (!isRelationBicyclic(secretType, minSlotCount,
                            secretLayout.getIntegerRelation())) {
      auto [toReplace, newSecretLayoutAttr] =
          convertToLayout(ctx, builder, op, secretOperand, secretLayout,
                          getBicyclicLayoutRelation(secretType, minSlotCount));
      debugAssignLayout(toReplace, newSecretLayoutAttr);
      assignedLayouts.insert({toReplace, newSecretLayoutAttr});
    }

    // Assign a generalized-diagonal layout to the plaintext matrix.
    IntegerRelation diagRelation =
        inputSecret
            ? getBicyclicDiagonalRelation(weightType, /*contractionDim=*/0,
                                          /*stride=*/mDim, minSlotCount)
            : getBicyclicDiagonalRelation(weightType, /*contractionDim=*/1,
                                          /*stride=*/pDim, minSlotCount);
    LayoutAttr weightLayout = getComposedLayoutAttr(weight);
    LayoutAttr diagLayoutAttr =
        LayoutAttr::getFromIntegerRelation(ctx, diagRelation);
    if (weightLayout != diagLayoutAttr) {
      auto [toReplace, newWeightLayoutAttr] =
          convertToLayout(ctx, builder, op, weight, weightLayout, diagRelation);
      debugAssignLayout(toReplace, newWeightLayoutAttr);
      assignedLayouts.insert({toReplace, newWeightLayoutAttr});
    }

    RankedTensorType outputType = cast<RankedTensorType>(result.getType());
    IntegerRelation outputLayoutResult =
        getBicyclicLayoutRelation(outputType, minSlotCount);
    LayoutAttr outputLayoutAttr =
        LayoutAttr::getFromIntegerRelation(ctx, outputLayoutResult);
    auto kernelInfoAttr =
        cloneKernelInfoWithResultShape(secretOperand, outputType.getShape());

    assignedLayouts.insert({result, outputLayoutAttr});
    setResultLayoutAttr(op, kernelInfoAttr);
    debugAssignLayout(result, outputLayoutAttr);

    auto kernelAttr = secret::KernelAttr::get(
        ctx, KernelName::MatmulBicyclicDiagonal, /*force=*/false);
    op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr);
    return success();
  }

  // Assign a per-row layout to the input matrix. Each row of the input matrix
  // will be packed into a separate ciphertext.
  auto lhsLayout = getComposedLayoutAttr(lhs);
  if (!isRelationPerRow(lhsType, minSlotCount,
                        lhsLayout.getIntegerRelation())) {
    // Insert a layout conversion op to make the matrix layout per-row
    auto [toReplace, newInputMatrixLayoutAttr] =
        convertToLayout(ctx, builder, op, lhs, lhsLayout,
                        getPerRowLayoutRelation(lhsType, minSlotCount));
    debugAssignLayout(toReplace, newInputMatrixLayoutAttr);
    assignedLayouts.insert({toReplace, newInputMatrixLayoutAttr});
  }

  // Assign or change the filter matrix a diagonal layout.
  auto rhsTransposeType = RankedTensorType::get(
      {rhsType.getDimSize(1), rhsType.getDimSize(0)}, rhsType.getElementType());
  auto rhsLayout = getComposedLayoutAttr(rhs);
  IntegerRelation rhsLayoutRelation = rhsLayout.getIntegerRelation();
  auto clonedFilterMatrixRelation = rhsLayoutRelation.clone();
  auto domainOffset =
      rhsLayoutRelation.getVarKindOffset(presburger::VarKind::Domain);
  clonedFilterMatrixRelation->swapVar(domainOffset, domainOffset + 1);
  if (!isRelationSquatDiagonal(rhsTransposeType, minSlotCount,
                               *clonedFilterMatrixRelation)) {
    // Insert a layout conversion op to make the matrix layout squat diagonal
    auto [toReplace, newFilterMatrixLayoutAttr] = convertToLayout(
        ctx, builder, op, rhs, rhsLayout,
        getDiagonalLayoutRelation(rhsTransposeType, minSlotCount));
    debugAssignLayout(toReplace, newFilterMatrixLayoutAttr);
    assignedLayouts.insert({toReplace, newFilterMatrixLayoutAttr});
  }

  // The output has the same per-row layout as the input matrix.
  RankedTensorType outputType = cast<RankedTensorType>(result.getType());
  IntegerRelation outputLayoutResult =
      getPerRowLayoutRelation(outputType, minSlotCount);
  LayoutAttr outputLayoutAttr =
      LayoutAttr::getFromIntegerRelation(ctx, outputLayoutResult);
  auto kernelInfoAttr =
      cloneKernelInfoWithResultShape(lhs, outputType.getShape());

  assignedLayouts.insert({result, outputLayoutAttr});
  setResultLayoutAttr(op, kernelInfoAttr);
  debugAssignLayout(result, outputLayoutAttr);

  // Add secret.kernel attribute for MatmulDiagonal
  auto kernelAttr =
      secret::KernelAttr::get(ctx, KernelName::MatmulDiagonal, /*force=*/false);
  op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr);

  return success();
}

LogicalResult LayoutPropagation::visitOperation(ReduceOp op) {
  MLIRContext* ctx = &getContext();
  mlir::IRRewriter builder(ctx);
  builder.setInsertionPoint(op);

  for (const auto& [tensor, result] :
       llvm::zip(op.getInputs(), op.getResults())) {
    LayoutAttr thisLayout = getComposedLayoutAttr(tensor);

    // enforce row-major layout
    RankedTensorType thisType = cast<RankedTensorType>(tensor.getType());
    if (!isRelationRowMajor(thisType, minSlotCount,
                            thisLayout.getIntegerRelation())) {
      LLVM_DEBUG(llvm::dbgs() << "ReduceOp tensor is not row major");
      auto [toReplace, newLayoutAttr] =
          convertToLayout(ctx, builder, op, tensor, thisLayout,
                          getRowMajorLayoutRelation(thisType, minSlotCount));
      debugAssignLayout(toReplace, newLayoutAttr);
      assignedLayouts.insert({toReplace, newLayoutAttr});
      thisLayout = newLayoutAttr;
    }

    // drop dimension on output
    LayoutAttr resultLayout =
        convertLayoutForReduce(thisLayout, op.getDimensions());
    assignedLayouts.insert({result, resultLayout});
    debugAssignLayout(result, resultLayout);
  }
  setResultLayoutAttr(op);
  return success();
}

LogicalResult LayoutPropagation::visitOperation(affine::AffineForOp op) {
  // Transfer the layout of the inits to the region iter args
  for (const auto& [init, iterArg, result] :
       llvm::zip(op.getInits(), op.getRegionIterArgs(), op.getResults())) {
    LayoutAttr layout;

    // The init may not have a layout if it is a non-secret initial value. In
    // that case we assign it a default layout.
    if (!assignedLayouts.contains(init)) {
      mlir::IRRewriter builder(&getContext());
      builder.setInsertionPoint(op);
      auto res = assignDefaultLayoutForOpOperand(op, init, builder);
      if (failed(res)) {
        return op->emitError()
               << "Failed to assign default layout to init " << init;
      }
      layout = cast<LayoutAttr>(res.value().getLayout());
    } else {
      layout = getComposedLayoutAttr(init);
    }
    assignedLayouts.insert({iterArg, layout});
    setAttributeAssociatedWith(
        iterArg, tensor_ext::TensorExtDialect::kLayoutAttrName, layout);
    debugAssignLayout(iterArg, layout);

    // The result of an AffineForOp also has the same type as its iter args.
    assignedLayouts.insert({result, layout});
    debugAssignLayout(result, layout);
  }

  setResultLayoutAttr(op);
  return success();
}

LogicalResult LayoutPropagation::visitOperation(tensor::ExtractOp op) {
  std::string relationStr = "{ [] -> [ct, slot] : ct = 0 and slot = 0 }";
  LayoutAttr scalarLayout = LayoutAttr::get(op.getContext(), relationStr);
  LLVM_DEBUG(llvm::dbgs() << "Assigning scalar layout " << scalarLayout
                          << " to value " << op.getResult() << "\n");
  Value result = op.getResult();
  auto kernelInfoAttr = cloneKernelInfoWithResultShape(op.getTensor(), {});
  assignedLayouts.insert({result, scalarLayout});
  debugAssignLayout(result, scalarLayout);
  setResultLayoutAttr(op, kernelInfoAttr);
  return success();
}

LogicalResult LayoutPropagation::visitOperation(tensor::InsertOp op) {
  // The layout of the result is the same as the layout of the input tensor.
  // The scalar input is handled by the generic logic in visitOperation,
  // which will assign a default scalar layout if it is secret and has no
  // layout.
  LayoutAttr tensorLayout = getComposedLayoutAttr(op.getDest());
  auto kernelInfoAttr = cloneKernelInfoWithResultShape(
      op.getScalar(), op.getDest().getType().getShape());
  Value result = op.getResult();
  assignedLayouts.insert({result, tensorLayout});
  debugAssignLayout(result, tensorLayout);
  setResultLayoutAttr(op, kernelInfoAttr);
  return success();
}

LogicalResult LayoutPropagation::visitOperation(tensor::InsertSliceOp op) {
  // Set the result layout to be the same as the dest tensor.
  if (!assignedLayouts.contains(op.getDest())) {
    return op->emitError() << "Destination tensor has no assigned layout";
  }
  LayoutAttr destLayout = getComposedLayoutAttr(op.getDest());
  auto kernelInfoAttr = cloneKernelInfoWithResultShape(
      op.getSource(), op.getDest().getType().getShape());

  Value result = op.getResult();
  assignedLayouts.insert({result, destLayout});
  debugAssignLayout(result, destLayout);
  setResultLayoutAttr(op, kernelInfoAttr);
  return success();
}

LogicalResult LayoutPropagation::visitOperation(tensor::PadOp op) {
  if (!assignedLayouts.contains(op.getSource())) {
    return op->emitError() << "Source tensor has no assigned layout";
  }

  // Extract static low and high paddings. If dynamic bounds are present, emit
  // error.
  bool hasDynamicBounds = false;
  for (int64_t val : op.getStaticLow()) {
    if (ShapedType::isDynamic(val)) {
      hasDynamicBounds = true;
    }
  }
  for (int64_t val : op.getStaticHigh()) {
    if (ShapedType::isDynamic(val)) {
      hasDynamicBounds = true;
    }
  }
  if (hasDynamicBounds || !op.getLow().empty() || !op.getHigh().empty()) {
    return op->emitError()
           << "Only static padding is supported for layout propagation";
  }

  Value padValue = op.getConstantPaddingValue();
  bool isZeroPad = padValue && (matchPattern(padValue, m_AnyZeroFloat()) ||
                                matchPattern(padValue, m_Zero()));
  if (!isZeroPad) {
    return op->emitError()
           << "layout propagation only supports zero-padding tensor.pad";
  }

  // Check if this pad is eligible to be folded forward into a conv op.
  // If so, we use the shifted relation (special case) to ensure the fold
  // pattern matches.
  bool isEligibleForConvFusion = false;
  RankedTensorType paddedType = op.getResultType();
  if (paddedType.getRank() == 3 && paddedType.getDimSize(0) == 1) {
    ArrayRef<int64_t> low = op.getStaticLow();
    ArrayRef<int64_t> high = op.getStaticHigh();
    if (low.size() == 3 && high.size() == 3) {
      if (low[0] == 0 && high[0] == 0 && low[1] == 0 && high[1] == 0 &&
          low[2] == high[2] && low[2] > 0) {
        isEligibleForConvFusion = true;
      }
    }
  }

  if (isEligibleForConvFusion) {
    // A zero-pad does not move any data: result[i + low] = source[i], so the
    // result layout is the source layout with each domain index shifted by the
    // low padding. Pad positions stay unmapped in the relation; unmapped points
    // are zero-filled when a layout is materialized, which matches the
    // zero-fill pad body.
    IntegerRelation padRelation =
        getComposedLayoutAttr(op.getSource()).getIntegerRelation();
    auto domainVarOffset =
        padRelation.getVarKindOffset(presburger::VarKind::Domain);
    for (auto [dim, low] : llvm::enumerate(op.getStaticLow())) {
      if (low != 0) {
        padRelation = shiftVar(padRelation, domainVarOffset + dim, low);
      }
    }

    LayoutAttr outputLayout =
        LayoutAttr::getFromIntegerRelation(op.getContext(), padRelation);
    Attribute kernelInfoAttr = cloneKernelInfoWithResultShape(
        op.getSource(), op.getResultType().getShape());
    assignedLayouts.insert({op.getResult(), outputLayout});
    debugAssignLayout(op.getResult(), outputLayout);
    setResultLayoutAttr(op, kernelInfoAttr);
    return success();
  }

  // General case: use getPaddingRelation and compose.
  SmallVector<int64_t> lowPadding = llvm::to_vector(op.getStaticLow());
  IntegerRelation sourceLayout =
      getComposedLayoutAttr(op.getSource()).getIntegerRelation();

  RankedTensorType unpaddedType = op.getSourceType();
  IntegerRelation paddingRel =
      getPaddingRelation(paddedType, unpaddedType, lowPadding);

  paddingRel.compose(sourceLayout);

  Attribute kernelInfoAttr =
      cloneKernelInfoWithResultShape(op.getSource(), paddedType.getShape());

  LayoutAttr outputLayout =
      LayoutAttr::getFromIntegerRelation(op.getContext(), paddingRel);
  assignedLayouts.insert({op.getResult(), outputLayout});
  debugAssignLayout(op.getResult(), outputLayout);
  setResultLayoutAttr(op, kernelInfoAttr);
  return success();
}

LogicalResult LayoutPropagation::visitOperation(tensor::ExtractSliceOp op) {
  // Assign the induced layout from extracting a slice from the source tensor.
  if (!assignedLayouts.contains(op.getSource())) {
    return op->emitError() << "Source tensor has no assigned layout";
  }
  IntegerRelation sourceLayout =
      getComposedLayoutAttr(op.getSource()).getIntegerRelation();

  FailureOr<IntegerRelation> maybeSliceExtractionLayout =
      getSliceExtractionRelation(op.getSourceType(), op.getResultType(),
                                 SmallVector<int64_t>(op.getStaticOffsets()),
                                 SmallVector<int64_t>(op.getStaticSizes()),
                                 SmallVector<int64_t>(op.getStaticStrides()));
  if (failed(maybeSliceExtractionLayout)) {
    return failure();
  }
  IntegerRelation sliceExtractionLayout = maybeSliceExtractionLayout.value();

  // Compose the inverted slice extraction layout with the source layout to
  // get the result slice layout.
  sliceExtractionLayout.inverse();
  sliceExtractionLayout.compose(sourceLayout);
  // If the slice extracted was not at offset zero, then the resulting slice may
  // be indexed at a non-zero ciphertext. For example, imagine extracting a
  // slice out of the second ciphertext. Then computing the inverse of the slice
  // extraction layout and composing that with the source relation would mean
  // that the slice would map to the second ciphertext. But a slice extracted
  // from a tensor.extract_slice op is always indexed starting from zero.
  // Reindexing the the resulting relation to start from ciphertext zero.
  auto ctVarOffset =
      sliceExtractionLayout.getVarKindOffset(presburger::VarKind::Range);
  auto ctLowerBound = sliceExtractionLayout.getConstantBound64(
      presburger::BoundType::LB, ctVarOffset);
  if (!ctLowerBound) {
    return op.emitError() << "failed to get constant bound on ciphertext index";
  }
  auto zeroIndexedSliceLayout =
      shiftVar(sliceExtractionLayout, ctVarOffset, -ctLowerBound.value());

  Attribute kernelInfoAttr = cloneKernelInfoWithResultShape(
      op.getSource(), op.getResultType().getShape());
  LayoutAttr outputLayout = LayoutAttr::getFromIntegerRelation(
      op.getContext(), zeroIndexedSliceLayout);
  assignedLayouts.insert({op.getResult(), outputLayout});
  debugAssignLayout(op.getResult(), outputLayout);
  setResultLayoutAttr(op, kernelInfoAttr);
  return success();
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
    Operation* op) {
  LLVM_DEBUG(llvm::dbgs() << "Checking for compatible argument layouts for: "
                          << op->getName() << "\n");
  return TypeSwitch<Operation*, CompatibilityResult>(op)
      // Trivially true ops
      .Case<func::FuncOp, GenericOp, YieldOp, affine::AffineForOp,
            affine::AffineYieldOp>(
          [&](auto op) { return CompatibilityResult{true, std::nullopt}; })
      // Ops with special rules
      .Case<DotOp, ReduceOp, MatvecOp, VecmatOp, MatmulOp, Conv1DOp,
            Conv1DNcwFcwOp, Conv2DOp, Conv2DNchwFchwOp, tensor::InsertSliceOp>(
          [&](auto op) { return hasCompatibleArgumentLayouts(op); })
      // By default, assume operands must all have the same layout.
      .Default([&](Operation* op) {
        std::optional<LayoutAttr> firstFoundLayout;
        // All operands require a layout unless the op implements
        // OperandLayoutRequirementOpInterface, which allows it to specify which
        // subset of operands require layouts.
        auto layoutReqIface = dyn_cast<OperandLayoutRequirementOpInterface>(op);

        for (auto& opOperand : op->getOpOperands()) {
          bool secretness = isSecret(opOperand.get(), solver);
          if (layoutReqIface && !layoutReqIface.operandRequiresLayout(
                                    opOperand.getOperandNumber(), secretness)) {
            continue;
          }

          if (!assignedLayouts.contains(opOperand.get())) {
            // If the operand has no layout, we can't propagate layout
            // information to the result.
            return CompatibilityResult{
                false, op->emitError("operand has no assigned layout")};
          }
          LayoutAttr layout = getComposedLayoutAttr(opOperand.get());

          if (!firstFoundLayout.has_value()) firstFoundLayout = layout;
          if (layout != firstFoundLayout.value()) {
            return CompatibilityResult{false, std::nullopt};
          }
        }

        return CompatibilityResult{true, std::nullopt};
      });
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(DotOp op) {
  LayoutAttr lhsLayout = getComposedLayoutAttr(op.getOperand(0));
  LayoutAttr rhsLayout = getComposedLayoutAttr(op.getOperand(1));

  if (lhsLayout != rhsLayout) {
    return {false, std::nullopt};
  }

  // Outs operand (operand 2) should be 0D.
  LayoutAttr outsLayout = getComposedLayoutAttr(op.getOperand(2));
  if (outsLayout.getIntegerRelation().getNumDomainVars() != 0) {
    return {false, std::nullopt};
  }

  return {true, std::nullopt};
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
    ReduceOp op) {
  // The arguments of a ReduceOp are the tensor(s) to reduce and the
  // initializer values for the reduction.
  for (const auto& [input, init] : llvm::zip(op.getInputs(), op.getInits())) {
    if (!assignedLayouts.contains(input)) {
      return {false, op->emitError("input tensor has no assigned layout")};
    }
    if (!assignedLayouts.contains(init)) {
      return {false,
              op->emitError("initializer tensor has no assigned layout")};
    }

    LayoutAttr inputLayout = getComposedLayoutAttr(input);
    LayoutAttr initLayout = getComposedLayoutAttr(init);
    LayoutAttr reducedInputLayout =
        convertLayoutForReduce(inputLayout, op.getDimensions());

    if (reducedInputLayout != initLayout) {
      return {false, std::nullopt};
    }
  }

  return {true, std::nullopt};
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
    VecmatOp op) {
  // Currently only support secret vectors and plaintext matrices.
  linalg::ContractionOpInterface vecmatOp =
      cast<linalg::ContractionOpInterface>(op.getOperation());
  Value vec = vecmatOp.lhs();
  Value mat = vecmatOp.rhs();
  if (isSecret(mat, solver) || !isSecret(vec, solver)) {
    return {false,
            op->emitError("Only secret vectors and plaintext matrices are "
                          "supported for linalg.vecmat")};
  }

  if (!assignedLayouts.contains(vec)) {
    return {false, op->emitError("vector operand has no assigned layout")};
  }
  return {true, std::nullopt};
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
    MatvecOp op) {
  // Currently only support secret vectors and plaintext matrices.
  linalg::ContractionOpInterface matvecOp =
      cast<linalg::ContractionOpInterface>(op.getOperation());
  Value vec = matvecOp.rhs();
  Value mat = matvecOp.lhs();
  if (isSecret(mat, solver) || !isSecret(vec, solver)) {
    return {false,
            op->emitError("Only secret vectors and plaintext matrices are "
                          "supported for linalg.vecmat")};
  }

  if (!assignedLayouts.contains(vec)) {
    return {false, op->emitError("vector operand has no assigned layout")};
  }
  return {true, std::nullopt};
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
    MatmulOp op) {
  Value lhs = op.getOperand(0);
  Value rhs = op.getOperand(1);

  if (!assignedLayouts.contains(lhs)) {
    return {false, op->emitError("LHS operand has no assigned layout")};
  }
  if (!assignedLayouts.contains(rhs)) {
    return {false, op->emitError("RHS operand has no assigned layout")};
  }

  return {true, std::nullopt};
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
    Conv1DOp op) {
  // Currently only support secret data and plaintext filters.
  Value data = op.getInputs().front();
  Value filter = op.getInputs().back();
  if (isSecret(filter, solver) || !isSecret(data, solver)) {
    return {false, op->emitError("Only secret data and plaintext filters are "
                                 "supported for linalg.conv1d")};
  }

  if (!assignedLayouts.contains(data)) {
    return {false, op->emitError("data operand has no assigned layout")};
  }
  return {true, std::nullopt};
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
    Conv2DOp op) {
  // Currently only support secret data and plaintext filters.
  Value data = op.getInputs().front();
  Value filter = op.getInputs().back();
  if (isSecret(filter, solver) || !isSecret(data, solver)) {
    return {false, op->emitError("Only secret data and plaintext filters are "
                                 "supported for linalg.conv2d")};
  }

  if (!assignedLayouts.contains(data)) {
    return {false, op->emitError("data operand has no assigned layout")};
  }
  return {true, std::nullopt};
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
    Conv1DNcwFcwOp op) {
  // Currently only support secret data and plaintext filters.
  Value data = op.getInputs().front();
  Value filter = op.getInputs().back();
  if (isSecret(filter, solver) || !isSecret(data, solver)) {
    return {false, op->emitError("Only secret data and plaintext filters are "
                                 "supported for linalg.conv_1d_ncw_fcw")};
  }

  if (!assignedLayouts.contains(data)) {
    return {false, op->emitError("data operand has no assigned layout")};
  }
  return {true, std::nullopt};
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
    Conv2DNchwFchwOp op) {
  // Currently only support secret data and plaintext filters.
  Value data = op.getInputs().front();
  Value filter = op.getInputs().back();
  if (isSecret(filter, solver) || !isSecret(data, solver)) {
    return {false, op->emitError("Only secret data and plaintext filters are "
                                 "supported for linalg.conv_2d_nchw_fchw")};
  }

  if (!assignedLayouts.contains(data)) {
    return {false, op->emitError("data operand has no assigned layout")};
  }
  return {true, std::nullopt};
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
    tensor::InsertSliceOp op) {
  // The arguments of a tensor::InsertSliceOp are the tensors to insert and the
  // tensor to insert into.
  auto insert = op.getSource();
  auto dest = op.getDest();

  if (!assignedLayouts.contains(insert)) {
    return {false, op->emitError("input tensor has no assigned layout")};
  }
  if (!assignedLayouts.contains(dest)) {
    return {false, op->emitError("destination tensor has no assigned layout")};
  }

  LayoutAttr insertLayout = getComposedLayoutAttr(insert);
  LayoutAttr destLayout = getComposedLayoutAttr(dest);
  auto destRelation = destLayout.getIntegerRelation();
  auto compatibleDestRelation = pushSliceLayoutThroughInsertSlice(
      SmallVector<int64_t>(op.getStaticSizes()), op.getResultType().getShape(),
      insertLayout.getIntegerRelation());
  assert(succeeded(compatibleDestRelation) && "expected to infer slice layout");

  // We compare layout attributes, which is a strict equality check on the
  // attribute string. This is more strict than comparing the IntegerRelations,
  // but equality checking on the relations is expensive.
  if (LayoutAttr::getFromIntegerRelation(
          op.getContext(), compatibleDestRelation.value()) != destLayout) {
    return {false, std::nullopt};
  }

  return {true, std::nullopt};
}

void LayoutPropagation::rectifyIncompatibleOperandLayouts(Operation* op) {
  LLVM_DEBUG({
    auto diag = op->emitRemark() << "Inserting layout conversion op due to "
                                    "disagreeing operand layouts";
    auto& note = diag.attachNote();
    for (auto operand : op->getOperands()) {
      LayoutAttr operandLayout;
      if (assignedLayouts.contains(operand))
        operandLayout = getComposedLayoutAttr(operand);
      note << "\n- Operand: " << operand << "; Layout: " << operandLayout;
    }
  });

  TypeSwitch<Operation*>(op)
      // These ops shouldn't rectify operand layouts
      .Case<func::FuncOp, func::ReturnOp, secret::GenericOp, secret::YieldOp>(
          [&](auto op) { return; })
      // Ops with special rules
      .Case<DotOp, ReduceOp, tensor::InsertOp, tensor::InsertSliceOp>(
          [&](auto op) { return rectifyIncompatibleOperandLayouts(op); })
      .Default([&](Operation* op) {
        // Default target layout is chosen arbitrarily as the first operand's
        // layout for now. A different pass is responsible for optimizing the
        // placement and mechanics of the layout conversion ops.
        mlir::IRRewriter builder(&getContext());
        const auto it = llvm::find_if(op->getOperands(), [this](Value pair) {
          return assignedLayouts.contains(pair);
        });
        LayoutAttr targetLayout = getComposedLayoutAttr(*it);

        for (auto& opOperand : op->getOpOperands()) {
          if (!assignedLayouts.contains(opOperand.get())) continue;
          LayoutAttr sourceLayout = getComposedLayoutAttr(opOperand.get());

          if (sourceLayout != targetLayout) {
            builder.setInsertionPoint(op);
            ConvertLayoutOp convertOp =
                ConvertLayoutOp::create(builder, op->getLoc(), opOperand.get(),
                                        sourceLayout, targetLayout);

            auto* lattice = solver->getOrCreateState<SecretnessLattice>(
                convertOp.getResult());
            lattice->getValue().setSecretness(
                isSecret(opOperand.get(), solver));

            // Layout of the result is the same as the target layout of the
            // conversion. Mostly this is done for consistency: all ops have
            // an attribute describing the layout of their results.
            OpBuilder builder(&getContext());
            assignedLayouts.insert({convertOp.getResult(), targetLayout});
            setResultLayoutAttr(convertOp,
                                findAttributeAssociatedWith(opOperand.get(),
                                                            kKernelInfoAttrName)
                                    .value_or(Attribute()));
            op->setOperand(opOperand.getOperandNumber(), convertOp.getResult());
          }
        }
      });
}

void LayoutPropagation::rectifyIncompatibleOperandLayouts(DotOp op) {
  mlir::IRRewriter builder(&getContext());
  builder.setInsertionPoint(op);

  LayoutAttr lhsLayout = getComposedLayoutAttr(op.getOperand(0));
  LayoutAttr rhsLayout = getComposedLayoutAttr(op.getOperand(1));

  if (lhsLayout != rhsLayout) {
    // Convert RHS to match LHS layout
    ConvertLayoutOp convertOp = ConvertLayoutOp::create(
        builder, op->getLoc(), op.getOperand(1), rhsLayout, lhsLayout);
    auto* lattice =
        solver->getOrCreateState<SecretnessLattice>(convertOp.getResult());
    lattice->getValue().setSecretness(isSecret(op.getOperand(1), solver));

    builder.replaceUsesWithIf(
        op.getOperand(1), convertOp.getResult(),
        [&](OpOperand& operand) { return operand.getOwner() == op; });
    assignedLayouts.insert({convertOp.getResult(), lhsLayout});
  }

  // Ensure Outs is 0D.
  LayoutAttr outsLayout = getComposedLayoutAttr(op.getOperand(2));
  if (outsLayout.getIntegerRelation().getNumDomainVars() != 0) {
    FailureOr<LayoutAttr> default0D =
        defaultLayoutForType(op.getOperand(2).getType());
    if (succeeded(default0D)) {
      ConvertLayoutOp convertOp =
          ConvertLayoutOp::create(builder, op->getLoc(), op.getOperand(2),
                                  outsLayout, default0D.value());
      auto* lattice =
          solver->getOrCreateState<SecretnessLattice>(convertOp.getResult());
      lattice->getValue().setSecretness(isSecret(op.getOperand(2), solver));

      builder.replaceUsesWithIf(
          op.getOperand(2), convertOp.getResult(),
          [&](OpOperand& operand) { return operand.getOwner() == op; });
      assignedLayouts.insert({convertOp.getResult(), default0D.value()});
    }
  }
}

void LayoutPropagation::rectifyIncompatibleOperandLayouts(ReduceOp op) {
  mlir::IRRewriter builder(&getContext());
  builder.setInsertionPoint(op);

  for (const auto& [input, init] : llvm::zip(op.getInputs(), op.getInits())) {
    LayoutAttr inputLayout = getComposedLayoutAttr(input);
    LayoutAttr initLayout = getComposedLayoutAttr(init);
    LayoutAttr reducedInputLayout =
        convertLayoutForReduce(inputLayout, op.getDimensions());

    if (reducedInputLayout != initLayout) {
      ConvertLayoutOp convertOp = ConvertLayoutOp::create(
          builder, op->getLoc(), init, initLayout, reducedInputLayout);
      Value toReplace = convertOp.getResult();
      builder.replaceUsesWithIf(init, toReplace, [&](OpOperand& operand) {
        return operand.getOwner() == op;
      });
      assignedLayouts.insert({toReplace, reducedInputLayout});
      setResultLayoutAttr(convertOp);
    }
  }
}

void LayoutPropagation::rectifyIncompatibleOperandLayouts(
    tensor::InsertSliceOp op) {
  // Update the dest tensor to align with the source tensor slice.
  LayoutAttr sliceLayout = getComposedLayoutAttr(op.getSource());
  auto maybeNewLayout = pushSliceLayoutThroughInsertSlice(
      SmallVector<int64_t>(op.getStaticSizes()), op.getResultType().getShape(),
      sliceLayout.getIntegerRelation());
  assert(succeeded(maybeNewLayout) && "expected to infer slice layout");
  auto newLayoutAttr = LayoutAttr::getFromIntegerRelation(
      op.getContext(), maybeNewLayout.value());

  if (!assignedLayouts.contains(op.getDest())) {
    assignedLayouts.insert({op.getDest(), newLayoutAttr});
    return;
  }

  mlir::IRRewriter builder(&getContext());
  builder.setInsertionPoint(op);
  LayoutAttr destLayout = getComposedLayoutAttr(op.getDest());
  auto convertLayoutOp = ConvertLayoutOp::create(
      builder, op->getLoc(), op.getDest(), destLayout, newLayoutAttr);
  Value toReplace = convertLayoutOp.getResult();
  builder.replaceUsesWithIf(op.getDest(), toReplace, [&](OpOperand& operand) {
    return operand.getOwner() == op;
  });
  assignedLayouts.insert({toReplace, newLayoutAttr});
  setResultLayoutAttr(convertLayoutOp);
}

void LayoutPropagation::rectifyIncompatibleOperandLayouts(tensor::InsertOp op) {
  // The scalar input's layout must be made to match the tensor input's layout.
  // 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 with ciphertext inputs,
  // so we instead incur the cost of a layout conversion before the insert.
  mlir::IRRewriter builder(&getContext());
  builder.setInsertionPoint(op);

  LayoutAttr destLayout = getComposedLayoutAttr(op.getDest());
  LayoutAttr scalarLayout = getComposedLayoutAttr(op.getScalar());
  IntegerRelation destRel = destLayout.getIntegerRelation();
  std::optional<int64_t> destNumCts = destRel.getConstantBound64(
      presburger::BoundType::UB,
      // First range var is the ciphertext index, always has lower bound of
      // zero
      destRel.getVarKindOffset(presburger::VarKind::Range));
  assert(destNumCts.has_value());

  std::string newScalarLayoutStr = llvm::formatv(
      "{ [] -> [ct, slot] : 0 <= ct <= {0} and 0 <= slot <= {1} }",
      destNumCts.value(), minSlotCount - 1);
  LayoutAttr newScalarLayout =
      LayoutAttr::get(op.getContext(), newScalarLayoutStr);

  ConvertLayoutOp convertOp = ConvertLayoutOp::create(
      builder, op->getLoc(), op.getScalar(), scalarLayout, newScalarLayout);
  Value toReplace = convertOp.getResult();
  builder.replaceUsesWithIf(op.getScalar(), toReplace, [&](OpOperand& operand) {
    return operand.getOwner() == op;
  });
  assignedLayouts.insert({toReplace, newScalarLayout});
  setResultLayoutAttr(convertOp);
}

void LayoutPropagation::passLayoutThroughOp(Operation* op) {
  // All inputs have the same layout, so just propagate it to all results
  LayoutAttr layout = getComposedLayoutAttr(op->getOperand(0));
  for (Value result : op->getResults()) {
    assignedLayouts.insert({result, layout});
    debugAssignLayout(result, layout);
  }
  setResultLayoutAttr(op);
}

FailureOr<LayoutAttr> LayoutPropagation::defaultLayoutForScalarType(
    Type scalarType) {
  auto* ctx = scalarType.getContext();
  OpBuilder builder(ctx);

  // Only support Int/Index/Float scalars,
  // fail for any other random type that might make it here.
  if (!isa<FloatType, IntegerType, IndexType>(scalarType)) {
    return failure();
  }

  std::string relationStr = llvm::formatv(
      "{ [] -> [ct, slot] : ct = 0 and 0 <= slot <= {0} }", minSlotCount - 1);
  return LayoutAttr::get(ctx, relationStr);
}

FailureOr<LayoutAttr> LayoutPropagation::defaultLayoutForType(Type type) {
  Type ty = type;
  if (SecretType secretType = dyn_cast<SecretType>(type)) {
    ty = secretType.getValueType();
  }

  RankedTensorType tensorType = dyn_cast<RankedTensorType>(ty);
  if (!tensorType) {
    LLVM_DEBUG(llvm::dbgs() << "defaultLayoutForType: not a tensor, calling "
                               "defaultLayoutForScalarType\n");
    return defaultLayoutForScalarType(ty);
  }

  LLVM_DEBUG(llvm::dbgs() << "defaultLayoutForType: tensor rank="
                          << tensorType.getRank() << "\n");
  if (tensorType.getRank() == 0) {
    LLVM_DEBUG(llvm::dbgs() << "defaultLayoutForType: rank 0 tensor, calling "
                               "defaultLayoutForScalarType\n");
    return defaultLayoutForScalarType(tensorType.getElementType());
  }

  // By default, each tensor is laid out in row-major order. The slots will be
  // repeated each N elements where N is the next power of two of the total
  // number of data elements.
  LLVM_DEBUG(llvm::dbgs() << "getting row-major layout map for type="
                          << tensorType << "\n");
  IntegerRelation relation =
      getRowMajorLayoutRelation(tensorType, minSlotCount);
  return LayoutAttr::getFromIntegerRelation(tensorType.getContext(), relation);
}

void LayoutPropagation::setResultLayoutAttr(Operation* op,
                                            Attribute kernelInfo) {
  // If no kernelInfo attribute is present, check if any of the operands have
  // one.
  if (!kernelInfo && op->getNumOperands() > 0) {
    for (Value operand : op->getOperands()) {
      auto info = findAttributeAssociatedWith(operand, kKernelInfoAttrName);
      if (succeeded(info)) {
        kernelInfo = info.value();
        break;
      }
    }
  }
  if (kernelInfo) {
    for (Value result : op->getResults()) {
      setAttributeAssociatedWith(result, kKernelInfoAttrName, kernelInfo);
    }
  }

  OpBuilder builder(&getContext());
  SmallVector<Attribute> resultLayouts = llvm::map_to_vector(
      op->getResults(),
      [&](Value result) -> Attribute { return assignedLayouts.at(result); });

  if (op->getNumResults() == 1) {
    op->setAttr(tensor_ext::TensorExtDialect::kLayoutAttrName,
                resultLayouts.front());
    return;
  }
  op->setAttr(tensor_ext::TensorExtDialect::kLayoutAttrName,
              builder.getArrayAttr(resultLayouts));
}

void LayoutPropagation::runOnOperation() {
  DataFlowSolver solver;
  dataflow::loadBaselineAnalyses(solver);
  solver.load<SecretnessAnalysis>();
  if (failed(solver.initializeAndRun(getOperation()))) {
    getOperation()->emitOpError() << "Failed to run secretness analysis.\n";
    signalPassFailure();
    return;
  }
  this->solver = &solver;

  LLVM_DEBUG(llvm::dbgs() << "Running layout propagation on operation: "
                          << getOperation()->getName() << "\n");
  WalkResult result =
      getOperation()->walk<WalkOrder::PreOrder>([&](Operation* op) {
        LogicalResult result = visitOperation(op);
        if (failed(result)) {
          return WalkResult::interrupt();
        }
        return WalkResult::advance();
      });

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

  RewritePatternSet cleanupPatterns(&getContext());
  cleanupPatterns.add<tensor_ext::FoldConvertLayoutIntoAssignLayoutPattern>(
      &getContext());
  (void)applyPatternsGreedily(getOperation(), std::move(cleanupPatterns));
};

}  // namespace heir
}  // namespace mlir
