@@ -16,12 +16,11 @@ limitations under the License.
1616#include " xla/codegen/math/rsqrt.h"
1717
1818#include < cstddef>
19+ #include < vector>
1920
2021#include " absl/log/check.h"
2122#include " absl/log/log.h"
22- #include " absl/status/status.h"
2323#include " absl/status/statusor.h"
24- #include " absl/strings/str_cat.h"
2524#include " llvm/ADT/APInt.h"
2625#include " llvm/IR/Argument.h"
2726#include " llvm/IR/BasicBlock.h"
@@ -38,8 +37,8 @@ limitations under the License.
3837#include " llvm/IR/Value.h"
3938#include " llvm/Support/Casting.h"
4039#include " llvm/Support/TypeSize.h"
40+ #include " llvm/Target/TargetMachine.h"
4141#include " xla/codegen/math/intrinsic.h"
42- #include " xla/service/llvm_ir/llvm_util.h"
4342#include " xla/xla_data.pb.h"
4443
4544namespace xla ::codegen::intrinsics {
@@ -68,19 +67,19 @@ static llvm::Value* NewtonRaphsonRsqrtIteration(llvm::IRBuilder<>& builder,
6867
6968struct RsqrtIntrinsic {
7069 llvm::Intrinsic::ID id;
71- int mask_bits; // Some avx512 calls require masks.
72- bool needs_insert_element;
70+ int mask_bits; // Some avx512 calls require masks.
71+ int needs_insert_element_size; // Some avx512 calls require padding.
7372
7473 static RsqrtIntrinsic ForF32 (size_t num_elements) {
7574 switch (num_elements) {
7675 case 1 :
77- return {llvm::Intrinsic::x86_sse_rsqrt_ss, 0 , true };
76+ return {llvm::Intrinsic::x86_sse_rsqrt_ss, 0 , 4 };
7877 case 4 :
79- return {llvm::Intrinsic::x86_sse_rsqrt_ps, 0 , false };
78+ return {llvm::Intrinsic::x86_sse_rsqrt_ps, 0 , 0 };
8079 case 8 :
81- return {llvm::Intrinsic::x86_avx_rsqrt_ps_256, 0 , false };
80+ return {llvm::Intrinsic::x86_avx_rsqrt_ps_256, 0 , 0 };
8281 case 16 :
83- return {llvm::Intrinsic::x86_avx512_rsqrt14_ps_512, 16 , false };
82+ return {llvm::Intrinsic::x86_avx512_rsqrt14_ps_512, 16 , 0 };
8483 default :
8584 LOG (FATAL) << " Unsupported vector width for rsqrt: "
8685 }
@@ -89,12 +88,18 @@ struct RsqrtIntrinsic {
8988 static RsqrtIntrinsic ForF64 (size_t num_elements) {
9089 // We assume AVX512 is available for F64.
9190 switch (num_elements) {
91+ case 1 :
92+ // Assuming AVX512 is available.
93+ // We don't use x86_avx512_rsqrt14_sd because it also requires padding
94+ // into <2 x double> vectors and it takes an additional source vector
95+ // for the upper bits of the result.
96+ return {llvm::Intrinsic::x86_avx512_rsqrt14_pd_128, 8 , 2 };
9297 case 2 :
93- return {llvm::Intrinsic::x86_avx512_rsqrt14_pd_128, 8 , false };
98+ return {llvm::Intrinsic::x86_avx512_rsqrt14_pd_128, 8 , 0 };
9499 case 4 :
95- return {llvm::Intrinsic::x86_avx512_rsqrt14_pd_256, 8 , false };
100+ return {llvm::Intrinsic::x86_avx512_rsqrt14_pd_256, 8 , 0 };
96101 case 8 :
97- return {llvm::Intrinsic::x86_avx512_rsqrt14_pd_512, 8 , false };
102+ return {llvm::Intrinsic::x86_avx512_rsqrt14_pd_512, 8 , 0 };
98103 default :
99104 LOG (FATAL) << " Unsupported vector width for rsqrt: "
100105 }
@@ -106,36 +111,39 @@ struct RsqrtIntrinsic {
106111 llvm::Intrinsic::getOrInsertDeclaration (module , id);
107112
108113 llvm::Value* y_approx;
109- if (needs_insert_element) {
114+ std::vector<llvm::Value*> args = {x};
115+ if (needs_insert_element_size > 0 ) {
116+ // Pad into a vector of size `needs_insert_element_size`.
110117 llvm::Type* sse_vec_type = llvm::VectorType::get (
111- x->getType ()->getScalarType (), llvm::ElementCount::getFixed (4 ));
118+ x->getType ()->getScalarType (),
119+ llvm::ElementCount::getFixed (needs_insert_element_size));
112120 llvm::Value* vec_x = llvm::UndefValue::get (sse_vec_type);
113121 vec_x = builder.CreateInsertElement (vec_x, x, builder.getInt32 (0 ));
114- llvm::Value* approx_vec =
115- builder.CreateCall (rsqrt_intrinsic, {vec_x}, " y_approx.vec"
116- y_approx = builder.CreateExtractElement (approx_vec, builder.getInt32 (0 ),
117- " y_approx"
118- } else if (mask_bits > 0 ) {
119- llvm::Value* dest = llvm::ConstantFP::get (x->getType (), 0.0 );
122+ args[0 ] = vec_x;
123+ }
124+ if (mask_bits > 0 ) {
125+ llvm::Value* src = llvm::ConstantFP::get (args[0 ]->getType (), 0.0 );
120126 llvm::Value* mask = llvm::ConstantInt::get (
121127 builder.getContext (), llvm::APInt (mask_bits, -1 , true ));
122- y_approx =
123- builder.CreateCall (rsqrt_intrinsic, {x, dest, mask}, " y_approx"
124-
125- } else {
126- y_approx = builder.CreateCall (rsqrt_intrinsic, {x}, " y_approx"
128+ args.push_back (src);
129+ args.push_back (mask);
130+ }
131+ y_approx = builder.CreateCall (rsqrt_intrinsic, args, " y_approx"
132+ if (needs_insert_element_size > 0 ) {
133+ // Extract the result from the padded vector.
134+ y_approx = builder.CreateExtractElement (y_approx, builder.getInt32 (0 ),
135+ " y_approx"
127136 }
128137 return y_approx;
129138 }
130139};
131140
132- absl::StatusOr<llvm::Function*> Rsqrt::CreateDefinition (llvm::Module* module ,
133- Type type) {
141+ absl::StatusOr<llvm::Function*> Rsqrt::CreateDefinition (
142+ llvm::Module* module , llvm::TargetMachine* target_machine, Type type) {
143+ CHECK (type.element_type () == F64 || type.element_type () == F32)
144+ << type.name ();
134145 llvm::Type* input_type = Type::TypeToIrType (type, module ->getContext ());
135146 CHECK (input_type != nullptr );
136- CHECK (input_type->isFloatingPointTy () || input_type->isVectorTy ());
137- CHECK (input_type->getScalarType ()->isFloatTy () ||
138- input_type->getScalarType ()->isDoubleTy ());
139147
140148 llvm::LLVMContext& context = module ->getContext ();
141149 llvm::IRBuilder<> builder (context);
@@ -151,30 +159,44 @@ absl::StatusOr<llvm::Function*> Rsqrt::CreateDefinition(llvm::Module* module,
151159 module ->getOrInsertFunction (Rsqrt::Name (type), function_type)
152160 .getCallee ());
153161
154- llvm::Argument* input_x_arg = func->getArg (0 );
155- input_x_arg ->setName (" x"
162+ llvm::Argument* x = func->getArg (0 );
163+ x ->setName (" x"
156164 llvm::BasicBlock* entry_bb = llvm::BasicBlock::Create (context, " entry"
157- llvm::Value* x = input_x_arg;
158165 builder.SetInsertPoint (entry_bb);
166+
167+ if ((type.element_type () == F64 &&
168+ !target_machine->getTargetFeatureString ().contains (" +avx512f"
169+ !target_machine->getTargetFeatureString ().contains (" +avx"
170+ LOG_EVERY_N (INFO, 1000 )
171+ << " avx not available, falling back to 1 / sqrt(x) for " name ();
172+ // We can't use the same approximation algorithm for F64 without AVX512 or
173+ // anything non-x86 and without avx.
174+ llvm::Value* one = llvm::ConstantFP::get (input_type, 1.0 );
175+ llvm::Value* sqrt_x =
176+ builder.CreateUnaryIntrinsic (llvm::Intrinsic::sqrt, x);
177+ llvm::Value* inv_sqrt_x = builder.CreateFDiv (one, sqrt_x, " inv_sqrt_x"
178+ builder.CreateRet (inv_sqrt_x);
179+ return func;
180+ }
181+
159182 RsqrtIntrinsic rsqrt_intrinsic = input_type->getScalarType ()->isFloatTy ()
160183 ? RsqrtIntrinsic::ForF32 (num_elements)
161184 : RsqrtIntrinsic::ForF64 (num_elements);
162185 llvm::Value* y_approx = rsqrt_intrinsic.CreateCall (builder, x);
163186
164187 llvm::Value* refined_result =
165- NewtonRaphsonRsqrtIteration (builder, input_x_arg , y_approx, input_type);
188+ NewtonRaphsonRsqrtIteration (builder, x , y_approx, input_type);
166189 if (input_type->getScalarType ()->isDoubleTy ()) {
167190 // Do an additional refinement step for F64.
168- refined_result = NewtonRaphsonRsqrtIteration (builder, input_x_arg,
169- refined_result, input_type);
191+ refined_result =
192+ NewtonRaphsonRsqrtIteration (builder, x, refined_result, input_type);
170193 }
171194
172- // Create a mask for special cases (denormals and infinities) to fall back
173- // to the intrinsic's result, matching Eigen's behavior.
174195 const llvm::fltSemantics& semantics =
175196 input_type->getScalarType ()->getFltSemantics ();
176- llvm::Constant* flt_min = llvm::ConstantFP::get (
177- input_type, llvm::APFloat::getSmallestNormalized (semantics));
197+ llvm::APFloat flt_min_val = llvm::APFloat::getSmallestNormalized (semantics);
198+ llvm::Constant* flt_min = llvm::ConstantFP::get (input_type, flt_min_val);
199+
178200 llvm::Constant* inf =
179201 llvm::ConstantFP::get (input_type, llvm::APFloat::getInf (semantics));
180202
@@ -183,11 +205,24 @@ absl::StatusOr<llvm::Function*> Rsqrt::CreateDefinition(llvm::Module* module,
183205 llvm::Value* use_hw_approx_mask =
184206 builder.CreateOr (lt_min_mask, inf_mask, " use_hw_approx_mask"
185207
208+ if (type.element_type () == F64) {
209+ // If the input is very large, the result should be 0.
210+ // This is effectively calculating 1.0 / FLT_MIN, which is the threshold
211+ // where rsqrt(x*x) would be close to the smallest representable number.
212+ llvm::APFloat rsqrt_is_zero_val = llvm::APFloat (semantics, 1 );
213+ rsqrt_is_zero_val.divide (flt_min_val, llvm::APFloat::rmTowardZero);
214+ llvm::Constant* rsqrt_is_zero =
215+ llvm::ConstantFP::get (input_type, rsqrt_is_zero_val);
216+ llvm::Value* rsqrt_is_zero_mask =
217+ builder.CreateFCmpOGE (x, rsqrt_is_zero, " rsqrt_is_zero_mask"
218+ use_hw_approx_mask =
219+ builder.CreateOr (use_hw_approx_mask, rsqrt_is_zero_mask);
220+ }
221+
186222 // If input is normal and finite, use the refined result. Otherwise, use the
187223 // raw hardware approximation.
188224 llvm::Value* result = builder.CreateSelect (use_hw_approx_mask, y_approx,
189225 refined_result, " result"
190-
191226 builder.CreateRet (result);
192227 return func;
193228}
0 commit comments