diff options
| author |  Mark Heffernan <meheff@google.com> | 2018-10-08 14:26:43 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-10-08 14:34:02 -0700 |
| commit | 396a8a4105edd409d0821c4d5d0b920b315ffb72 (patch) | |
| tree | 428350d427ffb29470e284077a2734b273b7cc4d /tensorflow/compiler/xla/service/hlo_instruction.h | |
| parent | bc5635dc3ac78007caee88fabd81d23ad945b637 (diff) | |
Add custom call with layout constraints.
Add a variant of CustomCall which specifies arbitrary layout constraints on the operands and result. The existing non-layout-constrained CustomCall is changed to have no layout preference and can now be assigned arbitrary layouts by layout assignment.
PiperOrigin-RevId: 216249615
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_instruction.h')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_instruction.h | 10 |
1 files changed, 10 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index 374862c4b6..44f776ebac 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -734,6 +734,16 @@ class HloInstruction { const Shape& shape, absl::Span<HloInstruction* const> operands, absl::string_view custom_call_target, absl::string_view opaque = ""); + // Overload which constrains the layouts of the operand and result. 'shape' + // and 'operand_shapes_with_layout' must have layouts. + // 'operand_shapes_with_layout' must have a compatible element for each + // operand. + static std::unique_ptr<HloInstruction> CreateCustomCall( + const Shape& shape, absl::Span<HloInstruction* const> operands, + absl::string_view custom_call_target, + absl::Span<const Shape> operand_shapes_with_layout, + absl::string_view opaque = ""); + // Creates a tuple instruction with the given elements. This is a convenience // wrapper around CreateVariadic. static std::unique_ptr<HloInstruction> CreateTuple( |