colin
/
mlir-hlo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Implement InferShapedTypeOpInterface and use inferReturnTypes for mhlo.imag and mhlo.real 

This makes the lhlo lowering work with dynamic shapes.

PiperOrigin-RevId: 334553472
This commit is contained in:
Benjamin Kramer 2020-09-30 02:01:45 -07:00 committed by TensorFlow MLIR Team
parent 39389587d2
commit c8919f8419
3 changed files with 40 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
View File

@ -193,12 +193,9 @@ def HLO_Expm1Op: HLO_UnaryElementwiseOp<"exponential_minus_one",
def HLO_FloorOp: HLO_UnaryElementwiseOp<"floor",
[NoSideEffect, SameOperandsAndResultType], HLO_FpTensor>, BASE_HLO_FloorOp;
def HLO_ImagOp: HLO_Op<
"imag", [NoSideEffect, SameOperandsAndResultShape]>, BASE_HLO_ImagOp {
let builders = [OpBuilder<
"OpBuilder &, OperationState &tblgen_state, Value val">];
let arguments = (ins HLO_ComplexTensor);
def HLO_ImagOp: HLO_UnaryElementwiseOp<"imag",
[NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>],
HLO_ComplexTensor>, BASE_HLO_ImagOp {
let results = (outs HLO_FpTensor);
let hasFolder = 1;
}
@ -237,12 +234,9 @@ def HLO_PopulationCountOp: HLO_UnaryElementwiseOp<"popcnt",
[NoSideEffect, SameOperandsAndResultType], HLO_IntTensor>,
BASE_HLO_PopulationCountOp;
def HLO_RealOp: HLO_Op<
"real", [NoSideEffect, SameOperandsAndResultShape]>, BASE_HLO_RealOp {
let builders = [OpBuilder<
"OpBuilder &, OperationState &tblgen_state, Value val">];
let arguments = (ins HLO_ComplexTensor);
def HLO_RealOp: HLO_UnaryElementwiseOp<"real",
[NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>],
HLO_ComplexTensor>, BASE_HLO_RealOp {
let results = (outs HLO_FpTensor);
let hasFolder = 1;
}

14
lib/Dialect/mhlo/IR/hlo_ops.cc
View File

@ -932,8 +932,11 @@ Type CreateRealType(Type type) {
}
} // namespace
void ImagOp::build(OpBuilder& builder, OperationState& state, Value val) {
build(builder, state, CreateRealType(val.getType()), val);
LogicalResult ImagOp::inferReturnTypes(
MLIRContext*, Optional<Location>, ValueRange operands, DictionaryAttr,
RegionRange, SmallVectorImpl<Type>& inferredReturnTypes) {
inferredReturnTypes.push_back(CreateRealType(operands[0].getType()));
return success();
}
OpFoldResult ImagOp::fold(ArrayRef<Attribute> operands) {
@ -945,8 +948,11 @@ OpFoldResult ImagOp::fold(ArrayRef<Attribute> operands) {
return {};
}
void RealOp::build(OpBuilder& builder, OperationState& state, Value val) {
build(builder, state, CreateRealType(val.getType()), val);
LogicalResult RealOp::inferReturnTypes(
MLIRContext*, Optional<Location>, ValueRange operands, DictionaryAttr,
RegionRange, SmallVectorImpl<Type>& inferredReturnTypes) {
inferredReturnTypes.push_back(CreateRealType(operands[0].getType()));
return success();
}
OpFoldResult RealOp::fold(ArrayRef<Attribute> operands) {

24
tests/hlo-legalize-to-lhlo.mlir
View File

@ -248,6 +248,18 @@ func @real(%operand: memref<2x2xcomplex<f32>>, %result: memref<2x2xf32>) {
// -----
// BOTH-LABEL: func @real_dyn
func @real_dyn(%operand: memref<?xcomplex<f32>>, %result: memref<?xf32>) {
%tensor_operand = tensor_load %operand : memref<?xcomplex<f32>>
%tensor_result = "mhlo.real"(%tensor_operand)
: (tensor<?xcomplex<f32>>) -> tensor<?xf32>
// BOTH: "lmhlo.real"(%{{.*}}, %{{.*}})
tensor_store %tensor_result, %result : memref<?xf32>
return
}
// -----
// BOTH-LABEL: func @imag
func @imag(%operand: memref<2x2xcomplex<f32>>, %result: memref<2x2xf32>) {
%tensor_operand = tensor_load %operand : memref<2x2xcomplex<f32>>
@ -260,6 +272,18 @@ func @imag(%operand: memref<2x2xcomplex<f32>>, %result: memref<2x2xf32>) {
// -----
// BOTH-LABEL: func @imag_dyn
func @imag_dyn(%operand: memref<?xcomplex<f32>>, %result: memref<?xf32>) {
%tensor_operand = tensor_load %operand : memref<?xcomplex<f32>>
%tensor_result = "mhlo.imag"(%tensor_operand)
: (tensor<?xcomplex<f32>>) -> tensor<?xf32>
// BOTH: "lmhlo.imag"(%{{.*}}, %{{.*}})
tensor_store %tensor_result, %result : memref<?xf32>
return
}
// -----
// BOTH-LABEL: func @iota
func @iota(%result: memref<10xi32>) {
%tensor_result = "mhlo.iota"()