Add rsqrt op to Standard dialect and lower it to LLVM dialect.

[akuegel]

Summary

This adds an rsqrt op to the standard dialect, and lowers
it as 1 / sqrt to the LLVM dialect.

[akuegel]

Sorry, I accidentally created a new patch instead of updating the old one. Abandoning this.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

def RankOp : Std_Op<"rank", [NoSideEffect]> {
   let hasFolder = 1;
   let assemblyFormat = "operands attr-dict `:` type(operands)";
 }
 
 def RemFOp : FloatArithmeticOp<"remf"> {
   let summary = "floating point division remainder operation";
 }
 
+def RsqrtOp : FloatUnaryOp<"rsqrt"> {
+  let summary = "reciprocal of sqrt (1 / sqrt of the specified value)";
+  let description = [{
+    The `rsqrt` operation computes the reciprocal of the square root. It takes
+    one operand and returns one result of the same type. This type may be a
+    float scalar type, a vector whose element type is float, or a tensor of
+    floats. It has no standard attributes.
+  }];
+}
+
 def SignedRemIOp : IntArithmeticOp<"remi_signed"> {
   let summary = "signed integer division remainder operation";
   let hasFolder = 1;
 }
 
 def UnsignedRemIOp : IntArithmeticOp<"remi_unsigned"> {
   let summary = "unsigned integer division remainder operation";
   let hasFolder = 1;

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     rewriter.replaceOpWithNewOp<LLVM::CallOp>(
         op, ArrayRef<Type>(), rewriter.getSymbolRefAttr(freeFunc), casted);
     return matchSuccess();
   }
 
   bool useAlloca;
 };
 
+// A `rsqrt` is converted into `1 / sqrt`.
+struct RsqrtOpLowering : public LLVMLegalizationPattern<RsqrtOp> {
+  using LLVMLegalizationPattern<RsqrtOp>::LLVMLegalizationPattern;
+
+  PatternMatchResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    OperandAdaptor<RsqrtOp> transformed(operands);
+    auto operandType =
+        transformed.operand().getType().dyn_cast_or_null<LLVM::LLVMType>();
+
+    if (!operandType)
+      return matchFailure();
+
+    auto loc = op->getLoc();
+    auto one = rewriter.create<LLVM::ConstantOp>(
+        loc, typeConverter.convertType(operandType),
+        rewriter.getF32FloatAttr(1.0));
+    auto sqrt = rewriter.create<LLVM::SqrtOp>(loc, transformed.operand());
+    rewriter.replaceOpWithNewOp<LLVM::FDivOp>(op, operandType, one, sqrt);
+    return matchSuccess();
+  }
+};
+
 // A `tanh` is converted into a call to the `tanh` function.
 struct TanhOpLowering : public LLVMLegalizationPattern<TanhOp> {
   using LLVMLegalizationPattern<TanhOp>::LLVMLegalizationPattern;
 
   PatternMatchResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
 
patterns.insert<
       IndexCastOpLowering,
       MulFOpLowering,
       MulIOpLowering,
       NegFOpLowering,
       OrOpLowering,
       PrefetchOpLowering,
       RemFOpLowering,
       ReturnOpLowering,
+      RsqrtOpLowering,
       SIToFPLowering,
       SelectOpLowering,
       ShiftLeftOpLowering,
       SignExtendIOpLowering,
       SignedDivIOpLowering,
       SignedRemIOpLowering,
       SignedShiftRightOpLowering,
       SplatOpLowering,

mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir

 // CHECK-LABEL: func @strided_memref(
 func @strided_memref(%ind: index) {
   %0 = alloc()[%ind] : memref<32x64xf32, affine_map<(i, j)[M] -> (32 + M * i + j)>>
   std.return
 }
 
 // -----
 
+// CHECK-LABEL: func @rsqrt(
+// CHECK-SAME: !llvm.float
+func @rsqrt(%arg0 : f32) {
+  // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1.000000e+00 : f32) : !llvm.float
+  // CHECK: %[[SQRT:.*]] = "llvm.intr.sqrt"(%arg0) : (!llvm.float) -> !llvm.float
+  // CHECK: %[[DIV:.*]] = llvm.fdiv %[[ONE]], %[[SQRT]] : !llvm.float
+  %0 = rsqrt %arg0 : f32
+  std.return
+}
+
+// -----
+
+// CHECK-LABEL: func @rsqrt_vector(
+// CHECK-SAME: !llvm<"<4 x float>">
+func @rsqrt_vector(%arg0 : vector<4xf32>) {
+  // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1.000000e+00 : f32) : !llvm<"<4 x float>">
+  // CHECK: %[[SQRT:.*]] = "llvm.intr.sqrt"(%arg0) : (!llvm<"<4 x float>">) -> !llvm<"<4 x float>">
+  // CHECK: %[[DIV:.*]] = llvm.fdiv %[[ONE]], %[[SQRT]] : !llvm<"<4 x float>">
+  %0 = rsqrt %arg0 : vector<4xf32>
+  std.return
+}
+
+// -----
+
 // This should not crash. The first operation cannot be converted, so the
 // second should not match. This attempts to convert `return` to `llvm.return`
 // and complains about non-LLVM types.
 func @unknown_source() -> i32 {
   %0 = "foo"() : () -> i32
   %1 = addi %0, %0 : i32
   // expected-error@+1 {{must be LLVM dialect type}}
   return %1 : i32
 }

mlir/test/IR/core-ops.mlir

^bb42(%t: tensor<4x4x?xf32>, %f: f32, %i: i32, %idx : index, %j: i64, %half: f16):
   %140 = sqrt %f : f32
 
   // CHECK: %{{[0-9]+}} = sqrt %cst_8 : vector<4xf32>
   %141 = sqrt %vcf32 : vector<4xf32>
 
   // CHECK: %{{[0-9]+}} = sqrt %arg0 : tensor<4x4x?xf32>
   %142 = sqrt %t : tensor<4x4x?xf32>
 
+  // CHECK: %{{[0-9]+}} = rsqrt %arg1 : f32
+  %143 = rsqrt %f : f32
+
   return
 }
 
 // CHECK-LABEL: func @affine_apply() {
 func @affine_apply() {
   %i = "std.constant"() {value = 0: index} : () -> index
   %j = "std.constant"() {value = 1: index} : () -> index