maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][Shape] Lower cstr_eq to shape_eq + assert 

Differential Revision: https://reviews.llvm.org/D97860
This commit is contained in:
Benjamin Kramer 2021-03-01 20:41:12 +01:00
parent 19aefd2d5d
commit 73cb58dc48
3 changed files with 21 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp
View File

@ -39,6 +39,7 @@ public:
void mlir::populateConvertShapeConstraintsConversionPatterns(
OwningRewritePatternList &patterns, MLIRContext *ctx) {
patterns.insert<CstrBroadcastableToRequire>(ctx);
patterns.insert<CstrEqToRequire>(ctx);
patterns.insert<ConvertCstrRequireOp>(ctx);
}

7
mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.td
View File

@ -24,4 +24,11 @@ def CstrBroadcastableToRequire : Pat<(Shape_CstrBroadcastableOp $shapes),
(Shape_IsBroadcastableOp $shapes),
(BroadcastableStringAttr))>;
def EqStringAttr : NativeCodeCall<[{
$_builder.getStringAttr("required equal shapes")
}]>;
def CstrEqToRequire : Pat<(Shape_CstrEqOp $shapes),
(Shape_CstrRequireOp (Shape_ShapeEqOp $shapes), (EqStringAttr))>;
#endif // MLIR_CONVERSION_SHAPETOSTANDARD_TD

13
mlir/test/Conversion/ShapeToStandard/convert-shape-constraints.mlir
View File

@ -14,6 +14,19 @@ func @cstr_broadcastable(%arg0: tensor<?xindex>, %arg1: tensor<?xindex>) -> !sha
return %witness : !shape.witness
}
// CHECK-LABEL: func @cstr_eq(
// CHECK-SAME: %[[LHS:.*]]: tensor<?xindex>,
// CHECK-SAME: %[[RHS:.*]]: tensor<?xindex>) -> !shape.witness {
// CHECK: %[[RET:.*]] = shape.const_witness true
// CHECK: %[[EQUAL_IS_VALID:.*]] = shape.shape_eq %[[LHS]], %[[RHS]]
// CHECK: assert %[[EQUAL_IS_VALID]], "required equal shapes"
// CHECK: return %[[RET]] : !shape.witness
// CHECK: }
func @cstr_eq(%arg0: tensor<?xindex>, %arg1: tensor<?xindex>) -> !shape.witness {
%witness = shape.cstr_eq %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>
return %witness : !shape.witness
}
// CHECK-LABEL: func @cstr_require
func @cstr_require(%arg0: i1) -> !shape.witness {
// CHECK: %[[RET:.*]] = shape.const_witness true