Skip to content

lower linalg.broadcast #3234

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

lower linalg.broadcast #3234

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 52 additions & 2 deletions lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -749,6 +749,55 @@ class ConvertLinalgReduce : public ConversionBase<linalg::ReduceOp> {
}
};

struct ConvertLinalgBroadcast : public ConversionBase<linalg::BroadcastOp> {
public:
using ConversionBase<linalg::BroadcastOp>::ConversionBase;

LogicalResult matchAndRewrite(
linalg::BroadcastOp op, OpAdaptor adaptor,
ContextAwareConversionPatternRewriter& rewriter) const final {
Value input = adaptor.getOperands()[0];
Value output = adaptor.getOperands()[1];

LayoutAttr insLayout = getLayoutAttr(input);
LayoutAttr outsLayout = getLayoutAttr(output);
if (!insLayout || !outsLayout) {
return rewriter.notifyMatchFailure(
op, "missing new layout attribute for inputs/outputs");
}

auto outputShape =
cast<RankedTensorType>(op->getResult(0).getType()).getShape();

auto dimensions = op.getDimensions();
auto maybeComposedLayout = getBroadcastTranslationRelation(
insLayout.getIntegerRelation(), outsLayout.getIntegerRelation(),
outputShape, dimensions);
if (failed(maybeComposedLayout)) {
return rewriter.notifyMatchFailure(op,
"failed to compute composed layout");
}
IntegerRelation composedLayout = maybeComposedLayout.value();

ImplicitLocOpBuilder b(op.getLoc(), rewriter);
LayoutAttr newLayoutAttr =
LayoutAttr::getFromIntegerRelation(getContext(), composedLayout);

auto resultCiphertextSemanticType =
cast<RankedTensorType>(getTypeConverter()->convertType(
op->getResult(0).getType(), outsLayout));

auto remapAndExtract = remapAndExtractResult(b, input, newLayoutAttr,
resultCiphertextSemanticType);

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

struct ConvertLinalgDot : public ConversionBase<linalg::DotOp> {
public:
using ConversionBase<linalg::DotOp>::ConversionBase;
Expand Down Expand Up @@ -2616,8 +2665,9 @@ struct ConvertToCiphertextSemantics
});

patterns.add<ConvertAnyAddingMaterializedAttr, ConvertConvertLayout,
ConvertFunc, ConvertLinalgMatmul, ConvertLinalgBatchMatmul,
ConvertLinalgReduce, ConvertLinalgDot, ConvertSecretGeneric,
ConvertFunc, ConvertLinalgBroadcast, ConvertLinalgMatmul,
ConvertLinalgBatchMatmul, ConvertLinalgReduce,
ConvertLinalgDot, ConvertSecretGeneric,
ConvertTensorCollapseShape, ConvertTensorExpandShape,
ConvertTensorExtractLayout, ConvertTensorExtractSlice,
ConvertTensorInsertLayout, ConvertTensorInsertSlice>(
Expand Down
123 changes: 116 additions & 7 deletions lib/Transforms/LayoutPropagation/LayoutPropagation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,7 @@ struct LayoutPropagation : impl::LayoutPropagationBase<LayoutPropagation> {
LogicalResult visitOperation(ExpandShapeOp op);
LogicalResult visitOperation(GenericOp op);
LogicalResult visitOperation(ReduceOp op);
LogicalResult visitOperation(linalg::BroadcastOp op);
LogicalResult visitOperation(Conv1DOp op);
LogicalResult visitOperation(Conv1DNcwFcwOp op);
LogicalResult visitOperation(Conv2DOp op);
Expand Down Expand Up @@ -193,6 +194,7 @@ struct LayoutPropagation : impl::LayoutPropagationBase<LayoutPropagation> {
CompatibilityResult hasCompatibleArgumentLayouts(Conv2DOp op);
CompatibilityResult hasCompatibleArgumentLayouts(Conv2DNchwFchwOp op);
CompatibilityResult hasCompatibleArgumentLayouts(ReduceOp op);
CompatibilityResult hasCompatibleArgumentLayouts(linalg::BroadcastOp op);
CompatibilityResult hasCompatibleArgumentLayouts(VecmatOp op);
CompatibilityResult hasCompatibleArgumentLayouts(MatvecOp op);
CompatibilityResult hasCompatibleArgumentLayouts(MatmulOp op);
Expand All @@ -204,9 +206,13 @@ struct LayoutPropagation : impl::LayoutPropagationBase<LayoutPropagation> {
// Op-specific overrides
void rectifyIncompatibleOperandLayouts(DotOp op);
void rectifyIncompatibleOperandLayouts(ReduceOp op);
void rectifyIncompatibleOperandLayouts(linalg::BroadcastOp op);
void rectifyIncompatibleOperandLayouts(tensor::InsertOp op);
void rectifyIncompatibleOperandLayouts(tensor::InsertSliceOp op);

FailureOr<LayoutAttr> inferBroadcastOutputLayout(LayoutAttr insLayout,
linalg::BroadcastOp op);

// Return the default layout for a given type
FailureOr<LayoutAttr> defaultLayoutForType(Type type);
FailureOr<LayoutAttr> defaultLayoutForScalarType(Type type);
Expand Down Expand Up @@ -338,9 +344,9 @@ LogicalResult LayoutPropagation::visitOperation(Operation* op) {
// secret ops
.Case<GenericOp, YieldOp>([&](auto op) { return visitOperation(op); })
// linalg ops
.Case<DotOp, MatvecOp, VecmatOp, ReduceOp, MatmulOp, BatchMatmulOp,
Conv1DOp, Conv1DNcwFcwOp, Conv2DOp, Conv2DNchwFchwOp>(
[&](auto op) { return visitOperation(op); })
.Case<DotOp, MatvecOp, VecmatOp, ReduceOp, linalg::BroadcastOp, MatmulOp,
BatchMatmulOp, Conv1DOp, Conv1DNcwFcwOp, Conv2DOp,
Conv2DNchwFchwOp>([&](auto op) { return visitOperation(op); })
// affine ops
.Case<affine::AffineForOp>([&](auto op) { return visitOperation(op); })
// tensor ops
Expand Down Expand Up @@ -1411,8 +1417,8 @@ CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
affine::AffineYieldOp>(
[&](auto op) { return CompatibilityResult{true, std::nullopt}; })
// Ops with special rules
.Case<DotOp, ReduceOp, MatvecOp, VecmatOp, Conv1DOp, Conv1DNcwFcwOp,
Conv2DOp, Conv2DNchwFchwOp, tensor::InsertSliceOp>(
.Case<DotOp, ReduceOp, linalg::BroadcastOp, MatvecOp, VecmatOp, Conv1DOp,
Conv1DNcwFcwOp, Conv2DOp, Conv2DNchwFchwOp, tensor::InsertSliceOp>(
[&](auto op) { return hasCompatibleArgumentLayouts(op); })
// By default, assume operands must all have the same layout.
.Default([&](Operation* op) {
Expand Down Expand Up @@ -1640,7 +1646,8 @@ void LayoutPropagation::rectifyIncompatibleOperandLayouts(Operation* op) {

TypeSwitch<Operation*>(op)
// Ops with special rules
.Case<DotOp, ReduceOp, tensor::InsertOp, tensor::InsertSliceOp>(
.Case<DotOp, ReduceOp, linalg::BroadcastOp, tensor::InsertOp,
tensor::InsertSliceOp>(
[&](auto op) { return rectifyIncompatibleOperandLayouts(op); })
.Default([&](Operation* op) {
// Default target layout is chosen arbitrarily as the first operand's
Expand Down Expand Up @@ -1932,7 +1939,109 @@ void LayoutPropagation::runOnOperation() {
cleanupPatterns.add<tensor_ext::FoldConvertLayoutIntoAssignLayoutPattern>(
&getContext());
(void)applyPatternsGreedily(getOperation(), std::move(cleanupPatterns));
};
}

FailureOr<LayoutAttr> LayoutPropagation::inferBroadcastOutputLayout(
LayoutAttr insLayout, linalg::BroadcastOp op) {
RankedTensorType destType =
cast<RankedTensorType>(op->getResult(0).getType());

auto dimensions = op.getDimensions();
SmallVector<unsigned int> addedDims(dimensions.begin(), dimensions.end());

// 1. Compute the slot span of the input layout
auto maybeSpan = getSlotSpan(insLayout.getIntegerRelation());
if (!maybeSpan.has_value()) {
return failure();
}
int64_t span = maybeSpan.value();

// 2. Sort the added dimensions descending for stride calculation
SmallVector<unsigned int> sortedDims = addedDims;
llvm::sort(sortedDims, std::greater<unsigned int>());

// 3. Compute strides
DenseMap<unsigned int, int64_t> strides;
int64_t currentStride = span;
for (unsigned int d : sortedDims) {
strides[d] = currentStride;
currentStride *= destType.getDimSize(d);
}

// 4. Apply broadcasting in ascending order of indices
llvm::sort(addedDims); // Sort ascending
IntegerRelation currentRel = insLayout.getIntegerRelation();
for (unsigned int d : addedDims) {
currentRel = broadcastDimension(currentRel, destType, d, strides[d]);
}

return LayoutAttr::getFromIntegerRelation(&getContext(), currentRel);
}

LogicalResult LayoutPropagation::visitOperation(linalg::BroadcastOp op) {
auto linalgOp = cast<linalg::LinalgOp>(op.getOperation());
LayoutAttr insLayout = getComposedLayoutAttr(linalgOp.getDpsInputs()[0]);
auto maybeResultLayout = inferBroadcastOutputLayout(insLayout, op);
if (failed(maybeResultLayout)) {
return op.emitError("failed to infer broadcast output layout");
}
LayoutAttr resultLayout = maybeResultLayout.value();
assignedLayouts.insert({op->getResult(0), resultLayout});
debugAssignLayout(op->getResult(0), resultLayout);
setResultLayoutAttr(op);
return success();
}

CompatibilityResult LayoutPropagation::hasCompatibleArgumentLayouts(
linalg::BroadcastOp op) {
auto linalgOp = cast<linalg::LinalgOp>(op.getOperation());
if (!assignedLayouts.contains(linalgOp.getDpsInputs()[0])) {
return {false, op->emitError("input tensor has no assigned layout")};
}
if (!assignedLayouts.contains(linalgOp.getDpsInits()[0])) {
return {false, op->emitError("init tensor has no assigned layout")};
}

LayoutAttr insLayout = getComposedLayoutAttr(linalgOp.getDpsInputs()[0]);
LayoutAttr outsLayout = getComposedLayoutAttr(linalgOp.getDpsInits()[0]);

auto maybeExpectedOutsLayout = inferBroadcastOutputLayout(insLayout, op);
if (failed(maybeExpectedOutsLayout)) {
return {false, op.emitError("failed to infer broadcast output layout")};
}

if (maybeExpectedOutsLayout.value() != outsLayout) {
return {false, std::nullopt};
}

return {true, std::nullopt};
}

void LayoutPropagation::rectifyIncompatibleOperandLayouts(
linalg::BroadcastOp op) {
auto linalgOp = cast<linalg::LinalgOp>(op.getOperation());
LayoutAttr insLayout = getComposedLayoutAttr(linalgOp.getDpsInputs()[0]);
LayoutAttr outsLayout = getComposedLayoutAttr(linalgOp.getDpsInits()[0]);

auto maybeExpectedOutsLayout = inferBroadcastOutputLayout(insLayout, op);
assert(succeeded(maybeExpectedOutsLayout) && "expected to infer layout");
LayoutAttr expectedOutsLayout = maybeExpectedOutsLayout.value();

if (expectedOutsLayout != outsLayout) {
mlir::IRRewriter builder(&getContext());
builder.setInsertionPoint(op);
// Convert the init operand to the expected layout.
Value init = linalgOp.getDpsInits()[0];
ConvertLayoutOp convertOp = ConvertLayoutOp::create(
builder, op->getLoc(), init, outsLayout, expectedOutsLayout);
Value toReplace = convertOp.getResult();
builder.replaceUsesWithIf(init, toReplace, [&](OpOperand& operand) {
return operand.getOwner() == op;
});
assignedLayouts.insert({toReplace, expectedOutsLayout});
setResultLayoutAttr(convertOp);
}
}

} // namespace heir
} // namespace mlir
41 changes: 41 additions & 0 deletions lib/Transforms/SoftmaxToCgfSoftmax/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
load("@heir//lib/Transforms:transforms.bzl", "add_heir_transforms")
load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
load("@rules_cc//cc:cc_library.bzl", "cc_library")

package(
default_applicable_licenses = ["@heir//:license"],
default_visibility = ["//visibility:public"],
)

cc_library(
name = "SoftmaxToCgfSoftmax",
srcs = ["SoftmaxToCgfSoftmax.cpp"],
hdrs = ["SoftmaxToCgfSoftmax.h"],
deps = [
":pass_inc_gen",
"@heir//lib/Dialect/MathExt/IR:Dialect",
"@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:DialectUtils",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:LinalgDialect",
"@llvm-project//mlir:MathDialect",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Support",
"@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:TransformUtils",
],
)

add_heir_transforms(
generated_target_name = "pass_inc_gen",
pass_name = "SoftmaxToCgfSoftmax",
td_file = "SoftmaxToCgfSoftmax.td",
)

cc_binary(
name = "cgf_softmax_fuzz",
srcs = ["cgf_softmax_fuzz.cpp"],
deps = [
"@com_google_absl//absl/random",
],
)
Loading
Loading