1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #include <vector> 17 18 #include "tensorflow/compiler/tf2xla/shape_util.h" 19 #include "tensorflow/compiler/tf2xla/type_util.h" 20 #include "tensorflow/compiler/tf2xla/xla_helpers.h" 21 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" 22 #include "tensorflow/compiler/tf2xla/xla_op_registry.h" 23 #include "tensorflow/compiler/xla/literal_util.h" 24 #include "tensorflow/core/framework/op_kernel.h" 25 #include "tensorflow/core/framework/partial_tensor_shape.h" 26 #include "tensorflow/core/framework/register_types.h" 27 #include "tensorflow/core/framework/tensor.h" 28 #include "tensorflow/core/framework/tensor_types.h" 29 #include "tensorflow/core/framework/types.h" 30 #include "tensorflow/core/kernels/bounds_check.h" 31 #include "tensorflow/core/kernels/concat_lib.h" 32 #include "tensorflow/core/lib/core/status.h" 33 #include "tensorflow/core/platform/types.h" 34 35 namespace tensorflow { 36 namespace { 37 38 // TODO(phawkins): implement double-sized windowed reductions in XLA and remove 39 // the type constraint. 40 constexpr std::array<DataType, 3> kScanOpTypes = { 41 {DT_HALF, DT_BFLOAT16, DT_FLOAT}}; 42 43 class ScanOp : public XlaOpKernel { 44 public: 45 ScanOp(OpKernelConstruction* ctx, bool sum) : XlaOpKernel(ctx), sum_(sum) { 46 OP_REQUIRES_OK(ctx, ctx->GetAttr("reverse", &reverse_)); 47 OP_REQUIRES_OK(ctx, ctx->GetAttr("exclusive", &exclusive_)); 48 } 49 50 void Compile(XlaOpKernelContext* ctx) override { 51 const TensorShape input_shape = ctx->InputShape(0); 52 const TensorShape tensor_axis_shape = ctx->InputShape(1); 53 54 OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(tensor_axis_shape), 55 errors::InvalidArgument("ScanOp: axis must be a scalar, not ", 56 tensor_axis_shape.DebugString())); 57 58 int64 axis; 59 OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(1, &axis)); 60 if (axis < 0) { 61 axis += input_shape.dims(); 62 } 63 OP_REQUIRES( 64 ctx, FastBoundsCheck(axis, input_shape.dims()), 65 errors::InvalidArgument("ScanOp: Expected scan axis in the range [", 66 -input_shape.dims(), ", ", input_shape.dims(), 67 "), but got ", axis)); 68 69 DataType dtype = ctx->input_type(0); 70 71 if (input_shape.num_elements() == 0) { 72 // Exit early if there is nothing to compute. 73 ctx->SetOutput(0, ctx->Input(0)); 74 return; 75 } 76 77 xla::ComputationBuilder* builder = ctx->builder(); 78 79 std::vector<int64> window_strides(input_shape.dims(), 1); 80 std::vector<int64> window_dims(input_shape.dims(), 1); 81 window_dims[axis] = input_shape.dim_size(axis); 82 83 std::vector<std::pair<int64, int64>> padding(input_shape.dims(), {0, 0}); 84 padding[axis].first = input_shape.dim_size(axis) - 1; 85 // In exclusive mode, add an extra padding element so there is a complete 86 // window of padding before the data starts. 87 if (exclusive_) { 88 ++padding[axis].first; 89 } 90 if (reverse_) { 91 std::swap(padding[axis].first, padding[axis].second); 92 } 93 94 xla::ComputationDataHandle input = ctx->Input(0); 95 xla::ComputationDataHandle init; 96 const xla::Computation* reducer; 97 if (sum_) { 98 init = XlaHelpers::Zero(builder, dtype); 99 reducer = ctx->GetOrCreateAdd(dtype); 100 } else { 101 init = XlaHelpers::One(builder, dtype); 102 reducer = ctx->GetOrCreateMul(dtype); 103 } 104 auto output = builder->ReduceWindowWithGeneralPadding( 105 ctx->Input(0), init, *reducer, window_dims, window_strides, padding); 106 107 // In exclusive mode, we have computed an extra element containing the sum 108 // of all the input elements. Slice off this extra "last" element. 109 if (exclusive_) { 110 if (reverse_) { 111 output = builder->SliceInDim(output, 1, input_shape.dim_size(axis) + 1, 112 1, axis); 113 114 } else { 115 output = 116 builder->SliceInDim(output, 0, input_shape.dim_size(axis), 1, axis); 117 } 118 } 119 ctx->SetOutput(0, output); 120 } 121 122 private: 123 const bool sum_; // True=cumulative sum. False=cumulative product. 124 bool reverse_; 125 bool exclusive_; 126 }; 127 128 class CumsumOp : public ScanOp { 129 public: 130 explicit CumsumOp(OpKernelConstruction* ctx) : ScanOp(ctx, /*sum=*/true) {} 131 }; 132 REGISTER_XLA_OP(Name("Cumsum") 133 .TypeConstraint("T", kScanOpTypes) 134 .CompileTimeConstInput("axis"), 135 CumsumOp); 136 137 class CumprodOp : public ScanOp { 138 public: 139 explicit CumprodOp(OpKernelConstruction* ctx) : ScanOp(ctx, /*sum=*/false) {} 140 }; 141 REGISTER_XLA_OP(Name("Cumprod") 142 .TypeConstraint("T", kScanOpTypes) 143 .CompileTimeConstInput("axis"), 144 CumprodOp); 145 146 } // anonymous namespace 147 } // namespace tensorflow 148
