Skip to content
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

Jump to bottom

Split verifiers and shape functions for ops with regions #401

Merged
merged 7 commits into from
Dec 10, 2022
Merged

Split verifiers and shape functions for ops with regions #401

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
28d2f89
Split verifier for ReduceOp and SortOp
Oct 28, 2022
13f6db1
fix some unintented change
Oct 28, 2022
1a32cc2
Move verifyReduceOp and verifySortOp end of file
Dec 7, 2022
ae7282b
split ReduceWindowOp and WhileOp
Dec 8, 2022
a939920
Remove unused args
Dec 8, 2022
0706fa2
fix clang-tidy
Dec 9, 2022
42fae15
Share the verification logic in shape function too
Dec 10, 2022
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
31 changes: 25 additions & 6 deletions stablehlo/dialect/StablehloOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2942,7 +2942,14 @@ LogicalResult ReduceWindowOp::inferReturnTypeComponents(
location, adaptor.getInputs(), adaptor.getInitValues(),
adaptor.getWindowDimensions(), adaptor.getWindowStrides(),
adaptor.getBaseDilations(), adaptor.getWindowDilations(),
adaptor.getPadding(), adaptor.getBody(), inferredReturnShapes);
adaptor.getPadding(), inferredReturnShapes);
}

LogicalResult ReduceWindowOp::verify() {
return hlo::verifyReduceWindowOp(getLoc(), getInputs(), getInitValues(),
getWindowDimensions(), getWindowStrides(),
getBaseDilations(), getWindowDilations(),
getPadding(), getBody());
}

// Get the operation used for reduction applied to `result_index`th result. Its
Expand Down Expand Up @@ -3302,7 +3309,12 @@ LogicalResult ReduceOp::inferReturnTypeComponents(
ReduceOp::Adaptor adaptor(operands, attributes, regions);
return hlo::inferReduceOp(location, adaptor.getInputs(),
adaptor.getInitValues(), adaptor.getDimensions(),
adaptor.getBody(), inferredReturnShapes);
inferredReturnShapes);
}

LogicalResult ReduceOp::verify() {
return hlo::verifyReduceOp(getLoc(), getInputs(), getInitValues(),
getDimensions(), getBody());
}

LogicalResult ReduceOp::reifyReturnTypeShapes(
Expand Down Expand Up @@ -3792,8 +3804,12 @@ LogicalResult SortOp::inferReturnTypeComponents(
DictionaryAttr attributes, RegionRange regions,
SmallVectorImpl<ShapedTypeComponents>& inferredReturnShapes) {
SortOp::Adaptor adaptor(operands, attributes, regions);
return hlo::inferSortOp(location, adaptor.getInputs(), adaptor.getDimension(),
adaptor.getComparator(), inferredReturnShapes);
return hlo::inferSortOp(location, adaptor.getInputs(), inferredReturnShapes);
}

LogicalResult SortOp::verify() {
return hlo::verifySortOp(getLoc(), getInputs(), getDimension(),
getComparator());
}

//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -4362,8 +4378,11 @@ LogicalResult WhileOp::inferReturnTypes(
DictionaryAttr attributes, RegionRange regions,
SmallVectorImpl<Type>& inferredReturnTypes) {
WhileOp::Adaptor adaptor(operands, attributes, regions);
return hlo::inferWhileOp(location, adaptor.getOperand(), adaptor.getCond(),
adaptor.getBody(), inferredReturnTypes);
return hlo::inferWhileOp(location, adaptor.getOperand(), inferredReturnTypes);
}

LogicalResult WhileOp::verify() {
return hlo::verifyWhileOp(getLoc(), getOperand(), getCond(), getBody());
}

/// Print a `while` op.
Expand Down
8 changes: 8 additions & 0 deletions stablehlo/dialect/StablehloOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1251,6 +1251,8 @@ def StableHLO_WhileOp: StableHLO_Op<"while", [

let results = (outs Variadic<HLO_TensorOrToken>);

let hasVerifier = 1;

let extraClassDeclaration = [{
// Method of OpAsmOpInterface used during custom printing to name the block
// arguments in the nested regions. We name both the condition and the body
Expand Down Expand Up @@ -1433,6 +1435,8 @@ def StableHLO_ReduceOp: StableHLO_ShapedInterfaceOp<"reduce", [

let hasCustomAssemblyFormat = 1;

let hasVerifier = 1;

// TODO(hinsu): Verify that the attached body arguments and results are
// compatible with reduce op's operands.
let regions = (region SizedRegion<1>:$body);
Expand Down Expand Up @@ -2593,6 +2597,8 @@ def StableHLO_SortOp : StableHLO_Op<"sort",
let builders = [
OpBuilder<(ins "ValueRange":$inputs, CArg<"int64_t", "-1">:$dimension,
CArg<"bool", "false">:$is_stable)>];

let hasVerifier = 1;
}

def StableHLO_ReverseOp: StableHLO_Op<"reverse",
Expand Down Expand Up @@ -2783,6 +2789,8 @@ def StableHLO_ReduceWindowOp: StableHLO_Op<"reduce_window", [

let regions = (region SizedRegion<1>:$body);

let hasVerifier = 1;

// Builder for non-variadic version of the operation.
let builders = [
OpBuilder<(ins "Type":$result_type, "Value":$operand,
Expand Down
Loading