message: Allow HLO version of Argmax as fallback when CustomCall is not supported
diff: @@ -30,7 +30,9 @@ limitations under the License.
  namespace tensorflow {
  namespace {
  -// The logic below uses a custom-call to implement argmax.
  +// The logic below uses a custom-call to implement argmax when possible. When
  +// custom-call is not allowed or input shapes are not supported, this kernel
  +// falls back to using XLA HLO native ArgMax.
  //
  // Also see b/29507024 for first-class XLA support for indexing ops.
  class ArgMaxCustomCallOp : public XlaOpKernel {
  @@ -64,16 +66,26 @@ class ArgMaxCustomCallOp : public XlaOpKernel {
  "Reduction axis ", dim,
  " is empty in shape: ", input_shape.DebugString()));
  + const DataType dtype = output_type(0);
  + xla::PrimitiveType output_type;
  + OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(dtype, &output_type));
  +
  + // Fall back to XLA ArgMax HLO when CustomCall is not allowed or when input
  + // shape isn't supported.
  + if (!XlaContext::Get(ctx).allow_cpu_custom_calls() ||
  + (input_dims != 1 && input_dims != 2)) {
  + xla::XlaOp output = XlaHelpers::ArgMax(ctx->Input(0), output_type, axis);
  + ctx->SetOutput(0, output);
  + return;
  + }
  +
  + xla::XlaOp output;
  // The output shape is the input shape contracted along axis.
  TensorShape output_shape;
  for (int d = 0; d < input_shape.dims() - 1; ++d) {
  output_shape.AddDim(input_shape.dim_size((d < axis) ? d : d + 1));
  }
  - // For now we use a custom-call, only for the 1d and 2d cases.
  - OP_REQUIRES(ctx, XlaContext::Get(ctx).allow_cpu_custom_calls(),
  - errors::InvalidArgument(
  - "ArgMax implementation requires a CustomCall on CPU"));
  xla::XlaBuilder& b = *ctx->builder();
  // XLA passes <out> to the function, so it is not included here.
  @@ -104,27 +116,14 @@ class ArgMaxCustomCallOp : public XlaOpKernel {
  }
  // Tell XLA to call the custom code, defined in
  - // index_ops_kernel_argmax_float_1d.cc.
  - xla::XlaOp output;
  - switch (input_shape.dims()) {
  - case 1:
  + // index_ops_kernel_argmax_float_{1, 2}d.cc.
  + if (input_dims == 1) {
  output = xla::CustomCallWithLayout(&b, "argmax_float_1d_xla_impl", args,
  xla_shape, arg_shapes);
  - break;
  - case 2:
  + } else {
  output = xla::CustomCallWithLayout(&b, "argmax_float_2d_xla_impl", args,
  xla_shape, arg_shapes);
  - break;
  - default:
  - OP_REQUIRES(ctx, false,
  - errors::Unimplemented(
  - "Argmax is only implemented for 1d and 2d tensors"
  - ", but got shape: ",
  - input_shape.DebugString()));
  }
  - const DataType dtype = output_type(0);
  - xla::PrimitiveType output_type;
  - OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(dtype, &output_type));
  output = xla::ConvertElementType(output, output_type);
  ctx->SetOutput(0, output);
  }
  