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());
}
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,
XlaOp rhs,
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,
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,
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);
StatusOr<XlaOp> AddInstruction(HloInstructionProto&& instr, HloOpcode opcode,
absl::Span<const XlaOp> operands = {});
virtual StatusOr<XlaOp> AddInstruction(HloInstructionProto&& instr,
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,
HloInstructionProto* instr);