From 3515e1f109fca4c3d7a79e82dad01b818b9073e0 Mon Sep 17 00:00:00 2001 From: Wenhao Zhang Date: Fri, 28 Aug 2026 22:20:57 -0700 Subject: [PATCH] feat: in-layout strided reduce for bicyclic/tricyclic layout Support in-layout strided reductions along CRT-based packings without requiring physical layout conversion, using binary span-doubling rotate-and-reduce and valid-prefix periodic replication. Also upgrade partial target relations to canonical total congruences in elementwise rectification. PiperOrigin-RevId: 972965062 --- lib/Dialect/Secret/IR/SecretOps.cpp | 2 + lib/Dialect/Secret/IR/SecretPatterns.cpp | 22 +- lib/Dialect/TensorExt/Transforms/BUILD | 2 + .../Transforms/ImplementShiftNetwork.cpp | 93 ++- lib/Kernel/BUILD | 2 + lib/Kernel/Kernel.cpp | 2 + lib/Kernel/Kernel.h | 2 + lib/Kernel/KernelImplementation.h | 39 +- lib/Kernel/KernelName.h | 4 + lib/Kernel/RotateAndReduceFuzzTest.cpp | 3 +- lib/Kernel/RotateAndReduceImplTest.cpp | 51 ++ .../AssignLayout.cpp | 29 +- .../ConvertToCiphertextSemantics.cpp | 382 +++++++++++- lib/Transforms/LayoutPropagation/BUILD | 1 + .../LayoutPropagation/LayoutPropagation.cpp | 553 +++++++++++++++--- lib/Transforms/LayoutPropagation/Utils.cpp | 67 ++- .../LayoutPropagation/UtilsTest.cpp | 24 + lib/Utils/Layout/Utils.cpp | 149 +++++ lib/Utils/Layout/Utils.h | 27 +- lib/Utils/Layout/UtilsTest.cpp | 67 +++ lib/Utils/RotationUtils.h | 19 +- .../Transforms/implement_shift_network.mlir | 31 +- .../batch_matmul_pt.mlir | 54 ++ .../broadcast.mlir | 17 + .../linalg_reduce.mlir | 74 ++- .../transpose.mlir | 17 + .../layout_optimization/pad_to_layout.mlir | 4 +- .../layout_propagation/batch_matmul_pt.mlir | 64 ++ .../layout_propagation/broadcast.mlir | 23 + .../layout_propagation/linalg_reduce.mlir | 38 +- .../layout_propagation/transpose.mlir | 39 ++ 31 files changed, 1704 insertions(+), 197 deletions(-) create mode 100644 tests/Transforms/convert_to_ciphertext_semantics/batch_matmul_pt.mlir create mode 100644 tests/Transforms/convert_to_ciphertext_semantics/broadcast.mlir create mode 100644 tests/Transforms/convert_to_ciphertext_semantics/transpose.mlir create mode 100644 tests/Transforms/layout_propagation/batch_matmul_pt.mlir create mode 100644 tests/Transforms/layout_propagation/broadcast.mlir create mode 100644 tests/Transforms/layout_propagation/transpose.mlir diff --git a/lib/Dialect/Secret/IR/SecretOps.cpp b/lib/Dialect/Secret/IR/SecretOps.cpp index fffb22459c..07bf233a37 100644 --- a/lib/Dialect/Secret/IR/SecretOps.cpp +++ b/lib/Dialect/Secret/IR/SecretOps.cpp @@ -589,6 +589,8 @@ GenericOp GenericOp::extractOpBeforeGeneric(Operation* opToExtract, })) { this->setOperandAttrsAttr( ArrayAttr::get(this->getContext(), newGenericArgAttrs)); + } else { + this->removeAllOperandAttrsAttr(); } }); rewriter.replaceOp(opToExtract, oldGenericNewBlockArgs); diff --git a/lib/Dialect/Secret/IR/SecretPatterns.cpp b/lib/Dialect/Secret/IR/SecretPatterns.cpp index d6718e7370..4fe11f8dd9 100644 --- a/lib/Dialect/Secret/IR/SecretPatterns.cpp +++ b/lib/Dialect/Secret/IR/SecretPatterns.cpp @@ -176,11 +176,15 @@ LogicalResult RemoveUnusedGenericArgs::matchAndRewrite( if (arg.use_empty()) { LLVM_DEBUG(llvm::dbgs() << arg << " has no uses; removing\n"); hasUnusedOps = true; + // Read the operand attrs BEFORE erasing the operand: the accessor + // treats a size-mismatched (stale) array as absent, so reading after + // the erase would both skip the rebuild and leave the stale array on + // the op. + auto attrs = op.getAllOperandAttrsAttr(); rewriter.modifyOpInPlace(op, [&]() { body->eraseArgument(i); op.getOperation()->eraseOperand(i); }); - auto attrs = op.getAllOperandAttrsAttr(); if (attrs) { SmallVector attrList; for (auto [j, attr] : llvm::enumerate(attrs)) { @@ -188,7 +192,12 @@ LogicalResult RemoveUnusedGenericArgs::matchAndRewrite( attrList.push_back(attr); } } - op.setOperandAttrsAttr(ArrayAttr::get(op.getContext(), attrList)); + // An empty array is not "no attrs"; remove it outright so later + // operand appends don't see a stale array. + if (attrList.empty()) + op.removeAllOperandAttrsAttr(); + else + op.setOperandAttrsAttr(ArrayAttr::get(op.getContext(), attrList)); } // Ensure the next iteration uses the right arg number @@ -263,11 +272,13 @@ LogicalResult RemoveNonSecretGenericArgs::matchAndRewrite( BlockArgument correspondingArg = body->getArgument(i); rewriter.replaceAllUsesWith(correspondingArg, op->getOperand(i)); + // Read the operand attrs BEFORE erasing the operand; see + // RemoveUnusedGenericArgs. + auto attrs = op.getAllOperandAttrsAttr(); rewriter.modifyOpInPlace(op, [&]() { body->eraseArgument(i); op.getOperation()->eraseOperand(i); }); - auto attrs = op.getAllOperandAttrsAttr(); if (attrs) { SmallVector attrList; for (auto [j, attr] : llvm::enumerate(attrs)) { @@ -275,7 +286,10 @@ LogicalResult RemoveNonSecretGenericArgs::matchAndRewrite( attrList.push_back(attr); } } - op.setOperandAttrsAttr(ArrayAttr::get(op.getContext(), attrList)); + if (attrList.empty()) + op.removeAllOperandAttrsAttr(); + else + op.setOperandAttrsAttr(ArrayAttr::get(op.getContext(), attrList)); } i--; } diff --git a/lib/Dialect/TensorExt/Transforms/BUILD b/lib/Dialect/TensorExt/Transforms/BUILD index 3240d20c9d..77eafcbb30 100644 --- a/lib/Dialect/TensorExt/Transforms/BUILD +++ b/lib/Dialect/TensorExt/Transforms/BUILD @@ -109,7 +109,9 @@ cc_library( ":pass_inc_gen", "@heir//lib/Dialect/TensorExt/IR:Dialect", "@heir//lib/Kernel:AbstractValue", + "@heir//lib/Kernel:ArithmeticDag", "@heir//lib/Kernel:IRMaterializingVisitor", + "@heir//lib/Utils", "@heir//lib/Utils:MathUtils", "@heir//lib/Utils/ADT:FrozenVector", "@heir//lib/Utils/Graph", diff --git a/lib/Dialect/TensorExt/Transforms/ImplementShiftNetwork.cpp b/lib/Dialect/TensorExt/Transforms/ImplementShiftNetwork.cpp index 9e361cb8d1..061dfcba77 100644 --- a/lib/Dialect/TensorExt/Transforms/ImplementShiftNetwork.cpp +++ b/lib/Dialect/TensorExt/Transforms/ImplementShiftNetwork.cpp @@ -2,11 +2,15 @@ #include #include +#include #include #include #include #include +#include +#include #include +#include #include #include #include @@ -16,15 +20,16 @@ #include "lib/Dialect/TensorExt/Transforms/RotationGroupKernel.h" #include "lib/Dialect/TensorExt/Transforms/ShiftScheme.h" #include "lib/Kernel/AbstractValue.h" +#include "lib/Kernel/ArithmeticDag.h" #include "lib/Kernel/IRMaterializingVisitor.h" #include "lib/Utils/ADT/FrozenVector.h" #include "lib/Utils/Graph/Graph.h" #include "lib/Utils/Layout/Utils.h" +#include "lib/Utils/Utils.h" +#include "llvm/include/llvm/ADT/DenseSet.h" // from @llvm-project #include "llvm/include/llvm/ADT/STLExtras.h" // from @llvm-project #include "llvm/include/llvm/Support/Debug.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project -#include "mlir/include/mlir/IR/Attributes.h" // from @llvm-project #include "mlir/include/mlir/IR/Builders.h" // from @llvm-project #include "mlir/include/mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project @@ -141,7 +146,7 @@ ShiftScheme VosVosErkinShiftNetworks::findBestShiftScheme( // iteration. SmallVector shiftOrder = initShiftOrder; - std::ranges::shuffle(shiftOrder.begin(), shiftOrder.end(), g); + std::shuffle(shiftOrder.begin(), shiftOrder.end(), g); ShiftStrategy strategy = evaluateShiftStrategy(mapping, shiftOrder); @@ -242,12 +247,6 @@ LogicalResult convertRemapOp(RemapOp op, "DenseIntElementsAttr"; } - ShiftScheme scheme = shiftNetworks.findShiftScheme(mapping); - auto rotationGroups = scheme.rotationGroups; - - assert(!rotationGroups.empty() && - "Shift network must have at least one group"); - b.setInsertionPointAfter(op); // Could add a special case here if the numCiphertexts == 1, using @@ -272,8 +271,80 @@ LogicalResult convertRemapOp(RemapOp op, ciphertexts.push_back(kernel::SSAValue(slice.getResult())); } - auto resultNodes = implementShiftNetwork(ciphertexts, mapping, scheme, - minSlotCount, dagElemType); + using NodeTy = kernel::ArithmeticDagNode; + SmallVector> resultNodes; + + const auto targetToSource = mapping.getTargetToSource(); + + // Count the rotations the direct path would emit. + llvm::DenseSet> uniqueShifts; + for (const auto& [target, source] : targetToSource) { + int64_t r = (source.slot - target.slot) % minSlotCount; + if (r < 0) r += minSlotCount; + if (r != 0) uniqueShifts.insert({source.ct, r}); + } + + // If the number of distinct rotation offsets is <= log2(N_slots), direct + // depth-1 rotation + plaintext masking uses at most `uniqueShifts` rotations + // at multiplicative depth 1, vs VVE's multi-stage log2(N) rotations and + // ~2·log2(N) masks. + int64_t maxDirectRotations = + static_cast(std::ceil(std::log2(minSlotCount))); + + if (static_cast(uniqueShifts.size()) <= maxDirectRotations) { + // Build rotation-mask groups. + std::map, std::vector> groups; + for (const auto& [target, source] : targetToSource) { + int64_t r = (source.slot - target.slot) % minSlotCount; + if (r < 0) r += minSlotCount; + auto& mask = groups + .try_emplace(std::make_tuple(target.ct, source.ct, r), + std::vector(minSlotCount, 0.0)) + .first->second; + mask[target.slot] = 1.0; + } + + SmallVector> cts; + for (const auto& ct : ciphertexts) { + cts.push_back(NodeTy::leaf(ct)); + } + + std::map, std::shared_ptr> + rotationCache; + auto maskType = makeTensorType(dagElemType, {1, minSlotCount}); + resultNodes.assign(numCiphertexts, nullptr); + for (auto& [key, mask] : groups) { + auto [targetCt, sourceCt, r] = key; + std::shared_ptr rotated; + if (r == 0) { + rotated = cts[sourceCt]; + } else { + auto [it, inserted] = rotationCache.try_emplace({sourceCt, r}, nullptr); + if (inserted) { + it->second = NodeTy::leftRotate(cts[sourceCt], r); + } + rotated = it->second; + } + auto [allZero, allOne] = allZeroAllOne(mask); + (void)allZero; + std::shared_ptr term = + allOne ? rotated + : NodeTy::mul(rotated, NodeTy::constantTensor(mask, maskType)); + resultNodes[targetCt] = resultNodes[targetCt] + ? NodeTy::add(resultNodes[targetCt], term) + : term; + } + std::vector zeros(minSlotCount, 0.0); + for (auto& node : resultNodes) { + if (!node) node = NodeTy::constantTensor(zeros, maskType); + } + } else { + ShiftScheme scheme = shiftNetworks.findShiftScheme(mapping); + assert(!scheme.rotationGroups.empty() && + "Shift network must have at least one group"); + resultNodes = implementShiftNetwork(ciphertexts, mapping, scheme, + minSlotCount, dagElemType); + } kernel::IRMaterializingVisitor visitor(singleCiphertextType); auto resultVectors = visitor.process(resultNodes, b); diff --git a/lib/Kernel/BUILD b/lib/Kernel/BUILD index 4c1af391bb..0b9d474f75 100644 --- a/lib/Kernel/BUILD +++ b/lib/Kernel/BUILD @@ -138,11 +138,13 @@ cc_test( ":KernelImplementation", ":RotationCountVisitor", "@googletest//:gtest_main", + "@heir//lib/Utils:RotationUtils", "@heir//lib/Utils/Layout:Codegen", "@heir//lib/Utils/Layout:Convolution", "@heir//lib/Utils/Layout:ConvolutionTestUtil", "@heir//lib/Utils/Layout:Evaluate", "@heir//lib/Utils/Layout:Utils", + "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], diff --git a/lib/Kernel/Kernel.cpp b/lib/Kernel/Kernel.cpp index a8e8f8f3e3..7e451954b5 100644 --- a/lib/Kernel/Kernel.cpp +++ b/lib/Kernel/Kernel.cpp @@ -36,6 +36,8 @@ std::string kernelNameAsStr(const KernelName& kernelName) { return "MatmulBicyclic"; case KernelName::MatmulBicyclicDiagonal: return "MatmulBicyclicDiagonal"; + case KernelName::BatchMatmulTricyclicDiagonal: + return "BatchMatmulTricyclicDiagonal"; case KernelName::BatchMatmulTricyclic: return "BatchMatmulTricyclic"; case KernelName::Dot: diff --git a/lib/Kernel/Kernel.h b/lib/Kernel/Kernel.h index 651ef90686..5ec76a5062 100644 --- a/lib/Kernel/Kernel.h +++ b/lib/Kernel/Kernel.h @@ -31,6 +31,8 @@ struct FieldParser { if (kernelName == "MatmulBicyclic") return heir::KernelName::MatmulBicyclic; if (kernelName == "MatmulBicyclicDiagonal") return heir::KernelName::MatmulBicyclicDiagonal; + if (kernelName == "BatchMatmulTricyclicDiagonal") + return heir::KernelName::BatchMatmulTricyclicDiagonal; if (kernelName == "BatchMatmulTricyclic") return heir::KernelName::BatchMatmulTricyclic; if (kernelName == "Dot") return heir::KernelName::Dot; diff --git a/lib/Kernel/KernelImplementation.h b/lib/Kernel/KernelImplementation.h index 9c21078671..b59eeb8351 100644 --- a/lib/Kernel/KernelImplementation.h +++ b/lib/Kernel/KernelImplementation.h @@ -77,13 +77,35 @@ std::enable_if_t::value, implementRotateAndReduceAccumulation( std::shared_ptr> vectorDag, int64_t period, int64_t steps, DagReducer reduceFunc) { + assert(steps >= 1 && "rotate-and-reduce needs at least one step"); using NodeTy = ArithmeticDagNode; - for (int64_t shiftSize = steps / 2; shiftSize > 0; shiftSize /= 2) { - auto rotated = NodeTy::leftRotate(vectorDag, shiftSize * period); - auto reduced = reduceFunc(vectorDag, rotated); - vectorDag = reduced; + if ((steps & (steps - 1)) == 0) { + for (int64_t shiftSize = steps / 2; shiftSize > 0; shiftSize /= 2) { + auto rotated = NodeTy::leftRotate(vectorDag, shiftSize * period); + auto reduced = reduceFunc(vectorDag, rotated); + vectorDag = reduced; + } + return vectorDag; } - return vectorDag; + std::shared_ptr> acc = nullptr; + std::shared_ptr> span = vectorDag; + int64_t spanLen = 1; + int64_t offset = 0; + int64_t remaining = steps; + while (remaining > 0) { + if (remaining & 1) { + auto term = + offset == 0 ? span : NodeTy::leftRotate(span, offset * period); + acc = acc ? reduceFunc(acc, term) : term; + offset += spanLen; + } + remaining >>= 1; + if (remaining > 0) { + span = reduceFunc(span, NodeTy::leftRotate(span, spanLen * period)); + spanLen <<= 1; + } + } + return acc; } template @@ -106,6 +128,13 @@ implementRotateAndReduceAccumulationRolled( using NodeTy = ArithmeticDagNode; using NodePtr = std::shared_ptr; + // Non-power-of-two steps fall back to the unrolled binary span-doubling + // form because the variable shift and rotation sequence cannot be expressed + // as a uniform loop with a single halved shift per iteration. + if ((steps & (steps - 1)) != 0) { + return implementRotateAndReduceAccumulation(vectorDag, period, steps, + reduceFunc); + } int64_t numIterations = static_cast(std::log2(steps)); if (numIterations <= 0) return vectorDag; diff --git a/lib/Kernel/KernelName.h b/lib/Kernel/KernelName.h index cd379d59f0..492d5df1bb 100644 --- a/lib/Kernel/KernelName.h +++ b/lib/Kernel/KernelName.h @@ -39,6 +39,10 @@ enum KernelName : int { // plaintext-lhs-secret-rhs cases. MatmulBicyclicDiagonal, + // Ciphertext-plaintext batch matmul by mapping the plaintext into n + // generalized diagonals. + BatchMatmulTricyclicDiagonal, + // Product and sum of two vectors, using a log2 rotate-and-reduce approach. Dot, }; diff --git a/lib/Kernel/RotateAndReduceFuzzTest.cpp b/lib/Kernel/RotateAndReduceFuzzTest.cpp index e4774ee4b4..3d4af5ad5a 100644 --- a/lib/Kernel/RotateAndReduceFuzzTest.cpp +++ b/lib/Kernel/RotateAndReduceFuzzTest.cpp @@ -177,7 +177,8 @@ FUZZ_TEST(RotateAndReduceFuzzTest, rotateAndReduceWithoutPlaintexts) .WithMinSize(1) .WithMaxSize(32), /*period=*/fuzztest::InRange(1L, 4L), - /*steps=*/fuzztest::ElementOf({1L, 2L, 4L, 8L, 16L}), + /*steps=*/ + fuzztest::ElementOf({1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L, 16L}), fuzztest::Arbitrary()); // Fuzz test for rotate and reduce with plaintexts and zero diagonals diff --git a/lib/Kernel/RotateAndReduceImplTest.cpp b/lib/Kernel/RotateAndReduceImplTest.cpp index d3485c346a..3059866230 100644 --- a/lib/Kernel/RotateAndReduceImplTest.cpp +++ b/lib/Kernel/RotateAndReduceImplTest.cpp @@ -1,6 +1,8 @@ #include #include #include +#include +#include #include #include "gtest/gtest.h" // from @googletest @@ -8,6 +10,8 @@ #include "lib/Kernel/ArithmeticDag.h" #include "lib/Kernel/EvalVisitor.h" #include "lib/Kernel/KernelImplementation.h" +#include "lib/Utils/RotationUtils.h" +#include "llvm/include/llvm/ADT/DenseSet.h" // from @llvm-project namespace mlir { namespace heir { @@ -313,6 +317,53 @@ TEST(RotateAndReduceImplTest, BroadcastedReduce_Masked_Stride) { } } +void collectRotations( + const std::shared_ptr>& node, + llvm::DenseSet& rotations, + std::unordered_set*>& visited) { + if (!node || !visited.insert(node.get()).second) return; + if (auto* rotate = + std::get_if>(&node->node_variant)) { + if (auto* scalar = + std::get_if(&rotate->shift->node_variant)) { + rotations.insert(static_cast(scalar->value)); + } + collectRotations(rotate->operand, rotations, visited); + collectRotations(rotate->shift, rotations, visited); + return; + } + if (auto* add = std::get_if>(&node->node_variant)) { + collectRotations(add->left, rotations, visited); + collectRotations(add->right, rotations, visited); + return; + } +} + +TEST(RotateAndReduceImplTest, TestPredictorSync) { + std::vector dummyVec = {1}; + LiteralValue val(dummyVec); + auto leaf = ArithmeticDagNode::leaf(val); + auto addReducer = [](std::shared_ptr> a, + std::shared_ptr> b) { + return ArithmeticDagNode::add(a, b); + }; + + for (int64_t period : {1, 3, 7}) { + for (int64_t steps = 1; steps <= 100; ++steps) { + auto dag = implementRotateAndReduceAccumulation( + leaf, period, steps, addReducer); + llvm::DenseSet actualRotations; + std::unordered_set*> visited; + collectRotations(dag, actualRotations, visited); + + llvm::DenseSet predicted = rotateAndReduceRotationIndices( + period, steps, /*hasPlaintexts=*/false); + EXPECT_EQ(actualRotations, predicted) + << "Mismatch for period=" << period << ", steps=" << steps; + } + } +} + } // namespace } // namespace kernel } // namespace heir diff --git a/lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp b/lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp index 6c0c9c88b0..ee2e0bfd88 100644 --- a/lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp +++ b/lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.cpp @@ -460,10 +460,22 @@ static FailureOr implementAssignLayoutStep( std::vector rawBuffer( static_cast(numTargetElements) * byteWidth, 0); + std::vector written(srcIsSplat ? 0 : numTargetElements, false); for (const auto& [domainPoint, rangePoint] : collector.points) { int64_t dstFlat = flatten(rangePoint, dstStrides); if (dstFlat < 0 || dstFlat >= numTargetElements) continue; int64_t srcFlat = srcIsSplat ? 0 : flatten(domainPoint, srcStrides); + if (!srcIsSplat && written[dstFlat] && + std::memcmp( + rawBuffer.data() + static_cast(dstFlat) * byteWidth, + srcRaw.data() + static_cast(srcFlat) * byteWidth, + byteWidth) != 0) { + return builder.emitError() + << "layout maps two distinct data values to the same slot " + << dstFlat << "; a non-replicated value cannot be packed " + << "into this (non-injective) layout"; + } + if (!srcIsSplat) written[dstFlat] = true; std::memcpy(rawBuffer.data() + static_cast(dstFlat) * byteWidth, srcRaw.data() + static_cast(srcFlat) * byteWidth, byteWidth); @@ -497,13 +509,21 @@ static FailureOr implementAssignLayoutStep( Attribute splatValue = srcIsSplat ? constantAttr.getSplatValue() : Attribute(); auto srcValues = constantAttr.getValues(); + std::vector written(srcIsSplat ? 0 : numTargetElements, false); for (const auto& [domainPoint, rangePoint] : collector.points) { int64_t dstFlat = flatten(rangePoint, dstStrides); if (dstFlat < 0 || dstFlat >= numTargetElements) continue; - packedValues[dstFlat] = srcIsSplat - ? splatValue - : srcValues[static_cast( - flatten(domainPoint, srcStrides))]; + Attribute val = srcIsSplat ? splatValue + : srcValues[static_cast( + flatten(domainPoint, srcStrides))]; + if (!srcIsSplat && written[dstFlat] && packedValues[dstFlat] != val) { + return builder.emitError() + << "layout maps two distinct data values to the same slot " + << dstFlat << "; a non-replicated value cannot be packed " + << "into this (non-injective) layout"; + } + if (!srcIsSplat) written[dstFlat] = true; + packedValues[dstFlat] = val; } auto packedConstantAttr = @@ -545,6 +565,7 @@ static FailureOr implementAssignLayoutStep( if (succeeded(folded)) { return folded.value(); } + return failure(); } auto zeroOp = arith::ConstantOp::create(builder, targetType, diff --git a/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp b/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp index 89ca8d9673..82ae06d7d2 100644 --- a/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp +++ b/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp @@ -701,17 +701,69 @@ class ConvertLinalgReduce : public ConversionBase { public: using ConversionBase::ConversionBase; - void rotateAndReduceKernel(linalg::ReduceOp op, OpAdaptor adaptor, - ContextAwareConversionPatternRewriter& rewriter, - Operation* innerOp) const { + LogicalResult rotateAndReduceKernel( + linalg::ReduceOp op, OpAdaptor adaptor, + ContextAwareConversionPatternRewriter& rewriter, + Operation* innerOp) const { auto input = op.getInputs()[0]; - auto originalShape = cast(input.getType()).getShape(); + auto inputType = cast(input.getType()); + auto originalShape = inputType.getShape(); // Pre-conditions enforce that the reduction operation occurs along the // inputs single axis unsigned steps = originalShape[op.getDimensions()[0]]; unsigned period = 1; + auto layoutAttr = cast(op->getAttr(kLayoutAttrName)); + auto convertedType = + getTypeConverter()->convertType(op.getResult(0).getType(), layoutAttr); + assert(convertedType && "expected ranked tensor type"); + int64_t numSlots = + cast(adaptor.getInputs()[0].getType()).getDimSize(1); + + // In-layout strided reduce: infer period from bicyclic or tricyclic layout. + int64_t dim = op.getDimensions()[0]; + auto maybeInputLayout = + getTypeConverter()->getContextualAttr(adaptor.getInputs()[0]); + if (failed(maybeInputLayout)) { + return op.emitOpError("input tensor has no assigned layout"); + } + auto inputLayout = dyn_cast(maybeInputLayout.value()); + if (!inputLayout) { + return op.emitOpError("input tensor layout attribute is invalid"); + } + + auto expected = convertLayoutForReduce(inputLayout, op.getDimensions()); + if (cast(op->getAttr(kLayoutAttrName)) != expected) { + return op.emitOpError() + << "result layout does not match the reduce of the input " + "layout; propagation and conversion disagree (got " + << op->getAttr(kLayoutAttrName) << " expected " << expected << ")"; + } + + const IntegerRelation& inputRel = inputLayout.getIntegerRelation(); + if (inputType.getRank() == 2 && dim == 1) { + if (isRelationBicyclic(inputType, numSlots, inputRel)) { + period = originalShape[0]; + } else if (!isRelationRowMajor(inputType, numSlots, inputRel)) { + return op.emitOpError( + "input layout is neither row-major nor the bicyclic congruence"); + } + } else if (inputType.getRank() == 3 && dim == 2) { + if (isRelationTricyclic(inputType, numSlots, inputRel)) { + period = originalShape[0] * originalShape[1]; + } else if (!isRelationRowMajor(inputType, numSlots, inputRel)) { + return op.emitOpError( + "input layout is neither row-major nor the tricyclic congruence"); + } + } else if (!isRelationRowMajor(inputType, numSlots, inputRel)) { + return op.emitOpError( + "unsupported reduce dimension for non-row-major input layout"); + } else if (dim != inputType.getRank() - 1) { + return op.emitOpError( + "row-major reduce only supported along innermost dimension"); + } + SSAValue vectorLeaf(adaptor.getInputs()[0]); kernel::DagType dagType = kernel::mlirTypeToDagType(input.getType()); @@ -720,16 +772,52 @@ class ConvertLinalgReduce : public ConversionBase { innerOp->getName().getStringRef().str()); rewriter.setInsertionPointAfter(op); - auto layoutAttr = cast(op->getAttr(kLayoutAttrName)); - auto convertedType = - getTypeConverter()->convertType(input.getType(), layoutAttr); - Value finalOutput = materializeKernel( rewriter, op.getLoc(), implementedKernel, convertedType, layoutAttr); // Add the initial value. Value result = adaptor.getInits()[0]; - addBiasAndReplace(rewriter, op, finalOutput, result, layoutAttr); + if (period == 1) { + addBiasAndReplace(rewriter, op, finalOutput, result, layoutAttr); + return success(); + } + + // In-layout strided reduce: the rotation circuit (period P_kept over + // `steps` hops) wraps garbage into the tail exactly like the diagonal + // matmul kernels, so it joins the same valid-prefix replication + // contract. The result's content is periodic with period P_kept (the + // sum at slot s depends only on s mod P_kept: the strided hops sweep + // every reduced residue), so P_kept is the replication period. + ImplicitLocOpBuilder b(op->getLoc(), rewriter); + Operation* addBias = + makeAppropriatelyTypedAddOp(b, op->getLoc(), finalOutput, result); + addBias->setAttr(kLayoutAttrName, cast(layoutAttr)); + setMaterializedAttr(addBias); + auto ctSemanticResultType = + cast(addBias->getResult(0).getType()); + int64_t reach = period * (steps - 1); + int64_t validPrefix = ctSemanticResultType.getDimSize(1) - reach; + LLVM_DEBUG(llvm::dbgs() + << "Strided reduce valid prefix: " << validPrefix << "\n"); + if (validPrefix < period) { + return op.emitOpError() + << "strided reduce reach " << reach << " exceeds the " << numSlots + << "-slot budget (validPrefix " << validPrefix << " < period " + << period << ")"; + } + // By CRT packing, numSlots >= period * steps. + // Since reach = period * (steps - 1), + // validPrefix = numSlots - reach >= period * steps - period * (steps - 1) = + // period. + assert( + validPrefix >= period && + "valid prefix must be at least period to allow periodic replication"); + Operation* replicated = replicateValidPrefixOfResult( + b, addBias->getResult(0), cast(layoutAttr), + /*inputPeriod=*/period, validPrefix); + setMaterializedAttr(replicated); + rewriter.replaceOp(op, replicated); + return success(); } LogicalResult matchAndRewrite( @@ -771,8 +859,7 @@ class ConvertLinalgReduce : public ConversionBase { op, "missing new layout attribute for input"); // Based on ImplementRotateAndReduce - rotateAndReduceKernel(op, adaptor, rewriter, innerOp); - return success(); + return rotateAndReduceKernel(op, adaptor, rewriter, innerOp); } }; @@ -829,6 +916,139 @@ struct ConvertLinalgDot : public ConversionBase { } }; +// Lowers linalg.broadcast under bicyclic or tricyclic packing where the +// broadcast acts as a zero-cost view without moving ciphertext data. +class ConvertLinalgBroadcast + : public ContextAwareOpConversionPattern { + public: + using ContextAwareOpConversionPattern< + linalg::BroadcastOp>::ContextAwareOpConversionPattern; + + LogicalResult matchAndRewrite( + linalg::BroadcastOp op, OpAdaptor adaptor, + ContextAwareConversionPatternRewriter& rewriter) const final { + auto resultLayout = + dyn_cast_or_null(op->getAttr(kLayoutAttrName)); + if (!resultLayout) { + return rewriter.notifyMatchFailure(op, + "op has no assigned layout attribute"); + } + + auto layoutLookup = + getTypeConverter()->getContextualAttr(adaptor.getInput()); + if (failed(layoutLookup)) { + return rewriter.notifyMatchFailure( + op, "input layout not found in contextual type converter"); + } + auto inputLayout = dyn_cast(layoutLookup.value()); + if (!inputLayout) { + return rewriter.notifyMatchFailure( + op, "input contextual attribute is not a LayoutAttr"); + } + + // Verify the free-view relation proof: the result relation must be a + // subset of the expected free broadcast relation. + auto resultType = cast(op->getResult(0).getType()); + IntegerRelation expectedRel = inputLayout.getIntegerRelation(); + SmallVector dims(op.getDimensions()); + llvm::sort(dims); + for (int64_t dim : dims) { + unsigned newVar = expectedRel.insertVar(presburger::VarKind::Domain, dim); + expectedRel.addBound(presburger::BoundType::LB, newVar, 0); + expectedRel.addBound(presburger::BoundType::UB, newVar, + resultType.getDimSize(dim) - 1); + } + + if (!resultLayout.getIntegerRelation().isSubsetOf(expectedRel)) { + return rewriter.notifyMatchFailure( + op, + "result layout relation is not a subset of the expected broadcast " + "relation"); + } + + Type convertedResultType = getTypeConverter()->convertType( + op->getResult(0).getType(), resultLayout); + if (!convertedResultType || + convertedResultType != adaptor.getInput().getType()) { + return rewriter.notifyMatchFailure( + op, "converted result ciphertext type does not match input type"); + } + + // Persist the result layout on a no-op cast. + auto castOp = UnrealizedConversionCastOp::create( + rewriter, op.getLoc(), convertedResultType, adaptor.getInput()); + setMaterializedAttr(castOp); + setAttributeAssociatedWith(castOp.getResult(0), kLayoutAttrName, + resultLayout); + rewriter.replaceOp(op, castOp); + return success(); + } +}; + +// A linalg.transpose whose result layout was derived by layout propagation +// (the input relation with permuted domain variables) packs to the exact +// same ciphertext contents as its input; at ciphertext semantics the op is +// a no-op forward, like the broadcast above. +class ConvertLinalgTranspose + : public ContextAwareOpConversionPattern { + public: + using ContextAwareOpConversionPattern< + linalg::TransposeOp>::ContextAwareOpConversionPattern; + + LogicalResult matchAndRewrite( + linalg::TransposeOp op, OpAdaptor adaptor, + ContextAwareConversionPatternRewriter& rewriter) const final { + auto resultLayout = + dyn_cast_or_null(op->getAttr(kLayoutAttrName)); + if (!resultLayout) { + return rewriter.notifyMatchFailure(op, + "op has no assigned layout attribute"); + } + + auto layoutLookup = + getTypeConverter()->getContextualAttr(adaptor.getInput()); + if (failed(layoutLookup)) { + return rewriter.notifyMatchFailure( + op, "input layout not found in contextual type converter"); + } + auto inputLayout = dyn_cast(layoutLookup.value()); + if (!inputLayout) { + return rewriter.notifyMatchFailure( + op, "input contextual attribute is not a LayoutAttr"); + } + + // Verify the transpose relation proof: the result relation must be a + // subset of the expected permuted input relation. + IntegerRelation expectedRel = getTransposedRelation( + inputLayout.getIntegerRelation(), op.getPermutation()); + if (!resultLayout.getIntegerRelation().isSubsetOf(expectedRel)) { + return rewriter.notifyMatchFailure( + op, + "result layout relation is not a subset of the expected transpose " + "relation"); + } + + Type convertedResultType = getTypeConverter()->convertType( + op->getResult(0).getType(), resultLayout); + if (!convertedResultType || + convertedResultType != adaptor.getInput().getType()) { + return rewriter.notifyMatchFailure( + op, "converted result ciphertext type does not match input type"); + } + // Persist the result layout on a no-op cast (mirroring the + // collapse/expand patterns): a bare forward leaves downstream + // consumers looking up the INPUT's contextual layout, whose domain + // order predates the transpose. + auto castOp = UnrealizedConversionCastOp::create( + rewriter, op.getLoc(), convertedResultType, adaptor.getInput()); + setMaterializedAttr(castOp); + setAttributeAssociatedWith(castOp.getResult(0), kLayoutAttrName, + resultLayout); + rewriter.replaceOp(op, castOp); + return success(); + } +}; + struct ConvertLinalgMatvecLayout : public ConversionBase { public: using ConversionBase::ConversionBase; @@ -2788,15 +3008,21 @@ struct ConvertLinalgMatmul << op << "\n"); // Determine if the lhs or rhs is the secret operand. + // Plaintext operands are packed via tensor_ext.assign_layout. + auto isCleartext = [](Value v) { + return v.getDefiningOp() != nullptr; + }; bool secretLhs = false; - auto genericOp = op->getParentOfType(); - if (genericOp) { + if (isCleartext(op.getInputs()[1]) && !isCleartext(op.getInputs()[0])) { + secretLhs = true; + } else if (!isCleartext(op.getInputs()[1]) && + isCleartext(op.getInputs()[0])) { + secretLhs = false; + } else if (auto genericOp = op->getParentOfType()) { if (auto blockArg = dyn_cast(op.getInputs()[0])) { if (blockArg.getArgNumber() < genericOp.getNumOperands()) { auto opArg = genericOp.getOperand(blockArg.getArgNumber()); - if (isa(opArg.getType())) { - secretLhs = true; - } + secretLhs = isa(opArg.getType()); } } } @@ -2990,8 +3216,9 @@ struct ConvertLinalgBatchMatmul rewriter.setInsertionPointAfter(op); ImplicitLocOpBuilder b(op.getLoc(), rewriter); - IRMaterializingVisitor visitor( - lhs.getType(), [&](Operation* createdOp) { setMaterializedAttr(op); }); + IRMaterializingVisitor visitor(lhs.getType(), [&](Operation* createdOp) { + setMaterializedAttr(createdOp); + }); Value finalOutput = visitor.process(implementedKernel, b)[0]; auto layoutAttr = cast(op->getAttr(kLayoutAttrName)); @@ -3028,9 +3255,114 @@ struct ConvertLinalgBatchMatmul rewriter.replaceOp(op, replicated); } + bool supportsTricyclicDiagonal(linalg::BatchMatmulOp op) const { + auto kernelAttr = op->getAttrOfType( + secret::SecretDialect::kKernelAttrName); + return kernelAttr && + kernelAttr.getName() == KernelName::BatchMatmulTricyclicDiagonal; + } + + // Ciphertext-plaintext batch matmul using tricyclic diagonal + // rotate-and-reduce. + void tricyclicDiagonalBatchKernel( + linalg::BatchMatmulOp op, OpAdaptor adaptor, + ContextAwareConversionPatternRewriter& rewriter) const { + LLVM_DEBUG(llvm::dbgs() + << "Converting linalg.batch_matmul with tricyclic diagonal " + "kernel: " + << op << "\n"); + + // Determine if the lhs or rhs is the secret operand. + // Plaintext operands are packed via tensor_ext.assign_layout. + auto isCleartext = [](Value v) { + return v.getDefiningOp() != nullptr; + }; + bool secretLhs = false; + if (isCleartext(op.getInputs()[1]) && !isCleartext(op.getInputs()[0])) { + secretLhs = true; + } else if (!isCleartext(op.getInputs()[1]) && + isCleartext(op.getInputs()[0])) { + secretLhs = false; + } else if (auto genericOp = op->getParentOfType()) { + if (auto blockArg = dyn_cast(op.getInputs()[0])) { + if (blockArg.getArgNumber() < genericOp.getNumOperands()) { + auto opArg = genericOp.getOperand(blockArg.getArgNumber()); + secretLhs = isa(opArg.getType()); + } + } + } + + TypedValue ct = cast>( + adaptor.getInputs()[secretLhs ? 0 : 1]); + SSAValue ctLeaf(ct); + TypedValue pt = cast>( + adaptor.getInputs()[secretLhs ? 1 : 0]); + SSAValue ptLeaf(pt); + + auto lhsType = cast(op.getInputs()[0].getType()); + auto rhsType = cast(op.getInputs()[1].getType()); + assert(lhsType.getDimSize(0) == rhsType.getDimSize(0) && + "batch matrix multiplication inputs must share the same batch " + "dimension"); + auto secretType = ct.getType(); + + auto dagType = kernel::mlirTypeToDagType(secretType); + int64_t period = secretLhs ? lhsType.getDimSize(0) * lhsType.getDimSize(1) + : rhsType.getDimSize(0) * rhsType.getDimSize(2); + int64_t steps = secretLhs ? lhsType.getDimSize(2) : rhsType.getDimSize(1); + std::string reduceOp = isa(secretType.getElementType()) + ? "arith.addf" + : "arith.addi"; + + std::shared_ptr> implementedKernel = + implementRotateAndReduce(ctLeaf, std::optional(ptLeaf), + period, steps, dagType, /*zeroDiagonals=*/{}, + reduceOp); + + rewriter.setInsertionPointAfter(op); + ImplicitLocOpBuilder b(op.getLoc(), rewriter); + IRMaterializingVisitor visitor(ct.getType(), [&](Operation* createdOp) { + setMaterializedAttr(createdOp); + }); + Value finalOutput = visitor.process(implementedKernel, b)[0]; + + auto layoutAttr = cast(op->getAttr(kLayoutAttrName)); + auto* finalOutputOp = finalOutput.getDefiningOp(); + finalOutputOp->setAttr(kLayoutAttrName, layoutAttr); + setMaterializedAttr(finalOutputOp); + + // Add the initial accumulator value. + Value result = adaptor.getOutputs()[0]; + + Operation* addBias = + makeAppropriatelyTypedAddOp(b, op->getLoc(), finalOutput, result); + addBias->setAttr(kLayoutAttrName, layoutAttr); + setMaterializedAttr(addBias); + + // Rebuild the full periodic output layout from the widest valid + // period-aligned window. + auto dataSemanticResultType = + cast(op->getResult(0).getType()); + int64_t reach = period * (steps - 1); + auto ctSemanticResultType = + cast(addBias->getResult(0).getType()); + int64_t validPrefix = ctSemanticResultType.getDimSize(1) - reach; + LLVM_DEBUG(llvm::dbgs() << "Tricyclic diagonal matmul valid prefix: " + << validPrefix << "\n"); + Operation* replicated = replicateValidPrefixOfResult( + b, addBias->getResult(0), layoutAttr, + dataSemanticResultType.getNumElements(), validPrefix); + setMaterializedAttr(replicated); + rewriter.replaceOp(op, replicated); + } + LogicalResult matchAndRewrite( linalg::BatchMatmulOp op, OpAdaptor adaptor, ContextAwareConversionPatternRewriter& rewriter) const final { + if (supportsTricyclicDiagonal(op)) { + tricyclicDiagonalBatchKernel(op, adaptor, rewriter); + return success(); + } if (supportsTricyclic(op, adaptor)) { tricyclicKernel(op, adaptor, rewriter); return success(); @@ -3082,13 +3414,13 @@ struct ConvertToCiphertextSemantics patterns.add< ConvertAnyAddingMaterializedAttr, ConvertBootstrap, - ConvertConvertLayout, ConvertFunc, ConvertLinalgMatmul, - ConvertLinalgBatchMatmul, ConvertLinalgReduce, ConvertLinalgDot, - ConvertSecretGeneric, ConvertTensorCollapseShape, - ConvertTensorExpandShape, ConvertTensorExtractLayout, - ConvertTensorExtractSlice, ConvertTensorPad, ConvertTensorInsertLayout, - ConvertTensorInsertSlice, PreserveLinalgMatvecAsLinearTransform>( - typeConverter, context); + ConvertConvertLayout, ConvertFunc, ConvertLinalgBroadcast, + ConvertLinalgMatmul, ConvertLinalgTranspose, ConvertLinalgBatchMatmul, + ConvertLinalgReduce, ConvertLinalgDot, ConvertSecretGeneric, + ConvertTensorCollapseShape, ConvertTensorExpandShape, + ConvertTensorExtractLayout, ConvertTensorExtractSlice, ConvertTensorPad, + ConvertTensorInsertLayout, ConvertTensorInsertSlice, + PreserveLinalgMatvecAsLinearTransform>(typeConverter, context); patterns.add(typeConverter, context, diff --git a/lib/Transforms/LayoutPropagation/BUILD b/lib/Transforms/LayoutPropagation/BUILD index eb58dd10f7..512d78561c 100644 --- a/lib/Transforms/LayoutPropagation/BUILD +++ b/lib/Transforms/LayoutPropagation/BUILD @@ -63,6 +63,7 @@ cc_library( hdrs = ["Utils.h"], deps = [ "@heir//lib/Dialect/TensorExt/IR:Dialect", + "@heir//lib/Utils/Layout:Utils", "@llvm-project//llvm:Support", "@llvm-project//mlir:Analysis", "@llvm-project//mlir:IR", diff --git a/lib/Transforms/LayoutPropagation/LayoutPropagation.cpp b/lib/Transforms/LayoutPropagation/LayoutPropagation.cpp index af4610771b..3ee7b20ca1 100644 --- a/lib/Transforms/LayoutPropagation/LayoutPropagation.cpp +++ b/lib/Transforms/LayoutPropagation/LayoutPropagation.cpp @@ -41,6 +41,7 @@ #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/Arith/IR/Arith.h" // from @llvm-project #include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/include/mlir/Dialect/Linalg/IR/Linalg.h" // from @llvm-project #include "mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h" // from @llvm-project @@ -68,6 +69,7 @@ namespace mlir { namespace heir { using linalg::BatchMatmulOp; +using linalg::BroadcastOp; using linalg::Conv1DNcwFcwOp; using linalg::Conv1DOp; using linalg::Conv2DNchwFchwOp; @@ -76,6 +78,7 @@ using linalg::DotOp; using linalg::MatmulOp; using linalg::MatvecOp; using linalg::ReduceOp; +using linalg::TransposeOp; using linalg::VecmatOp; using presburger::IntegerRelation; using secret::GenericOp; @@ -274,6 +277,8 @@ struct LayoutPropagation : impl::LayoutPropagationBase { // Op-specific transfer functions LogicalResult visitOperation(CollapseShapeOp op); LogicalResult visitOperation(ExpandShapeOp op); + LogicalResult visitOperation(BroadcastOp op); + LogicalResult visitOperation(TransposeOp op); LogicalResult visitOperation(GenericOp op); LogicalResult visitOperation(ReduceOp op); LogicalResult visitOperation(Conv1DOp op); @@ -311,6 +316,7 @@ struct LayoutPropagation : impl::LayoutPropagationBase { CompatibilityResult hasCompatibleArgumentLayouts(VecmatOp op); CompatibilityResult hasCompatibleArgumentLayouts(MatvecOp op); CompatibilityResult hasCompatibleArgumentLayouts(MatmulOp op); + CompatibilityResult hasCompatibleArgumentLayouts(BatchMatmulOp op); CompatibilityResult hasCompatibleArgumentLayouts(tensor::InsertSliceOp op); // Insert conversion ops to rectify incompatible operand layouts @@ -372,7 +378,44 @@ struct LayoutPropagation : impl::LayoutPropagationBase { FailureOr LayoutPropagation::assignDefaultLayoutForOpOperand( Operation* op, Value operand, IRRewriter& builder) { - FailureOr layout = defaultLayoutForType(operand.getType()); + FailureOr layout = failure(); + auto existingLayout = findAttributeAssociatedWith( + operand, tensor_ext::TensorExtDialect::kLayoutAttrName); + if (succeeded(existingLayout)) { + if (auto layoutAttr = dyn_cast(*existingLayout)) { + layout = layoutAttr; + } + } + + // Inherit layout from matching sibling operand to avoid conversions. + if (failed(layout)) { + bool isSplatConstant = false; + if (auto cst = operand.getDefiningOp()) { + if (auto dense = dyn_cast(cst.getValue())) { + isSplatConstant = dense.isSplat(); + } + } + for (Value other : op->getOperands()) { + if (other == operand || other.getType() != operand.getType()) continue; + auto it = assignedLayouts.find(other); + if (it == assignedLayouts.end()) continue; + if (auto otherAttr = dyn_cast(it->second)) { + // A non-injective (view) layout cannot represent an arbitrary value of + // this type, as packing would collapse distinct elements; allow only + // splat constants to inherit view layouts. + if (!isSplatConstant && + !isRelationInjective(otherAttr.getIntegerRelation())) { + continue; + } + layout = otherAttr; + break; + } + } + } + + if (failed(layout)) { + layout = defaultLayoutForType(operand.getType()); + } if (failed(layout)) { return failure(); } @@ -461,9 +504,9 @@ LogicalResult LayoutPropagation::visitOperation(Operation* op) { // secret ops .Case([&](auto op) { return visitOperation(op); }) // linalg ops - .Case( - [&](auto op) { return visitOperation(op); }) + .Case([&](auto op) { return visitOperation(op); }) // affine ops .Case([&](auto op) { return visitOperation(op); }) // tensor ops @@ -485,7 +528,17 @@ LogicalResult LayoutPropagation::visitOperation(func::FuncOp op) { // assign_layout ops and materialized to plaintexts server-side. continue; } - FailureOr layout = defaultLayoutForType(arg.getType()); + FailureOr layout = failure(); + auto existingLayout = findAttributeAssociatedWith( + arg, tensor_ext::TensorExtDialect::kLayoutAttrName); + if (succeeded(existingLayout)) { + if (auto layoutAttr = dyn_cast(*existingLayout)) { + layout = layoutAttr; + } + } + if (failed(layout)) { + layout = defaultLayoutForType(arg.getType()); + } if (failed(layout)) { return op->emitOpError() << "Failed to assign default layout to func argument " << arg; @@ -718,6 +771,64 @@ LogicalResult LayoutPropagation::visitOperation(VecmatOp op) { return success(); } +LogicalResult LayoutPropagation::visitOperation(BroadcastOp op) { + auto input = op.getInput(); + LayoutAttr inputLayout = getComposedLayoutAttr(input); + Value result = op->getResult(0); + auto resultType = cast(result.getType()); + MLIRContext* ctx = &getContext(); + + // The result layout is the input relation with one new free domain variable + // per broadcast dimension (bounded [0, size)), not otherwise constrained by + // the (ct, slot) map: every broadcasted index reads its source element's + // slot, so the packed data is unchanged and the broadcast itself materializes + // as a free view. A consumer that needs the operand in a materialized + // (bijective) layout gets a convert_layout via the ordinary rectification + // path. + IntegerRelation relation = inputLayout.getIntegerRelation(); + SmallVector dims(op.getDimensions()); + llvm::sort(dims); + for (int64_t dim : dims) { + unsigned newVar = relation.insertVar(presburger::VarKind::Domain, dim); + relation.addBound(presburger::BoundType::LB, newVar, 0); + relation.addBound(presburger::BoundType::UB, newVar, + resultType.getDimSize(dim) - 1); + } + + LayoutAttr resultLayoutAttr = + LayoutAttr::getFromIntegerRelation(ctx, relation); + Attribute kernelInfoAttr = + cloneKernelInfoWithResultShape(input, resultType.getShape()); + assignedLayouts.insert({result, resultLayoutAttr}); + setResultLayoutAttr(op, kernelInfoAttr); + debugAssignLayout(result, resultLayoutAttr); + return success(); +} + +LogicalResult LayoutPropagation::visitOperation(TransposeOp op) { + // The result layout is the input relation with its domain variables + // permuted: result index (j_0, ..., j_r) reads the slot of input index + // (j_perm[0], ..., j_perm[r]), so the packed data is unchanged and the + // transpose materializes as a retype. Consumers needing a materialized + // layout get a convert_layout from the ordinary rectification path. + auto input = op.getInput(); + LayoutAttr inputLayout = getComposedLayoutAttr(input); + IntegerRelation relation = getTransposedRelation( + inputLayout.getIntegerRelation(), op.getPermutation()); + + Value result = op->getResult(0); + auto resultType = cast(result.getType()); + MLIRContext* ctx = &getContext(); + LayoutAttr resultLayoutAttr = + LayoutAttr::getFromIntegerRelation(ctx, relation); + Attribute kernelInfoAttr = + cloneKernelInfoWithResultShape(input, resultType.getShape()); + assignedLayouts.insert({result, resultLayoutAttr}); + setResultLayoutAttr(op, kernelInfoAttr); + debugAssignLayout(result, resultLayoutAttr); + return success(); +} + LogicalResult LayoutPropagation::visitOperation(MatvecOp op) { auto matvecOp = cast(*op); auto matrix = matvecOp.lhs(); @@ -1283,34 +1394,154 @@ LogicalResult LayoutPropagation::visitOperation(BatchMatmulOp op) { 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 hDim = 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 = lhsType.getDimSize(0) == rhsType.getDimSize(0); - bool batchEqual = (hLhs == hRhs); - bool lhsCoprime = std::gcd(hLhs, mDim) == 1 && std::gcd(hLhs, nDim) == 1 && + bool lhsCoprime = std::gcd(hDim, mDim) == 1 && std::gcd(hDim, nDim) == 1 && std::gcd(mDim, nDim) == 1; - bool rhsCoprime = std::gcd(hRhs, nDim) == 1 && std::gcd(hRhs, pDim) == 1 && + bool rhsCoprime = std::gcd(hDim, nDim) == 1 && std::gcd(hDim, pDim) == 1 && std::gcd(nDim, pDim) == 1; - bool outputCoprime = std::gcd(hLhs, mDim) == 1 && std::gcd(hLhs, pDim) == 1 && + bool outputCoprime = std::gcd(hDim, mDim) == 1 && std::gcd(hDim, pDim) == 1 && std::gcd(mDim, pDim) == 1; + LLVM_DEBUG(llvm::dbgs() << "lhs=" << lhs << ";\nrhs=" << rhs << "\n"); + + // Tricyclic ct-pt / pt-ct batch matmul with generalized-diagonal-layout + // plaintext: the secret operand keeps its tricyclic packing, the weight + // packs as encode-time diagonals, and the BSGS reach is period * (steps - 1). + if (batchEqual && + ((inputSecret && !filterSecret) || (!inputSecret && filterSecret))) { + bool secretLhs = inputSecret; + Value secretOperand = secretLhs ? lhs : rhs; + Value weightOperand = secretLhs ? rhs : lhs; + RankedTensorType secretType = secretLhs ? lhsType : rhsType; + RankedTensorType weightType = secretLhs ? rhsType : lhsType; + RankedTensorType bmmOutputType = cast(result.getType()); + + // Coprimality is required for the secret operand and result layouts. + // The cleartext weight does not require internal coprimality as it + // packs into diagonals. + bool isCoprime = outputCoprime && (secretLhs ? lhsCoprime : rhsCoprime); + if (isCoprime) { + LayoutAttr secretLayout = getComposedLayoutAttr(secretOperand); + IntegerRelation secretTricyclic = + getTricyclicLayoutRelation(secretType, minSlotCount); + + // The broadcast bicyclic input contains the h-free relation (the + // (dim1, dim2) sub-tensor with the batch var free). Proving h-free + // containment allows refining the layout to tricyclic in place. + IntegerRelation hFree = getBicyclicLayoutRelation( + RankedTensorType::get( + {secretType.getDimSize(1), secretType.getDimSize(2)}, + secretType.getElementType()), + minSlotCount); + unsigned hVar = hFree.insertVar(presburger::VarKind::Domain, 0); + hFree.addBound(presburger::BoundType::LB, hVar, 0); + hFree.addBound(presburger::BoundType::UB, hVar, hDim - 1); + + bool isTricyclic = isRelationTricyclic(secretType, minSlotCount, + secretLayout.getIntegerRelation()); + bool isBroadcast = + !isTricyclic && + isRelationSubset(hFree, secretLayout.getIntegerRelation()); + bool needsConvert = !isTricyclic && !isBroadcast; + + int64_t ctStride = secretLhs ? mDim : pDim; + int64_t period = hDim * ctStride; + int64_t paddedFreeDim = secretLhs ? pDim : mDim; + int64_t contractionDim = secretLhs ? 1 : 2; + + int64_t reach = period * (nDim - 1); + if (reach + bmmOutputType.getNumElements() > minSlotCount) { + return builder.notifyMatchFailure( + op, "slot count budget exceeded for batch matmul diagonal reach"); + } + + if (isBroadcast) { + // The broadcast bicyclic input contains the h-free relation; + // refine layout in place without layout conversion. + LayoutAttr refined = + LayoutAttr::getFromIntegerRelation(ctx, secretTricyclic); + assignedLayouts[secretOperand] = refined; + setAttributeAssociatedWith( + secretOperand, tensor_ext::TensorExtDialect::kLayoutAttrName, + refined); + debugAssignLayout(secretOperand, refined); + } else if (needsConvert) { + auto [toReplace, newSecretLayoutAttr] = convertToLayout( + ctx, builder, op, secretOperand, secretLayout, secretTricyclic); + debugAssignLayout(toReplace, newSecretLayoutAttr); + assignedLayouts.insert({toReplace, newSecretLayoutAttr}); + } + + auto assignOrUpdateLayout = [&](Value val, const IntegerRelation& rel) { + LayoutAttr layout = LayoutAttr::getFromIntegerRelation(ctx, rel); + if (auto assignOp = val.getDefiningOp(); + assignOp && assignOp->hasOneUse()) { + assignOp.setLayoutAttr(layout); + assignedLayouts[val] = layout; + setAttributeAssociatedWith( + val, tensor_ext::TensorExtDialect::kLayoutAttrName, layout); + debugAssignLayout(val, layout); + return val; + } + builder.setInsertionPoint(op); + AssignLayoutOp assignOp = + AssignLayoutOp::create(builder, op->getLoc(), val, layout); + setAttributeAssociatedWith( + assignOp.getResult(), tensor_ext::TensorExtDialect::kLayoutAttrName, + layout); + Value toReplace = assignOp.getResult(); + builder.replaceUsesWithIf(val, toReplace, [&](OpOperand& other) { + return other.getOwner() == op; + }); + assignedLayouts.insert({toReplace, layout}); + debugAssignLayout(toReplace, layout); + return toReplace; + }; + + // Pack plaintext weight into diagonals walking the contraction dimension. + IntegerRelation diagRelation = getTricyclicDiagonalRelation( + weightType, contractionDim, ctStride, paddedFreeDim, minSlotCount); + LayoutAttr weightLayout = getComposedLayoutAttr(weightOperand); + LayoutAttr diagLayoutAttr = + LayoutAttr::getFromIntegerRelation(ctx, diagRelation); + if (weightLayout != diagLayoutAttr) { + assignOrUpdateLayout(weightOperand, diagRelation); + } + + LayoutAttr outputLayoutAttr = LayoutAttr::getFromIntegerRelation( + ctx, getTricyclicLayoutRelation(bmmOutputType, minSlotCount)); + assignedLayouts.insert({result, outputLayoutAttr}); + debugAssignLayout(result, outputLayoutAttr); + + // The kernel adds the accumulator (init) directly to the output via SIMD + // addition. Align the init layout with the result layout at compile time + // so cleartext constants (zeros/bias) re-pack freely without conversions. + Value init = op.getOutputs()[0]; + LayoutAttr initLayout = getComposedLayoutAttr(init); + if (initLayout != outputLayoutAttr) { + assignOrUpdateLayout(init, outputLayoutAttr.getIntegerRelation()); + } + + setResultLayoutAttr(op); + auto kernelAttr = secret::KernelAttr::get( + ctx, KernelName::BatchMatmulTricyclicDiagonal, /*force=*/false); + op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr); + return success(); + } + } + + if (!inputSecret || !filterSecret) { + return builder.notifyMatchFailure( + op, + "batch matrix multiplication with cleartext operand requires pairwise " + "coprime dimensions"); + } + // Tricyclic ct-ct batch matmul. if (inputSecret && filterSecret && batchEqual && lhsCoprime && rhsCoprime && outputCoprime) { @@ -1416,7 +1647,10 @@ LogicalResult LayoutPropagation::visitOperation(MatmulOp op) { return success(); } - // Bicyclic ct-pt matmul with generalized-diagonal-layout plaintext. + // Bicyclic ct-pt / pt-ct matmul with generalized-diagonal-layout + // plaintext: the secret operand keeps its bicyclic packing, the weight + // packs as encode-time diagonals, and the BSGS reach is period * (steps - 1) + // with period the secret operand's packed row (or column) count. if (outputCoprime && ((inputSecret && !filterSecret && lhsCoprime) || (!inputSecret && filterSecret && rhsCoprime))) { Value secretOperand = inputSecret ? lhs : rhs; @@ -1434,35 +1668,69 @@ LogicalResult LayoutPropagation::visitOperation(MatmulOp op) { assignedLayouts.insert({toReplace, newSecretLayoutAttr}); } - // Assign a generalized-diagonal layout to the plaintext matrix. + // The weight packs as the generalized-diagonal matrix; a conversion + // on a constant folds into its assign_layout. IntegerRelation diagRelation = inputSecret - ? getBicyclicDiagonalRelation(weightType, /*contractionDim=*/0, + ? getBicyclicDiagonalRelation(weightType, + /*contractionDim=*/0, /*stride=*/mDim, minSlotCount) - : getBicyclicDiagonalRelation(weightType, /*contractionDim=*/1, + : getBicyclicDiagonalRelation(weightType, + /*contractionDim=*/1, /*stride=*/pDim, minSlotCount); + auto assignOrUpdateLayout = [&](Value val, const IntegerRelation& rel) { + LayoutAttr layout = LayoutAttr::getFromIntegerRelation(ctx, rel); + if (auto assignOp = val.getDefiningOp(); + assignOp && assignOp->hasOneUse()) { + assignOp.setLayoutAttr(layout); + assignedLayouts[val] = layout; + setAttributeAssociatedWith( + val, tensor_ext::TensorExtDialect::kLayoutAttrName, layout); + debugAssignLayout(val, layout); + return val; + } + builder.setInsertionPoint(op); + AssignLayoutOp assignOp = + AssignLayoutOp::create(builder, op->getLoc(), val, layout); + setAttributeAssociatedWith(assignOp.getResult(), + tensor_ext::TensorExtDialect::kLayoutAttrName, + layout); + Value toReplace = assignOp.getResult(); + builder.replaceUsesWithIf(val, toReplace, [&](OpOperand& other) { + return other.getOwner() == op; + }); + assignedLayouts.insert({toReplace, layout}); + debugAssignLayout(toReplace, layout); + return toReplace; + }; + 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}); + assignOrUpdateLayout(weight, diagRelation); } RankedTensorType outputType = cast(result.getType()); - IntegerRelation outputLayoutResult = - getBicyclicLayoutRelation(outputType, minSlotCount); - LayoutAttr outputLayoutAttr = - LayoutAttr::getFromIntegerRelation(ctx, outputLayoutResult); + LayoutAttr outputLayoutAttr = LayoutAttr::getFromIntegerRelation( + ctx, getBicyclicLayoutRelation(outputType, minSlotCount)); auto kernelInfoAttr = cloneKernelInfoWithResultShape(secretOperand, outputType.getShape()); assignedLayouts.insert({result, outputLayoutAttr}); - setResultLayoutAttr(op, kernelInfoAttr); debugAssignLayout(result, outputLayoutAttr); + // The kernel adds the accumulator (init) directly to the output via SIMD + // addition. Align the init layout with the result layout at compile time + // so cleartext constants (zeros/bias) re-pack freely without conversions. + Value init = op.getOutputs()[0]; + LayoutAttr initLayout = getComposedLayoutAttr(init); + if (initLayout != outputLayoutAttr) { + assignOrUpdateLayout(init, outputLayoutAttr.getIntegerRelation()); + } + + setResultLayoutAttr(op, kernelInfoAttr); + auto kernelAttr = secret::KernelAttr::get( ctx, KernelName::MatmulBicyclicDiagonal, /*force=*/false); op->setAttr(secret::SecretDialect::kKernelAttrName, kernelAttr); @@ -1533,8 +1801,66 @@ LogicalResult LayoutPropagation::visitOperation(ReduceOp op) { // enforce row-major layout RankedTensorType thisType = cast(tensor.getType()); - if (!isRelationRowMajor(thisType, minSlotCount, - thisLayout.getIntegerRelation())) { + LayoutAttr resultLayout; + + // In-layout strided reduce (LKAA25 App. F): a bicyclic/tricyclic layout + // reduced along the axis mapping to one cyclic factor sums with a + // strided rotate-and-reduce - NO conversion - and the result comes + // out replicated along the reduced axis, which is the layout the + // following divide/normalize broadcast wants (so that conversion + // disappears too). + bool fitsOneCt = thisType.getNumElements() <= minSlotCount; + bool isBicyclic = fitsOneCt && thisType.getRank() == 2 && + op.getDimensions() == ArrayRef{1} && + isRelationBicyclic(thisType, minSlotCount, + thisLayout.getIntegerRelation()); + bool isTricyclic = fitsOneCt && thisType.getRank() == 3 && + op.getDimensions() == ArrayRef{2} && + isRelationTricyclic(thisType, minSlotCount, + thisLayout.getIntegerRelation()); + // In-layout strided reduce is currently restricted to floating-point + // element types; integer reductions require zero-padding semantics. + bool isCyclic = (isBicyclic || isTricyclic) && op.getInputs().size() == 1 && + isa(thisType.getElementType()); + + if (!isCyclic && op.getInputs().size() == 1 && + isa(thisType.getElementType()) && fitsOneCt) { + std::optional target; + if (thisType.getRank() == 2 && + op.getDimensions() == ArrayRef{1} && + thisType.getDimSize(0) > 1 && thisType.getDimSize(1) > 1 && + std::gcd(thisType.getDimSize(0), thisType.getDimSize(1)) == 1) { + target = getBicyclicLayoutRelation(thisType, minSlotCount); + } else if (thisType.getRank() == 3 && + op.getDimensions() == ArrayRef{2} && + thisType.getDimSize(0) > 1 && thisType.getDimSize(1) > 1 && + thisType.getDimSize(2) > 1 && + std::gcd(thisType.getDimSize(0), thisType.getDimSize(1)) == + 1 && + std::gcd(thisType.getDimSize(0), thisType.getDimSize(2)) == + 1 && + std::gcd(thisType.getDimSize(1), thisType.getDimSize(2)) == + 1) { + target = getTricyclicLayoutRelation(thisType, minSlotCount); + } + if (target) { + LayoutAttr canonical = LayoutAttr::getFromIntegerRelation(ctx, *target); + if (canonical != thisLayout) { + auto [toReplace, newLayoutAttr] = + convertToLayout(ctx, builder, op, tensor, thisLayout, *target); + debugAssignLayout(toReplace, newLayoutAttr); + assignedLayouts.insert({toReplace, newLayoutAttr}); + thisLayout = newLayoutAttr; + isCyclic = true; + } else { + isCyclic = true; + } + } + } + + // enforce row-major layout if not cyclic + if (!isCyclic && !isRelationRowMajor(thisType, minSlotCount, + thisLayout.getIntegerRelation())) { LLVM_DEBUG(llvm::dbgs() << "ReduceOp tensor is not row major"); auto [toReplace, newLayoutAttr] = convertToLayout(ctx, builder, op, tensor, thisLayout, @@ -1544,12 +1870,43 @@ LogicalResult LayoutPropagation::visitOperation(ReduceOp op) { thisLayout = newLayoutAttr; } - // drop dimension on output - LayoutAttr resultLayout = - convertLayoutForReduce(thisLayout, op.getDimensions()); + // Both the cyclic and row-major cases: drop the reduced dims from + // the input relation. For the canonical congruence this leaves + // exactly the kept-moduli replication relation (see + // TestBicyclicReduceProjection). + resultLayout = convertLayoutForReduce(thisLayout, op.getDimensions()); assignedLayouts.insert({result, resultLayout}); debugAssignLayout(result, resultLayout); } + + for (const auto& [init, result] : llvm::zip(op.getInits(), op.getResults())) { + LayoutAttr initLayout = getComposedLayoutAttr(init); + LayoutAttr targetLayout = cast(assignedLayouts.at(result)); + if (initLayout != targetLayout) { + if (auto assignOp = init.getDefiningOp(); + assignOp && assignOp->hasOneUse()) { + assignOp.setLayoutAttr(targetLayout); + assignedLayouts[init] = targetLayout; + setAttributeAssociatedWith( + init, tensor_ext::TensorExtDialect::kLayoutAttrName, targetLayout); + debugAssignLayout(init, targetLayout); + } else { + builder.setInsertionPoint(op); + AssignLayoutOp newAssignOp = + AssignLayoutOp::create(builder, op->getLoc(), init, targetLayout); + setAttributeAssociatedWith( + newAssignOp.getResult(), + tensor_ext::TensorExtDialect::kLayoutAttrName, targetLayout); + Value toReplace = newAssignOp.getResult(); + builder.replaceUsesWithIf(init, toReplace, [&](OpOperand& other) { + return other.getOwner() == op; + }); + assignedLayouts.insert({toReplace, targetLayout}); + debugAssignLayout(toReplace, targetLayout); + } + } + } + setResultLayoutAttr(op); return success(); } @@ -1776,6 +2133,15 @@ LogicalResult LayoutPropagation::visitOperation(tensor::ExtractSliceOp op) { return success(); } +template +static CompatibilityResult requireInputLayout( + OpT op, const DenseMap& assignedLayouts) { + if (!assignedLayouts.contains(op.getInput())) { + return {false, op->emitError("input operand has no assigned layout")}; + } + return {true, std::nullopt}; +} + CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts( Operation* op) { LLVM_DEBUG(llvm::dbgs() << "Checking for compatible argument layouts for: " @@ -1783,11 +2149,15 @@ CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts( return TypeSwitch(op) // Trivially true ops .Case( + affine::AffineYieldOp, ConvertLayoutOp>( [&](auto op) { return CompatibilityResult{true, std::nullopt}; }) + // Ops requiring input layout only + .Case( + [&](auto op) { return requireInputLayout(op, assignedLayouts); }) // Ops with special rules - .Case( + .Case( [&](auto op) { return hasCompatibleArgumentLayouts(op); }) // By default, assume operands must all have the same layout. .Default([&](Operation* op) { @@ -1841,27 +2211,11 @@ CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(DotOp op) { 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())) { + for (Value input : op.getInputs()) { 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}; } @@ -1918,6 +2272,21 @@ CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts( return {true, std::nullopt}; } +CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts( + BatchMatmulOp 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. @@ -2030,8 +2399,8 @@ void LayoutPropagation::rectifyIncompatibleOperandLayouts(Operation* op) { TypeSwitch(op) // These ops shouldn't rectify operand layouts - .Case( - [&](auto op) { return; }) + .Case([&](auto op) { return; }) // Ops with special rules .Case( [&](auto op) { return rectifyIncompatibleOperandLayouts(op); }) @@ -2045,6 +2414,37 @@ void LayoutPropagation::rectifyIncompatibleOperandLayouts(Operation* op) { }); LayoutAttr targetLayout = getComposedLayoutAttr(*it); + // A PARTIAL target relation (e.g. a pad-composed layout, whose + // domain excludes the pad indices) cannot serve as an elementwise + // rectification target: constants whose support includes the pad + // region (padding corrections, pins) are unrepresentable under + // it, and downstream in-layout kernels demand total congruences + // anyway. Upgrade such targets to the canonical layout for the + // type; converting the partial-layout operand to it zero-fills + // the pad slots (the physical repack the pad design prescribes). + if (auto targetType = dyn_cast((*it).getType()); + targetType && targetType.hasStaticShape()) { + const IntegerRelation& targetRel = targetLayout.getIntegerRelation(); + unsigned domainOffset = + targetRel.getVarKindOffset(presburger::VarKind::Domain); + bool partial = false; + // Check whether the relation restricts any domain variable to an + // upper bound strictly smaller than the tensor dimension size. + for (unsigned d = 0; d < targetRel.getNumDomainVars(); ++d) { + auto ub = targetRel.getConstantBound64(presburger::BoundType::UB, + domainOffset + d); + if (ub && *ub + 1 < targetType.getDimSize(d)) { + partial = true; + break; + } + } + if (partial) { + FailureOr canonical = defaultLayoutForType(targetType); + if (succeeded(canonical) && canonical.value() != targetLayout) + targetLayout = canonical.value(); + } + } + for (auto& opOperand : op->getOpOperands()) { if (!assignedLayouts.contains(opOperand.get())) continue; LayoutAttr sourceLayout = getComposedLayoutAttr(opOperand.get()); @@ -2117,28 +2517,7 @@ void LayoutPropagation::rectifyIncompatibleOperandLayouts(DotOp op) { } } -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(ReduceOp op) {} void LayoutPropagation::rectifyIncompatibleOperandLayouts( tensor::InsertSliceOp op) { diff --git a/lib/Transforms/LayoutPropagation/Utils.cpp b/lib/Transforms/LayoutPropagation/Utils.cpp index 695a0e5904..3b6adfb0dc 100644 --- a/lib/Transforms/LayoutPropagation/Utils.cpp +++ b/lib/Transforms/LayoutPropagation/Utils.cpp @@ -7,6 +7,7 @@ #include #include +#include "lib/Utils/Layout/Utils.h" #include "llvm/include/llvm/ADT/ArrayRef.h" // from @llvm-project #include "llvm/include/llvm/ADT/STLExtras.h" // from @llvm-project #include "llvm/include/llvm/ADT/SmallVector.h" // from @llvm-project @@ -14,6 +15,7 @@ #include "mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h" // from @llvm-project #include "mlir/include/mlir/IR/Attributes.h" // from @llvm-project #include "mlir/include/mlir/IR/BuiltinAttributes.h" // from @llvm-project +#include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/include/mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project @@ -136,6 +138,70 @@ SmallVector shiftByRemoved(ArrayRef dims, LayoutAttr convertLayoutForReduce(LayoutAttr inputLayout, ArrayRef dimsToReduce) { + const presburger::IntegerRelation& rel = inputLayout.getIntegerRelation(); + unsigned domainOffset = rel.getVarKindOffset(presburger::VarKind::Domain); + unsigned rangeOffset = rel.getVarKindOffset(presburger::VarKind::Range); + unsigned numDomainVars = rel.getNumDomainVars(); + MLIRContext* context = inputLayout.getContext(); + + // If the input layout has a tricyclic CRT layout (rank 3) and is reduced + // along dimension 2: + // Directly constructing the bicyclic relation for the remaining dimensions + // avoids Fourier-Motzkin projection of the merged CRT modulo constraint, + // which loses divisibility information and degenerates into an unconstrained + // relation. + auto slotUb = + rel.getConstantBound64(presburger::BoundType::UB, rangeOffset + 1); + if (slotUb.has_value() && numDomainVars == 3 && + dimsToReduce == ArrayRef{2}) { + int64_t numSlots = slotUb.value() + 1; + SmallVector shape; + for (unsigned i = 0; i < 3; ++i) { + auto ub = + rel.getConstantBound64(presburger::BoundType::UB, domainOffset + i); + if (ub.has_value()) { + shape.push_back(ub.value() + 1); + } + } + if (shape.size() == 3) { + RankedTensorType tensorType = + RankedTensorType::get(shape, Float32Type::get(context)); + if (isRelationTricyclic(tensorType, numSlots, rel)) { + RankedTensorType reducedType = RankedTensorType::get( + {tensorType.getDimSize(0), tensorType.getDimSize(1)}, + Float32Type::get(context)); + return LayoutAttr::getFromIntegerRelation( + context, getBicyclicLayoutRelation(reducedType, numSlots)); + } + } + } + + // If the input layout has a bicyclic CRT layout (rank 2) and is reduced + // along dimension 1: + if (slotUb.has_value() && numDomainVars == 2 && + dimsToReduce == ArrayRef{1}) { + int64_t numSlots = slotUb.value() + 1; + SmallVector shape; + for (unsigned i = 0; i < 2; ++i) { + auto ub = + rel.getConstantBound64(presburger::BoundType::UB, domainOffset + i); + if (ub.has_value()) { + shape.push_back(ub.value() + 1); + } + } + if (shape.size() == 2) { + RankedTensorType matrixType = + RankedTensorType::get(shape, Float32Type::get(context)); + if (isRelationBicyclic(matrixType, numSlots, rel)) { + presburger::IntegerRelation bicyclicRel = + getBicyclicLayoutRelation(matrixType, numSlots); + bicyclicRel.projectOut(1, 1); + return LayoutAttr::getFromIntegerRelation(context, + std::move(bicyclicRel)); + } + } + } + std::unique_ptr clonedRelation = inputLayout.getIntegerRelation().clone(); @@ -148,7 +214,6 @@ LayoutAttr convertLayoutForReduce(LayoutAttr inputLayout, clonedRelation->projectOut(dimIndex, 1); } - MLIRContext* context = inputLayout.getContext(); return LayoutAttr::getFromIntegerRelation(context, std::move(*clonedRelation)); } diff --git a/lib/Transforms/LayoutPropagation/UtilsTest.cpp b/lib/Transforms/LayoutPropagation/UtilsTest.cpp index 6b797d3ba6..2378725168 100644 --- a/lib/Transforms/LayoutPropagation/UtilsTest.cpp +++ b/lib/Transforms/LayoutPropagation/UtilsTest.cpp @@ -153,6 +153,30 @@ TEST(UtilsTest, TestReduceLayoutManyReductions) { EXPECT_TRUE(isRelationEqual(reducedRelation, expectedRelation)); } +TEST(UtilsTest, TestReduceTricyclicLayout) { + MLIRContext context; + context.loadDialect(); + + // Tricyclic layout for 2x33x37 tensor in 32768 slots. + RankedTensorType tensorType = + RankedTensorType::get({2, 33, 37}, Float32Type::get(&context)); + presburger::IntegerRelation relation = + getTricyclicLayoutRelation(tensorType, 32768); + LayoutAttr layout = LayoutAttr::getFromIntegerRelation(&context, relation); + + SmallVector dimsToReduce = {2}; + LayoutAttr reducedLayout = convertLayoutForReduce(layout, dimsToReduce); + presburger::IntegerRelation reducedRelation = + reducedLayout.getIntegerRelation(); + + EXPECT_EQ(reducedRelation.getNumDomainVars(), 2); + EXPECT_EQ(reducedRelation.getNumRangeVars(), 2); + + RankedTensorType reducedType = + RankedTensorType::get({2, 33}, Float32Type::get(&context)); + EXPECT_TRUE(isRelationBicyclic(reducedType, 32768, reducedRelation)); +} + TEST(UtilsTest, TestFoldConvSpatialPadding) { MLIRContext context; Type elementType = IndexType::get(&context); diff --git a/lib/Utils/Layout/Utils.cpp b/lib/Utils/Layout/Utils.cpp index b05d686640..e27bf3bc52 100644 --- a/lib/Utils/Layout/Utils.cpp +++ b/lib/Utils/Layout/Utils.cpp @@ -545,6 +545,61 @@ presburger::IntegerRelation getPerRowLayoutRelation(RankedTensorType matrixType, return result; } +presburger::IntegerRelation getTricyclicDiagonalRelation( + RankedTensorType weightType, int64_t contractionDim, int64_t ctStride, + int64_t paddedFreeDim, int64_t numSlots) { + int64_t rank = weightType.getRank(); + assert(rank == 3 && "tricyclic diagonal relation requires a rank-3 weight"); + assert( + (contractionDim == 1 || contractionDim == 2) && + "contractionDim must be 1 (ct-pt, RHS weight) or 2 (pt-ct, LHS weight)"); + int64_t h = weightType.getDimSize(0); + int64_t freeDim = (contractionDim == 1) ? 2 : 1; + int64_t n = weightType.getDimSize(contractionDim); + int64_t p = weightType.getDimSize(freeDim); + assert(p <= paddedFreeDim && + "paddedFreeDim must cover the weight's true size"); + + IntegerRelation result(PresburgerSpace::getRelationSpace( + rank, /*numRange=*/2, /*numSymbol=*/0, /*numLocals=*/0)); + + int domainOffset = result.getVarKindOffset(VarKind::Domain); + int rangeOffset = result.getVarKindOffset(VarKind::Range); + int contractionVarIndex = domainOffset + contractionDim; + int freeVarIndex = domainOffset + freeDim; + int diagVarIndex = rangeOffset; + int slotVarIndex = rangeOffset + 1; + + addBounds(result, domainOffset, 0, h - 1); + addBounds(result, contractionVarIndex, 0, n - 1); + addBounds(result, freeVarIndex, 0, p - 1); + addBounds(result, diagVarIndex, 0, n - 1); + addBounds(result, slotVarIndex, 0, numSlots - 1); + + // contractionIdx = (slot + diag * h * ctStride) mod n + SmallVector nCoeffs(result.getNumCols(), 0); + nCoeffs[slotVarIndex] = 1; + nCoeffs[diagVarIndex] = h * ctStride; + auto nMod = addModConstraint(result, nCoeffs, n); + addConstraint(result, {{nMod, 1}, {contractionVarIndex, -1}}, + /*equality=*/true); + + // freeIdx = slot mod paddedFreeDim (slots whose residue >= p are unmapped and + // zero-filled by AssignLayout) + SmallVector pCoeffs(result.getNumCols(), 0); + pCoeffs[slotVarIndex] = 1; + auto pMod = addModConstraint(result, pCoeffs, paddedFreeDim); + addConstraint(result, {{pMod, 1}, {freeVarIndex, -1}}, /*equality=*/true); + + // hIdx = slot mod h + SmallVector hCoeffs(result.getNumCols(), 0); + hCoeffs[slotVarIndex] = 1; + auto hMod = addModConstraint(result, hCoeffs, h); + addConstraint(result, {{hMod, 1}, {domainOffset, -1}}, /*equality=*/true); + + return result; +} + bool isRelationSquatDiagonal(RankedTensorType matrixType, int64_t minSlotCount, const presburger::IntegerRelation& relation) { IntegerRelation diagonalRelation = @@ -1305,6 +1360,7 @@ static std::optional tryIslEqual( const presburger::IntegerRelation& relation1, const presburger::IntegerRelation& relation2) { isl_ctx* ctx = isl_ctx_alloc(); + isl_ctx_set_max_operations(ctx, 100000); isl_map* map1 = isl_map_from_basic_map(convertRelationToBasicMap(relation1, ctx)); isl_map* map2 = @@ -1321,6 +1377,36 @@ static std::optional tryIslEqual( return equal == isl_bool_true; } +presburger::IntegerRelation getTransposedRelation( + const presburger::IntegerRelation& relation, + ArrayRef permutation) { + assert(permutation.size() == relation.getNumDomainVars() && + "permutation size must match relation domain rank"); + presburger::IntegerRelation result = relation; + unsigned domainOffset = result.getVarKindOffset(presburger::VarKind::Domain); + unsigned numDomain = permutation.size(); + + SmallVector posToElem(numDomain); + SmallVector elemToPos(numDomain); + for (unsigned i = 0; i < numDomain; ++i) { + posToElem[i] = i; + elemToPos[i] = i; + } + for (unsigned targetIdx = 0; targetIdx < numDomain; ++targetIdx) { + int64_t targetElem = permutation[targetIdx]; + int64_t curPos = elemToPos[targetElem]; + if (curPos != static_cast(targetIdx)) { + result.swapVar(domainOffset + targetIdx, domainOffset + curPos); + int64_t elemAtTarget = posToElem[targetIdx]; + posToElem[targetIdx] = targetElem; + posToElem[curPos] = elemAtTarget; + elemToPos[targetElem] = targetIdx; + elemToPos[elemAtTarget] = curPos; + } + } + return result; +} + bool isRelationEqual(const presburger::IntegerRelation& relation1, const presburger::IntegerRelation& relation2) { // Structural equality, in a few nanoseconds. @@ -1336,6 +1422,69 @@ bool isRelationEqual(const presburger::IntegerRelation& relation1, return islResult.value_or(false); } +bool isRelationSubset(const presburger::IntegerRelation& relation1, + const presburger::IntegerRelation& relation2) { + if (relation1.getNumDomainVars() != relation2.getNumDomainVars() || + relation1.getNumRangeVars() != relation2.getNumRangeVars()) { + return false; + } + if (relation1.isObviouslyEqual(relation2)) return true; + + // Local/div variables (e.g. modular CRT congruence relations): ISL quantifier + // elimination on existential basic maps is incomplete or fails. For bounded + // tensor domains, evaluate containment over discrete point pairs. + if (relation1.getNumLocalVars() > 0 || relation2.getNumLocalVars() > 0) { + PointPairCollector p1(relation1.getNumDomainVars(), + relation1.getNumRangeVars()); + enumeratePoints(relation1, p1); + std::sort(p1.points.begin(), p1.points.end()); + p1.points.erase(std::unique(p1.points.begin(), p1.points.end()), + p1.points.end()); + + PointPairCollector p2(relation2.getNumDomainVars(), + relation2.getNumRangeVars()); + enumeratePoints(relation2, p2); + std::sort(p2.points.begin(), p2.points.end()); + p2.points.erase(std::unique(p2.points.begin(), p2.points.end()), + p2.points.end()); + + if (p1.points.size() > p2.points.size()) return false; + return std::includes(p2.points.begin(), p2.points.end(), p1.points.begin(), + p1.points.end()); + } + + isl_ctx* ctx = isl_ctx_alloc(); + isl_ctx_set_max_operations(ctx, 100000); + isl_basic_map* bmap1 = convertRelationToBasicMap(relation1, ctx); + isl_basic_map* bmap2 = convertRelationToBasicMap(relation2, ctx); + isl_map* map1 = bmap1 ? isl_map_from_basic_map(bmap1) : nullptr; + isl_map* map2 = bmap2 ? isl_map_from_basic_map(bmap2) : nullptr; + + isl_bool isSubset = + (map1 && map2) ? isl_map_is_subset(map1, map2) : isl_bool_error; + + isl_map_free(map1); + isl_map_free(map2); + isl_ctx_free(ctx); + + return isSubset == isl_bool_true; +} + +bool isRelationInjective(const presburger::IntegerRelation& relation) { + isl_ctx* ctx = isl_ctx_alloc(); + isl_ctx_set_max_operations(ctx, 100000); + isl_basic_map* bmap = convertRelationToBasicMap(relation, ctx); + if (!bmap) { + isl_ctx_free(ctx); + return false; + } + isl_map* map = isl_map_from_basic_map(bmap); + isl_bool injective = isl_map_is_injective(map); + isl_map_free(map); + isl_ctx_free(ctx); + return injective == isl_bool_true; +} + bool isDenseLayout(const presburger::IntegerRelation& relation, RankedTensorType type) { isl_ctx* ctx = isl_ctx_alloc(); diff --git a/lib/Utils/Layout/Utils.h b/lib/Utils/Layout/Utils.h index 0fdd1ea102..f6a9376761 100644 --- a/lib/Utils/Layout/Utils.h +++ b/lib/Utils/Layout/Utils.h @@ -9,13 +9,11 @@ #include "mlir/include/mlir/Analysis/Presburger/IntegerRelation.h" // from @llvm-project #include "mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h" // from @llvm-project #include "mlir/include/mlir/Dialect/Arith/Utils/Utils.h" // from @llvm-project -#include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project #include "mlir/include/mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/include/mlir/Support/LLVM.h" // from @llvm-project // ISL #include "include/isl/ctx.h" // from @isl -#include "include/isl/map.h" // from @isl namespace mlir { namespace heir { @@ -107,6 +105,15 @@ presburger::IntegerRelation getPeriodicReplicationRelation( presburger::IntegerRelation getPerRowLayoutRelation(RankedTensorType matrixType, int64_t minSlotCount); +// Returns the diagonal packing relation for the rank-3 plaintext operand +// of the batch ciphertext-plaintext matmul: +// BatchDiag'(B, c)_k = +// B[k mod h][(k + c*h*ctStride) mod n][k mod paddedFreeDim] +// contractionDim specifies the contraction axis (1 for ct-pt, 2 for pt-ct). +presburger::IntegerRelation getTricyclicDiagonalRelation( + RankedTensorType weightType, int64_t contractionDim, int64_t ctStride, + int64_t paddedFreeDim, int64_t numSlots); + // Returns true if the given relation is a squat diagonal layout for the given // matrix type and ciphertext semantic shape. bool isRelationSquatDiagonal(RankedTensorType matrixType, int64_t minSlotCount, @@ -268,6 +275,12 @@ FailureOr getSliceExtractionRelation( SmallVector offsets, SmallVector sizes, SmallVector strides); +// Returns the relation corresponding to a transpose by permuting its domain +// variables according to `permutation`: result domain index i corresponds to +// original domain index `permutation[i]`. +presburger::IntegerRelation getTransposedRelation( + const presburger::IntegerRelation& relation, ArrayRef permutation); + // Tests whether two layout relations describe the same set of points. // // This check is one-sided: `true` means the relations are provably equal, but @@ -278,6 +291,16 @@ FailureOr getSliceExtractionRelation( bool isRelationEqual(const presburger::IntegerRelation& relation1, const presburger::IntegerRelation& relation2); +// Tests whether relation1 is a subset of relation2 (i.e. every point in +// relation1 is also in relation2). +bool isRelationSubset(const presburger::IntegerRelation& relation1, + const presburger::IntegerRelation& relation2); + +// Returns true if the relation maps no two domain points to the same +// range point (i.e. an arbitrary value of the domain type is +// representable). Conservative: unknown/failure returns false. +bool isRelationInjective(const presburger::IntegerRelation& relation); + // Returns true if the given relation is surjective onto the given tensor type. // This tests that the range set of the relation covers all points of the given // tensor type. This is used to test if a layout is dense, so that the layout diff --git a/lib/Utils/Layout/UtilsTest.cpp b/lib/Utils/Layout/UtilsTest.cpp index 445d17c699..e794f6aa19 100644 --- a/lib/Utils/Layout/UtilsTest.cpp +++ b/lib/Utils/Layout/UtilsTest.cpp @@ -1055,7 +1055,74 @@ TEST(UtilsTest, TestGetPaddingRelation) { // p = 1 => s = -1, out of bounds EXPECT_FALSE(rel.containsPointNoLocal({1, -1}).has_value()); } +TEST(UtilsTest, TestRelationSubset) { + // `from`: 2x4 box + auto from = getIntegerRelationFromIslStr( + "{ [i0, i1] -> [ct, slot] : ct = 0 and slot = i0 + 4*i1 and " + "0 <= i0 <= 3 and 0 <= i1 <= 1 }") + .value(); + // `to`: 1x4 sub-box (i1 = 0) + auto to = getIntegerRelationFromIslStr( + "{ [i0, i1] -> [ct, slot] : ct = 0 and slot = i0 and 0 <= i0 " + "<= 3 and i1 = 0 }") + .value(); + EXPECT_TRUE(isRelationSubset(to, from)); + EXPECT_FALSE(isRelationSubset(from, to)); + + // `outside`: point not in `from` + auto outside = getIntegerRelationFromIslStr( + "{ [i0, i1] -> [ct, slot] : ct = 0 and slot = i0 + 5 and " + "0 <= i0 <= 3 and i1 = 0 }") + .value(); + EXPECT_FALSE(isRelationSubset(outside, from)); +} + +TEST(UtilsTest, TestRelationInjective) { + MLIRContext context; + RankedTensorType type = + RankedTensorType::get({3, 5}, IndexType::get(&context)); + IntegerRelation bicyclic = getBicyclicLayoutRelation(type, 1024); + EXPECT_TRUE(isRelationInjective(bicyclic)); + + auto nonInjective = getIntegerRelationFromIslStr( + "{ [i0, i1] -> [ct, slot] : ct = 0 and slot = i1 and " + "0 <= i0 <= 1 and 0 <= i1 <= 4 }") + .value(); + EXPECT_FALSE(isRelationInjective(nonInjective)); +} +TEST(UtilsTest, TricyclicCtPtDiagonal2x5x7) { + MLIRContext context; + int64_t numSlots = 105; + int64_t ctStride = 3; + int64_t paddedFreeDim = 7; + int64_t contractionDim = 1; + RankedTensorType weightType = + RankedTensorType::get({2, 5, 7}, IndexType::get(&context)); + IntegerRelation relation = getTricyclicDiagonalRelation( + weightType, contractionDim, ctStride, paddedFreeDim, numSlots); + + EXPECT_TRUE(relation.containsPointNoLocal({0, 0, 0, 0, 0}).has_value()); + EXPECT_TRUE(relation.containsPointNoLocal({1, 1, 1, 0, 1}).has_value()); + EXPECT_FALSE(relation.containsPointNoLocal({0, 0, 0, 0, 1}).has_value()); +} + +TEST(UtilsTest, TestBicyclicReduceProjection) { + MLIRContext context; + RankedTensorType type = + RankedTensorType::get({33, 65}, Float32Type::get(&context)); + IntegerRelation rel = getBicyclicLayoutRelation(type, 8192); + EXPECT_TRUE(isRelationBicyclic(type, 8192, rel)); + + auto reducedExpected = + getIntegerRelationFromIslStr( + "{ [i0] -> [ct, slot] : ct = 0 and (-i0 + slot) mod 33 = 0 and 0 " + "<= i0 <= 32 and 0 <= slot <= 8191 }") + .value(); + IntegerRelation projected = rel; + projected.projectOut(1, 1); + EXPECT_TRUE(isRelationEqual(projected, reducedExpected)); +} } // namespace } // namespace heir } // namespace mlir diff --git a/lib/Utils/RotationUtils.h b/lib/Utils/RotationUtils.h index 96079cbbef..a7a94835f7 100644 --- a/lib/Utils/RotationUtils.h +++ b/lib/Utils/RotationUtils.h @@ -64,8 +64,23 @@ inline llvm::DenseSet rotateAndReduceRotationIndices( llvm::DenseSet result; if (!hasPlaintexts) { // Matches implementRotateAndReduceAccumulation - for (int64_t shiftSize = steps / 2; shiftSize > 0; shiftSize /= 2) { - result.insert(shiftSize * period); + if ((steps & (steps - 1)) == 0) { + for (int64_t shiftSize = steps / 2; shiftSize > 0; shiftSize /= 2) { + result.insert(shiftSize * period); + } + } else { + int64_t spanLen = 1, offset = 0, remaining = steps; + while (remaining > 0) { + if (remaining & 1) { + if (offset != 0) result.insert(offset * period); + offset += spanLen; + } + remaining >>= 1; + if (remaining > 0) { + result.insert(spanLen * period); + spanLen <<= 1; + } + } } return result; } diff --git a/tests/Dialect/TensorExt/Transforms/implement_shift_network.mlir b/tests/Dialect/TensorExt/Transforms/implement_shift_network.mlir index 73d4baa3e8..6346d815a0 100644 --- a/tests/Dialect/TensorExt/Transforms/implement_shift_network.mlir +++ b/tests/Dialect/TensorExt/Transforms/implement_shift_network.mlir @@ -15,31 +15,16 @@ func.func @test_no_conflicts(%0: tensor<1x64xi32>) -> tensor<1x64xi32> { return %1 : tensor<1x64xi32> } -// This test has a larger set of rotations because the Vos-Vos-Erkin method -// forces each rotation to be decomposed into power-of-two rotations, even if -// it could be done in a single rotation. In this case it's a (left-)rotation -// by 63 which is equivalent to a single (right)-rotation by -1, which requires -// all power-of-two components to be used. -// -// TODO(#2263): this test should only produce one rotation by -1 +// When the permutation is a single rotation, direct depth-1 rotation is +// selected, emitting a single rotation by 63 (or -1). // // CHECK: func.func @test_no_conflicts2 // CHECK-SAME: (%[[ARG0:.*]]: tensor<1x64xi32>) -> tensor<1x64xi32> -// CHECK: %[[SLICE:.*]] = tensor.extract_slice %[[ARG0]][0, 0] [1, 64] [1, 1] : tensor<1x64xi32> to tensor<1x64xi32> -// CHECK: %[[C1:.*]] = arith.constant 1 : index -// CHECK: %[[ROT0:.*]] = tensor_ext.rotate %[[SLICE]], %[[C1]] : tensor<1x64xi32>, index -// CHECK: %[[C2:.*]] = arith.constant 2 : index -// CHECK: %[[ROT1:.*]] = tensor_ext.rotate %[[ROT0]], %[[C2]] : tensor<1x64xi32>, index -// CHECK: %[[C4:.*]] = arith.constant 4 : index -// CHECK: %[[ROT2:.*]] = tensor_ext.rotate %[[ROT1]], %[[C4]] : tensor<1x64xi32>, index -// CHECK: %[[C8:.*]] = arith.constant 8 : index -// CHECK: %[[ROT3:.*]] = tensor_ext.rotate %[[ROT2]], %[[C8]] : tensor<1x64xi32>, index -// CHECK: %[[C16:.*]] = arith.constant 16 : index -// CHECK: %[[ROT4:.*]] = tensor_ext.rotate %[[ROT3]], %[[C16]] : tensor<1x64xi32>, index -// CHECK: %[[C32:.*]] = arith.constant 32 : index -// CHECK: %[[ROT5:.*]] = tensor_ext.rotate %[[ROT4]], %[[C32]] : tensor<1x64xi32>, index -// CHECK: %[[INSERT:.*]] = tensor.insert_slice %[[ROT5]] into %[[ARG0]][0, 0] [1, 64] [1, 1] : tensor<1x64xi32> into tensor<1x64xi32> -// CHECK: return %[[INSERT]] : tensor<1x64xi32> +// CHECK: %[[SLICE2:.*]] = tensor.extract_slice %[[ARG0]][0, 0] [1, 64] [1, 1] : tensor<1x64xi32> to tensor<1x64xi32> +// CHECK: %[[C63:.*]] = arith.constant 63 : index +// CHECK: %[[ROT2:.*]] = tensor_ext.rotate %[[SLICE2]], %[[C63]] : tensor<1x64xi32>, index +// CHECK: %[[INSERT2:.*]] = tensor.insert_slice %[[ROT2]] into %[[ARG0]][0, 0] [1, 64] [1, 1] : tensor<1x64xi32> into tensor<1x64xi32> +// CHECK: return %[[INSERT2]] : tensor<1x64xi32> #map2 = #tensor_ext.layout<"{ [ct1, slot1] -> [ct2, slot2] : ct1 = 0 and ct2 = 0 and ((slot1 + 1) - slot2) mod 64 = 0 and slot1 >= 0 and 63 >= slot1 and slot2 >= 0 and 63 >= slot2 }"> func.func @test_no_conflicts2(%0: tensor<1x64xi32>) -> tensor<1x64xi32> { %1 = tensor_ext.remap %0 {permutation = #map2} : tensor<1x64xi32> @@ -80,7 +65,7 @@ func.func @multi_ciphertext_swap_cts(%0: tensor<4x64xi32>) -> tensor<4x64xi32> { // properly with multi-ciphertext inputs. // // CHECK: func.func @multi_ciphertext_complex -// CHECK-COUNT-16: tensor_ext.rotate +// CHECK-COUNT-4: tensor_ext.rotate #map5 = #tensor_ext.layout<"{ [ct1, slot1] -> [ct2, slot2] : (ct1 - ct2) mod 4 = 3 and (slot1 - slot2) mod 64 = 5 and 0 <= ct1 <= 3 and 0 <= ct2 <= 3 and 0 <= slot1 <= 63 and 0 <= slot2 <= 63 }"> func.func @multi_ciphertext_complex(%0: tensor<4x64xi32>) -> tensor<4x64xi32> { %1 = tensor_ext.remap %0 {permutation = #map5} : tensor<4x64xi32> diff --git a/tests/Transforms/convert_to_ciphertext_semantics/batch_matmul_pt.mlir b/tests/Transforms/convert_to_ciphertext_semantics/batch_matmul_pt.mlir new file mode 100644 index 0000000000..c53bdbd213 --- /dev/null +++ b/tests/Transforms/convert_to_ciphertext_semantics/batch_matmul_pt.mlir @@ -0,0 +1,54 @@ +// RUN: heir-opt --layout-propagation=min-slot-count=1024 --convert-to-ciphertext-semantics=min-slot-count=1024 %s | FileCheck %s + +#bicyclic = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (-10i0 - 6i1 + slot) mod 15 = 0 and 0 <= i0 <= 2 and 0 <= i1 <= 4 and 0 <= slot <= 1023 }"> +#tricyclic_2_3_5 = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : ct = 0 and (-15i0 - 10i1 - 6i2 + slot) mod 30 = 0 and 0 <= i0 <= 1 and 0 <= i1 <= 2 and 0 <= i2 <= 4 and 0 <= slot <= 1023 }"> +#tricyclic_2_5_7 = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : ct = 0 and (-35i0 - 56i1 - 50i2 + slot) mod 70 = 0 and 0 <= i0 <= 1 and 0 <= i1 <= 4 and 0 <= i2 <= 6 and 0 <= slot <= 1023 }"> + +module { + // CHECK: @batch_matmul_broadcast + // CHECK-NOT: linalg.batch_matmul + // CHECK: tensor_ext.rotate + // CHECK: arith.mulf + // CHECK: tensor_ext.remap + func.func @batch_matmul_broadcast(%arg0: !secret.secret> {tensor_ext.layout = #bicyclic}, %arg1: tensor<2x5x7xf32>) -> !secret.secret> { + %cst = arith.constant dense<0.000000e+00> : tensor<2x3x7xf32> + %empty = tensor.empty() : tensor<2x3x5xf32> + %0 = secret.generic(%arg0: !secret.secret> {tensor_ext.layout = #bicyclic}) { + ^body(%input0: tensor<3x5xf32>): + %bcast = linalg.broadcast ins(%input0 : tensor<3x5xf32>) outs(%empty : tensor<2x3x5xf32>) dimensions = [0] + %1 = linalg.batch_matmul ins(%bcast, %arg1 : tensor<2x3x5xf32>, tensor<2x5x7xf32>) outs(%cst : tensor<2x3x7xf32>) -> tensor<2x3x7xf32> + secret.yield %1 : tensor<2x3x7xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } + + // CHECK: @batch_matmul_ctpt + // CHECK-NOT: linalg.batch_matmul + // CHECK: tensor_ext.rotate + // CHECK: arith.mulf + // CHECK: tensor_ext.remap + func.func @batch_matmul_ctpt(%arg0: !secret.secret> {tensor_ext.layout = #tricyclic_2_3_5}, %arg1: tensor<2x5x7xf32>) -> !secret.secret> { + %cst = arith.constant dense<0.000000e+00> : tensor<2x3x7xf32> + %0 = secret.generic(%arg0: !secret.secret> {tensor_ext.layout = #tricyclic_2_3_5}) { + ^body(%input0: tensor<2x3x5xf32>): + %1 = linalg.batch_matmul ins(%input0, %arg1 : tensor<2x3x5xf32>, tensor<2x5x7xf32>) outs(%cst : tensor<2x3x7xf32>) -> tensor<2x3x7xf32> + secret.yield %1 : tensor<2x3x7xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } + + // CHECK: @batch_matmul_ptct + // CHECK-NOT: linalg.batch_matmul + // CHECK: tensor_ext.rotate + // CHECK: arith.mulf + // CHECK: tensor_ext.remap + func.func @batch_matmul_ptct(%arg0: tensor<2x3x5xf32>, %arg1: !secret.secret> {tensor_ext.layout = #tricyclic_2_5_7}) -> !secret.secret> { + %cst = arith.constant dense<0.000000e+00> : tensor<2x3x7xf32> + %0 = secret.generic(%arg1: !secret.secret> {tensor_ext.layout = #tricyclic_2_5_7}) { + ^body(%input0: tensor<2x5x7xf32>): + %1 = linalg.batch_matmul ins(%arg0, %input0 : tensor<2x3x5xf32>, tensor<2x5x7xf32>) outs(%cst : tensor<2x3x7xf32>) -> tensor<2x3x7xf32> + secret.yield %1 : tensor<2x3x7xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } +} diff --git a/tests/Transforms/convert_to_ciphertext_semantics/broadcast.mlir b/tests/Transforms/convert_to_ciphertext_semantics/broadcast.mlir new file mode 100644 index 0000000000..31dcad6271 --- /dev/null +++ b/tests/Transforms/convert_to_ciphertext_semantics/broadcast.mlir @@ -0,0 +1,17 @@ +// RUN: heir-opt --layout-propagation=min-slot-count=8192 --convert-to-ciphertext-semantics=min-slot-count=8192 %s | FileCheck %s + +#bicyclic = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (-66i0 + 65i1 + slot) mod 2145 = 0 and 0 <= i0 <= 32 and 0 <= i1 <= 64 and 0 <= slot <= 8191 }"> + +// CHECK: func.func @broadcast_2d_to_3d +func.func @broadcast_2d_to_3d(%arg0: !secret.secret> {tensor_ext.layout = #bicyclic}) -> !secret.secret> { + %init = arith.constant dense<0.000000e+00> : tensor<2x33x65xf32> + // CHECK: secret.generic + // CHECK-NEXT: ^body(%[[IN:.*]]: tensor<1x8192xf32>): + // CHECK: secret.yield %[[IN]] : tensor<1x8192xf32> + %0 = secret.generic(%arg0: !secret.secret>) { + ^body(%input0: tensor<33x65xf32>): + %broadcasted = linalg.broadcast ins(%input0 : tensor<33x65xf32>) outs(%init : tensor<2x33x65xf32>) dimensions = [0] + secret.yield %broadcasted : tensor<2x33x65xf32> + } -> !secret.secret> + return %0 : !secret.secret> +} diff --git a/tests/Transforms/convert_to_ciphertext_semantics/linalg_reduce.mlir b/tests/Transforms/convert_to_ciphertext_semantics/linalg_reduce.mlir index 11954d8299..613f224d7a 100644 --- a/tests/Transforms/convert_to_ciphertext_semantics/linalg_reduce.mlir +++ b/tests/Transforms/convert_to_ciphertext_semantics/linalg_reduce.mlir @@ -1,31 +1,59 @@ -// RUN: heir-opt %s --split-input-file --convert-to-ciphertext-semantics=min-slot-count=1024 | FileCheck %s -// Test that a 8 length vector gets reduced. -// CHECK: func.func @main -// CHECK-NOT: linalg.reduce -// CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index -// CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index -// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index -// CHECK-DAG: %[[ASSIGN:.*]] = arith.constant dense<0{{.*}}> : tensor<1x1024xf32> -// CHECK: tensor_ext.rotate %{{.*}}, %[[c4]] -// CHECK: tensor_ext.rotate %{{.*}}, %[[c2]] -// CHECK: tensor_ext.rotate %{{.*}}, %[[c1]] -// CHECK: arith.addf %{{.*}}, %[[ASSIGN]] -#layout = #tensor_ext.layout<"{ [] -> [ct, slot] : ct = 0 and (slot) mod 8 = 0 and 0 <= slot <= 1023 }"> -#layout1 = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and (-i0 + slot) mod 8 = 0 and 0 <= i0 <= 7 and 0 <= slot <= 1023 }"> +// RUN: heir-opt --layout-propagation=min-slot-count=8192 --convert-to-ciphertext-semantics=min-slot-count=8192 %s | FileCheck %s + +#bicyclic = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (65i0 - 66i1 + slot) mod 2145 = 0 and 0 <= i0 <= 32 and 0 <= i1 <= 64 and 0 <= slot <= 8191 }"> + module { - func.func @main(%arg0: !secret.secret> {tensor_ext.layout = #layout1}, %arg1: !secret.secret> {tensor_ext.layout = #layout1}) -> (!secret.secret> {tensor_ext.layout = #layout}) { + // Test that an 8-length vector gets reduced. + // CHECK: func.func @main + // CHECK-NOT: linalg.reduce + // CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index + // CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index + // CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index + // CHECK-DAG: %[[ASSIGN:.*]] = arith.constant dense<0{{.*}}> : tensor<1x8192xf32> + // CHECK: tensor_ext.rotate %{{.*}}, %[[c4]] + // CHECK: tensor_ext.rotate %{{.*}}, %[[c2]] + // CHECK: tensor_ext.rotate %{{.*}}, %[[c1]] + // CHECK: arith.addf %{{.*}}, %[[ASSIGN]] + func.func @main(%arg0: !secret.secret>, %arg1: !secret.secret>) -> !secret.secret> { %cst = arith.constant dense<0.000000e+00> : tensor - %0 = secret.generic(%arg0: !secret.secret> {tensor_ext.layout = #layout1}, %arg1: !secret.secret> {tensor_ext.layout = #layout1}) { -^body(%input0: tensor<8xf32>, %input1: tensor<8xf32>): - %1 = arith.mulf %input0, %input1 {tensor_ext.layout = #layout1} : tensor<8xf32> - %2 = tensor_ext.assign_layout %cst {layout = #layout, tensor_ext.layout = #layout} : tensor - %reduced = linalg.reduce ins(%1 : tensor<8xf32>) outs(%2 : tensor) dimensions = [0] {tensor_ext.layout = #layout} + %0 = secret.generic(%arg0: !secret.secret>, %arg1: !secret.secret>) { + ^body(%input0: tensor<8xf32>, %input1: tensor<8xf32>): + %1 = arith.mulf %input0, %input1 : tensor<8xf32> + %reduced = linalg.reduce ins(%1 : tensor<8xf32>) outs(%cst : tensor) dimensions = [0] (%in: f32, %init: f32) { - %3 = arith.addf %in, %init : f32 - linalg.yield %3 : f32 + %2 = arith.addf %in, %init : f32 + linalg.yield %2 : f32 } secret.yield %reduced : tensor - } -> (!secret.secret> {tensor_ext.layout = #layout}) + } -> !secret.secret> return %0 : !secret.secret> } + + // CHECK: func.func @reduce_bicyclic + // CHECK-DAG: %[[C33:.*]] = arith.constant 33 : index + // CHECK-DAG: %[[C66:.*]] = arith.constant 66 : index + // CHECK-DAG: %[[C132:.*]] = arith.constant 132 : index + // CHECK-DAG: %[[C264:.*]] = arith.constant 264 : index + // CHECK-DAG: %[[C528:.*]] = arith.constant 528 : index + // CHECK-DAG: %[[C1056:.*]] = arith.constant 1056 : index + // CHECK: tensor_ext.rotate {{.*}}, %[[C33]] + // CHECK: tensor_ext.rotate {{.*}}, %[[C66]] + // CHECK: tensor_ext.rotate {{.*}}, %[[C132]] + // CHECK: tensor_ext.rotate {{.*}}, %[[C264]] + // CHECK: tensor_ext.rotate {{.*}}, %[[C528]] + // CHECK: tensor_ext.rotate {{.*}}, %[[C1056]] + // CHECK: tensor_ext.remap + func.func @reduce_bicyclic(%arg0: !secret.secret> {tensor_ext.layout = #bicyclic}) -> !secret.secret> { + %cst = arith.constant dense<0.000000e+00> : tensor<33xf32> + %0 = secret.generic(%arg0: !secret.secret>) { + ^body(%input0: tensor<33x65xf32>): + %reduced = linalg.reduce ins(%input0 : tensor<33x65xf32>) outs(%cst : tensor<33xf32>) dimensions = [1] + (%in: f32, %init: f32) { + %1 = arith.addf %in, %init : f32 + linalg.yield %1 : f32 + } + secret.yield %reduced : tensor<33xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } } diff --git a/tests/Transforms/convert_to_ciphertext_semantics/transpose.mlir b/tests/Transforms/convert_to_ciphertext_semantics/transpose.mlir new file mode 100644 index 0000000000..9fe2b1e749 --- /dev/null +++ b/tests/Transforms/convert_to_ciphertext_semantics/transpose.mlir @@ -0,0 +1,17 @@ +// RUN: heir-opt --layout-propagation=min-slot-count=8192 --convert-to-ciphertext-semantics=min-slot-count=8192 %s | FileCheck %s + +#bicyclic = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (-66i0 + 65i1 + slot) mod 2145 = 0 and 0 <= i0 <= 32 and 0 <= i1 <= 64 and 0 <= slot <= 8191 }"> + +// CHECK: func.func @transpose_2d +func.func @transpose_2d(%arg0: !secret.secret> {tensor_ext.layout = #bicyclic}) -> !secret.secret> { + %init = arith.constant dense<0.000000e+00> : tensor<65x33xf32> + // CHECK: secret.generic + // CHECK-NEXT: ^body(%[[IN:.*]]: tensor<1x8192xf32>): + // CHECK: secret.yield %[[IN]] : tensor<1x8192xf32> + %0 = secret.generic(%arg0: !secret.secret>) { + ^body(%input0: tensor<33x65xf32>): + %transposed = linalg.transpose ins(%input0 : tensor<33x65xf32>) outs(%init : tensor<65x33xf32>) permutation = [1, 0] + secret.yield %transposed : tensor<65x33xf32> + } -> !secret.secret> + return %0 : !secret.secret> +} diff --git a/tests/Transforms/layout_optimization/pad_to_layout.mlir b/tests/Transforms/layout_optimization/pad_to_layout.mlir index acc9118f18..ca02cd27e2 100644 --- a/tests/Transforms/layout_optimization/pad_to_layout.mlir +++ b/tests/Transforms/layout_optimization/pad_to_layout.mlir @@ -4,8 +4,8 @@ #layout1 = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : i0 = 0 and ct = 0 and (-48i1 - i2 + slot) mod 512 = 0 and 0 <= i1 <= 9 and 0 <= i2 <= 4095 - 48i1 and i2 <= 47 and 0 <= slot <= 4095 and 4096*floor((-512 + 48i1 + i2)/4096) <= -4096 + 48i1 + i2 }"> #layout2 = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (-48i0 - i1 + slot) mod 512 = 0 and 0 <= i0 <= 9 and 0 <= i1 <= 47 and 0 <= slot <= 4095 and 4096*floor((-512 + 48i0 + i1)/4096) <= -4096 + 48i0 + i1 }"> -// CHECK: #[[layout:.*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (1 - 48i0 - i1 + slot) mod 512 = 0 and 0 <= i0 <= 9 and 0 < i1 <= 48 and 0 <= slot <= 1023 and 1024*floor((511 + 48i0 + i1)/1024) < 48i0 + i1 }"> -// CHECK: #[[layout1:.*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : exists (e0, e1, e2, e3, e4: i0 = 0 and ct = 0 and 512e3 = -i1 + slot + 512e1 and 0 <= i1 <= 1023 and 0 <= slot <= 1023 and -511 + i1 - 1024e0 <= 512e1 <= i1 - 1024e0 and -511 + i1 - 512e1 <= 1024e2 <= i1 - 512e1 and 0 <= e4 <= 9 and -47 + i1 - 512e1 <= 48e4 <= i1 - 512e1) }"> +// CHECK: #[[layout:.*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (1 - 48i0 - i1 + slot) mod 512 = 0 and 0 <= i0 <= 9 and 0 < i1 <= 48 and 0 <= slot <= 4095 and 4096*floor((-513 + 48i0 + i1)/4096) <= -4097 + 48i0 + i1 }"> +// CHECK: #[[layout1:.*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : exists (e0, e1, e2, e3, e4: i0 = 0 and ct = 0 and 512e3 = -i1 + slot + 512e1 and 0 <= i1 <= 4095 and 0 <= slot <= 4095 and -4607 + i1 - 4096e0 <= 512e1 <= -4096 + i1 - 4096e0 and -4607 + i1 - 512e1 <= 4096e2 <= -4096 + i1 - 512e1 and 0 <= e4 <= 9 and -47 + i1 - 512e1 <= 48e4 <= i1 - 512e1) }"> module attributes {backend.lattigo, scheme.ckks} { // CHECK: func.func @tcresnet8small diff --git a/tests/Transforms/layout_propagation/batch_matmul_pt.mlir b/tests/Transforms/layout_propagation/batch_matmul_pt.mlir new file mode 100644 index 0000000000..2daca15fbd --- /dev/null +++ b/tests/Transforms/layout_propagation/batch_matmul_pt.mlir @@ -0,0 +1,64 @@ +// RUN: heir-opt --layout-propagation=min-slot-count=1024 --mlir-print-local-scope %s | FileCheck %s + +#bicyclic = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (-10i0 - 6i1 + slot) mod 15 = 0 and 0 <= i0 <= 2 and 0 <= i1 <= 4 and 0 <= slot <= 1023 }"> +#tricyclic_2_3_5 = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : ct = 0 and (-15i0 - 10i1 - 6i2 + slot) mod 30 = 0 and 0 <= i0 <= 1 and 0 <= i1 <= 2 and 0 <= i2 <= 4 and 0 <= slot <= 1023 }"> +#tricyclic_2_5_7 = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : ct = 0 and (-35i0 - 56i1 - 50i2 + slot) mod 70 = 0 and 0 <= i0 <= 1 and 0 <= i1 <= 4 and 0 <= i2 <= 6 and 0 <= slot <= 1023 }"> + +module { + // CHECK: func @batch_matmul_broadcast + // CHECK-NOT: tensor_ext.convert_layout + // CHECK: linalg.batch_matmul + // CHECK-SAME: #secret.kernel> {tensor_ext.layout = #bicyclic}, %arg1: tensor<2x5x7xf32>) -> !secret.secret> { + %cst = arith.constant dense<0.000000e+00> : tensor<2x3x7xf32> + %empty = tensor.empty() : tensor<2x3x5xf32> + %0 = secret.generic(%arg0: !secret.secret> {tensor_ext.layout = #bicyclic}) { + ^body(%input0: tensor<3x5xf32>): + %bcast = linalg.broadcast ins(%input0 : tensor<3x5xf32>) outs(%empty : tensor<2x3x5xf32>) dimensions = [0] + %1 = linalg.batch_matmul ins(%bcast, %arg1 : tensor<2x3x5xf32>, tensor<2x5x7xf32>) outs(%cst : tensor<2x3x7xf32>) -> tensor<2x3x7xf32> + secret.yield %1 : tensor<2x3x7xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } + + // CHECK: func @batch_matmul_ctpt + // CHECK-NOT: tensor_ext.convert_layout + // CHECK: linalg.batch_matmul + // CHECK-SAME: #secret.kernel> {tensor_ext.layout = #tricyclic_2_3_5}, %arg1: tensor<2x5x7xf32>) -> !secret.secret> { + %cst = arith.constant dense<0.000000e+00> : tensor<2x3x7xf32> + %0 = secret.generic(%arg0: !secret.secret> {tensor_ext.layout = #tricyclic_2_3_5}) { + ^body(%input0: tensor<2x3x5xf32>): + %1 = linalg.batch_matmul ins(%input0, %arg1 : tensor<2x3x5xf32>, tensor<2x5x7xf32>) outs(%cst : tensor<2x3x7xf32>) -> tensor<2x3x7xf32> + secret.yield %1 : tensor<2x3x7xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } + + // CHECK: func @batch_matmul_ptct + // CHECK-NOT: tensor_ext.convert_layout + // CHECK: linalg.batch_matmul + // CHECK-SAME: #secret.kernel, %arg1: !secret.secret> {tensor_ext.layout = #tricyclic_2_5_7}) -> !secret.secret> { + %cst = arith.constant dense<0.000000e+00> : tensor<2x3x7xf32> + %0 = secret.generic(%arg1: !secret.secret> {tensor_ext.layout = #tricyclic_2_5_7}) { + ^body(%input1: tensor<2x5x7xf32>): + %1 = linalg.batch_matmul ins(%arg0, %input1 : tensor<2x3x5xf32>, tensor<2x5x7xf32>) outs(%cst : tensor<2x3x7xf32>) -> tensor<2x3x7xf32> + secret.yield %1 : tensor<2x3x7xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } + + // CHECK: func @batch_matmul_unit_dim + // CHECK: linalg.batch_matmul + // CHECK-SAME: #secret.kernel>, %arg1: tensor<1x5x7xf32>) -> !secret.secret> { + %cst = arith.constant dense<0.000000e+00> : tensor<1x3x7xf32> + %0 = secret.generic(%arg0: !secret.secret>) { + ^body(%input0: tensor<1x3x5xf32>): + %1 = linalg.batch_matmul ins(%input0, %arg1 : tensor<1x3x5xf32>, tensor<1x5x7xf32>) outs(%cst : tensor<1x3x7xf32>) -> tensor<1x3x7xf32> + secret.yield %1 : tensor<1x3x7xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } +} diff --git a/tests/Transforms/layout_propagation/broadcast.mlir b/tests/Transforms/layout_propagation/broadcast.mlir new file mode 100644 index 0000000000..849a489f05 --- /dev/null +++ b/tests/Transforms/layout_propagation/broadcast.mlir @@ -0,0 +1,23 @@ +// RUN: heir-opt --layout-propagation=min-slot-count=8192 %s | FileCheck %s + +#bicyclic = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (-66i0 + 65i1 + slot) mod 2145 = 0 and 0 <= i0 <= 32 and 0 <= i1 <= 64 and 0 <= slot <= 8191 }"> + +// CHECK-DAG: #[[RES_LAYOUT:.*]] = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : ct = 0 and (-66i1 + 65i2 + slot) mod 2145 = 0 and 0 <= i0 <= 1 and 0 <= i1 <= 32 and 0 <= i2 <= 64 and 0 <= slot <= 8191 }"> + +module { + // CHECK: @broadcast_2d_to_3d + // CHECK: linalg.broadcast + // CHECK-SAME: dimensions = [0] + // CHECK-SAME: tensor_ext.layout = #[[RES_LAYOUT]] + // CHECK-NOT: tensor_ext.convert_layout + // CHECK: return + func.func @broadcast_2d_to_3d(%arg0: !secret.secret> {tensor_ext.layout = #bicyclic}) -> !secret.secret> { + %init = arith.constant dense<0.000000e+00> : tensor<2x33x65xf32> + %0 = secret.generic(%arg0: !secret.secret>) { + ^body(%input0: tensor<33x65xf32>): + %broadcasted = linalg.broadcast ins(%input0 : tensor<33x65xf32>) outs(%init : tensor<2x33x65xf32>) dimensions = [0] + secret.yield %broadcasted : tensor<2x33x65xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } +} diff --git a/tests/Transforms/layout_propagation/linalg_reduce.mlir b/tests/Transforms/layout_propagation/linalg_reduce.mlir index a31a9539c8..c4690ea806 100644 --- a/tests/Transforms/layout_propagation/linalg_reduce.mlir +++ b/tests/Transforms/layout_propagation/linalg_reduce.mlir @@ -1,12 +1,16 @@ -// RUN: heir-opt --layout-propagation --fold-convert-layout-into-assign-layout %s | FileCheck %s +// RUN: heir-opt --layout-propagation=min-slot-count=8192 %s | FileCheck %s + +// CHECK-DAG: #[[reduced_layout:.*]] = #tensor_ext.layout<"{ [] -> [ct, slot] : ct = 0 and 0 <= slot <= 8191 }"> +// CHECK-DAG: #[[input_layout:.*]] = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and (-i0 + slot) mod 8 = 0 and 0 <= i0 <= 7 and 0 <= slot <= 8191 }"> +// CHECK-DAG: #[[REDUCED_LAYOUT:.*]] = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and (-i0 + slot) mod 33 = 0 and 0 <= i0 <= 32 and 0 <= slot <= 8191 }"> + +#bicyclic = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (65i0 - 66i1 + slot) mod 2145 = 0 and 0 <= i0 <= 32 and 0 <= i1 <= 64 and 0 <= slot <= 8191 }"> -// CHECK: #[[reduced_layout:.*]] = #tensor_ext.layout<"{ [] -> [ct, slot] : ct = 0 and 0 <= slot <= 1023 }"> -// CHECK: #[[input_layout:.*]] = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and (-i0 + slot) mod 8 = 0 and 0 <= i0 <= 7 and 0 <= slot <= 1023 }"> -// CHECK: @main -// CHECK-SAME: %{{.*}}: !secret.secret> {{{.*}}tensor_ext.layout = #[[input_layout]]}, -// CHECK-SAME: %{{.*}}: !secret.secret> {{{.*}}tensor_ext.layout = #[[input_layout]]} -// CHECK-SAME: -> (!secret.secret> {tensor_ext.layout = #[[reduced_layout]]}) module { + // CHECK: func.func @main + // CHECK-SAME: %{{.*}}: !secret.secret> {{{.*}}tensor_ext.layout = #[[input_layout]]}, + // CHECK-SAME: %{{.*}}: !secret.secret> {{{.*}}tensor_ext.layout = #[[input_layout]]} + // CHECK-SAME: -> (!secret.secret> {tensor_ext.layout = #[[reduced_layout]]}) func.func @main(%arg0: !secret.secret>, %arg1: !secret.secret>) -> !secret.secret> { // CHECK-DAG: %[[cst:.*]] = arith.constant // CHECK-DAG: tensor_ext.assign_layout %[[cst]] @@ -23,4 +27,24 @@ module { } -> !secret.secret> return %0 : !secret.secret> } + + // CHECK: func.func @reduce_bicyclic + // CHECK: ^body(%[[INPUT:.*]]: tensor<33x65xf32>): + // CHECK: %[[INIT:.*]] = tensor_ext.assign_layout + // CHECK: linalg.reduce ins(%[[INPUT]] : tensor<33x65xf32>) outs(%[[INIT]] : tensor<33xf32>) + // CHECK-NOT: tensor_ext.convert_layout + // CHECK-SAME: tensor_ext.layout = #[[REDUCED_LAYOUT]] + func.func @reduce_bicyclic(%arg0: !secret.secret> {tensor_ext.layout = #bicyclic}) -> !secret.secret> { + %cst = arith.constant dense<0.000000e+00> : tensor<33xf32> + %0 = secret.generic(%arg0: !secret.secret>) { + ^body(%input0: tensor<33x65xf32>): + %reduced = linalg.reduce ins(%input0 : tensor<33x65xf32>) outs(%cst : tensor<33xf32>) dimensions = [1] + (%in: f32, %init: f32) { + %1 = arith.addf %in, %init : f32 + linalg.yield %1 : f32 + } + secret.yield %reduced : tensor<33xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } } diff --git a/tests/Transforms/layout_propagation/transpose.mlir b/tests/Transforms/layout_propagation/transpose.mlir new file mode 100644 index 0000000000..8a3a68d0c1 --- /dev/null +++ b/tests/Transforms/layout_propagation/transpose.mlir @@ -0,0 +1,39 @@ +// RUN: heir-opt --layout-propagation=min-slot-count=8192 %s | FileCheck %s + +#bicyclic = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (-66i0 + 65i1 + slot) mod 2145 = 0 and 0 <= i0 <= 32 and 0 <= i1 <= 64 and 0 <= slot <= 8191 }"> +#tricyclic = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : ct = 0 and (2145i0 - 2080i1 - 66i2 + slot) mod 4290 = 0 and 0 <= i0 <= 1 and 0 <= i1 <= 32 and 0 <= i2 <= 64 and 0 <= slot <= 8191 }"> + +// CHECK-DAG: #[[LAYOUT_2D:.*]] = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : ct = 0 and (65i0 - 66i1 + slot) mod 2145 = 0 and 0 <= i0 <= 64 and 0 <= i1 <= 32 and 0 <= slot <= 8191 }"> +// CHECK-DAG: #[[LAYOUT_3D:.*]] = #tensor_ext.layout<"{ [i0, i1, i2] -> [ct, slot] : ct = 0 and (-2080i0 + 2145i1 - 66i2 + slot) mod 4290 = 0 and 0 <= i0 <= 32 and 0 <= i1 <= 1 and 0 <= i2 <= 64 and 0 <= slot <= 8191 }"> + +module { + // CHECK: func.func @transpose_2d + // CHECK: linalg.transpose + // CHECK-SAME: permutation = [1, 0] + // CHECK-SAME: tensor_ext.layout = #[[LAYOUT_2D]] + // CHECK: return + func.func @transpose_2d(%arg0: !secret.secret> {tensor_ext.layout = #bicyclic}) -> !secret.secret> { + %init = arith.constant dense<0.000000e+00> : tensor<65x33xf32> + %0 = secret.generic(%arg0: !secret.secret>) { + ^body(%input0: tensor<33x65xf32>): + %transposed = linalg.transpose ins(%input0 : tensor<33x65xf32>) outs(%init : tensor<65x33xf32>) permutation = [1, 0] + secret.yield %transposed : tensor<65x33xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } + + // CHECK: func.func @transpose_3d + // CHECK: linalg.transpose + // CHECK-SAME: permutation = [1, 0, 2] + // CHECK-SAME: tensor_ext.layout = #[[LAYOUT_3D]] + // CHECK: return + func.func @transpose_3d(%arg0: !secret.secret> {tensor_ext.layout = #tricyclic}) -> !secret.secret> { + %init = arith.constant dense<0.000000e+00> : tensor<33x2x65xf32> + %0 = secret.generic(%arg0: !secret.secret>) { + ^body(%input0: tensor<2x33x65xf32>): + %transposed = linalg.transpose ins(%input0 : tensor<2x33x65xf32>) outs(%init : tensor<33x2x65xf32>) permutation = [1, 0, 2] + secret.yield %transposed : tensor<33x2x65xf32> + } -> !secret.secret> + return %0 : !secret.secret> + } +}