experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Override AddInstruction method in MlirHloBuilder

Browse Source 
This way all the supported ops can return an error if MLIR builder is used.

PiperOrigin-RevId: 327091802
Change-Id: I2df09131f89022f21243b173501cfb8dde0573c2
This commit is contained in:
Smit Hinsu 2020-08-17 13:51:58 -07:00 committed by TensorFlower Gardener
parent 1cba239bdd
commit eb377d252e
3 changed files with 17 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/compiler/mlir/xla/ir/mlir_hlo_builder.cc
View File

@ -351,6 +351,13 @@ StatusOr<XlaOp> MlirHloBuilder::InDimBroadcast(
return MakeXlaOp(op.getResult()); return MakeXlaOp(op.getResult());
} }
StatusOr<XlaOp> MlirHloBuilder::AddInstruction(
HloInstructionProto&& instr, HloOpcode opcode,
absl::Span<const XlaOp> operands) {
return Unimplemented("MlirHloBuilder does not support op %s",
HloOpcodeString(opcode));
}
StatusOr<XlaOp> MlirHloBuilder::Compare(const Shape& shape, XlaOp lhs, StatusOr<XlaOp> MlirHloBuilder::Compare(const Shape& shape, XlaOp lhs,
XlaOp rhs, XlaOp rhs,
ComparisonDirection direction) { ComparisonDirection direction) {

3
tensorflow/compiler/mlir/xla/ir/mlir_hlo_builder.h
View File

@ -196,6 +196,9 @@ class MlirHloBuilder : public XlaBuilder {
const Shape& shape, XlaOp operand, const Shape& shape, XlaOp operand,
absl::Span<const int64> broadcast_dimensions) override; absl::Span<const int64> broadcast_dimensions) override;
StatusOr<XlaOp> AddInstruction(HloInstructionProto&& instr, HloOpcode opcode,
absl::Span<const XlaOp> operands) override;
StatusOr<XlaOp> Compare(const Shape& shape, XlaOp lhs, XlaOp rhs, StatusOr<XlaOp> Compare(const Shape& shape, XlaOp lhs, XlaOp rhs,
ComparisonDirection direction) override; ComparisonDirection direction) override;

9
tensorflow/compiler/xla/client/xla_builder.h
View File

@ -784,8 +784,13 @@ class XlaBuilder {
XlaOp RemoveDynamicDimension(XlaOp operand, int64 dimension); XlaOp RemoveDynamicDimension(XlaOp operand, int64 dimension);
StatusOr<XlaOp> AddInstruction(HloInstructionProto&& instr, HloOpcode opcode, virtual StatusOr<XlaOp> AddInstruction(HloInstructionProto&& instr,
absl::Span<const XlaOp> operands = {}); HloOpcode opcode,
absl::Span<const XlaOp> operands);
StatusOr<XlaOp> AddInstruction(HloInstructionProto&& instr,
HloOpcode opcode) {
return AddInstruction(std::move(instr), opcode, /*operands=*/{});
}
void AddCalledComputation(const XlaComputation& computation, void AddCalledComputation(const XlaComputation& computation,
HloInstructionProto* instr); HloInstructionProto* instr);