| 125 | } |
| 126 | |
| 127 | XlaOp BuildCholesky(XlaOp a, int64 block_size, |
| 128 | PrecisionConfig::Precision precision) { |
| 129 | XlaBuilder* builder = a.builder(); |
| 130 | return builder->ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 131 | TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); |
| 132 | const int ndims = a_shape.rank(); |
| 133 | if (ndims < 2) { |
| 134 | return InvalidArgument( |
| 135 | "Argument to Cholesky must have rank >= 2; shape was %s", |
| 136 | a_shape.ToString()); |
| 137 | } |
| 138 | |
| 139 | const int64 n = ShapeUtil::GetDimension(a_shape, -1); |
| 140 | if (n != ShapeUtil::GetDimension(a_shape, -2)) { |
| 141 | return InvalidArgument( |
| 142 | "Argument to Cholesky must be batched square matrices; got shape %s", |
| 143 | ShapeUtil::HumanString(a_shape)); |
| 144 | } |
| 145 | |
| 146 | if (primitive_util::IsComplexType(a_shape.element_type())) { |
| 147 | return Unimplemented( |
| 148 | "Complex types are not implemented in Cholesky; got shape %s", |
| 149 | ShapeUtil::HumanString(a_shape)); |
| 150 | } |
| 151 | |
| 152 | if (block_size < 1) { |
| 153 | return InvalidArgument( |
| 154 | "block_size argument to Cholesky must be >= 1; got %d", block_size); |
| 155 | } |
| 156 | |
| 157 | // Blocked left-looking Cholesky factorization. |
| 158 | // Algorithm 1 from |
| 159 | // Haidar, Azzam, et al. "High-performance Cholesky factorization for |
| 160 | // GPU-only execution." Proceedings of General Purpose GPUs. ACM, 2017. |
| 161 | XlaOp l = ZerosLike(a); |
| 162 | XlaOp seen_error = ConstantR0<bool>(builder, false); |
| 163 | for (int64 i = 0; i < n; i += block_size) { |
| 164 | int64 k = std::min(block_size, n - i); |
| 165 | if (i > 0) { |
| 166 | // TODO(phawkins): consider implementing SYRK for the diagonal part of |
| 167 | // the panel. |
| 168 | // a[i:, i:i+k] -= np.dot(l[i:, :i], np.transpose(l[i:i+k, :i])) |
| 169 | auto lhs = SliceInMinorDims(l, {i, 0}, {n, i}); |
| 170 | auto rhs = SliceInMinorDims(l, {i, 0}, {i + k, i}); |
| 171 | auto delta = BatchDot(lhs, false, rhs, true, precision); |
| 172 | auto before = SliceInMinorDims(a, {i, i}, {n, i + k}); |
| 173 | a = UpdateSliceInMinorDims(a, before - delta, {i, i}); |
| 174 | } |
| 175 | |
| 176 | // l[i:i+k, i:i+k] = cholesky_unblocked(a[i:i+k, i:i+k]) |
| 177 | auto x = SliceInMinorDims(a, {i, i}, {i + k, i + k}); |
| 178 | XlaOp factorized; |
| 179 | XlaOp factorized_error; |
| 180 | std::tie(factorized, factorized_error) = CholeskyUnblocked(x, precision); |
| 181 | seen_error = Or(seen_error, factorized_error); |
| 182 | l = UpdateSliceInMinorDims(l, factorized, {i, i}); |
| 183 | |
| 184 | if (i + k < n) { |
no test coverage detected