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]

Fix ClangTidy warnings.

[ftynse]

Mostly good, let's make sure it works for other types than f32.

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

1688

This may not necessarily work if the operand type is not float. Maybe you could dispatch to different attributes based on type?

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

23

Could we also have a test for f64/double?

[akuegel]

Don't hard-code constant to have type F32.

[akuegel]

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

1688

I tried to fix this by using getFloatAttr(operandType, 1.0) instead. Not sure whether operandType is the actual type it expects here though. Now I get llvm.float / llvm.double as types for the constants, before it was f32.

[akuegel]

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

1688

Ah, and if I run the test without -c opt, I also get an assert failure. So as I feared, this is not the right fix yet. Will try to figure out how to make it use the F32 / F64 type as required.

[akuegel]

Fix type issues.

[akuegel]

PTAL. The tests seem to work now.

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

48

Does this look ok? I am not sure whether the operand type of constant is allowed to be different from the result type. But seems there is no assert that is triggered, so maybe it supports broadcasts?

[ftynse]

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

48

I think you need something like llvm.mlir.constant(dense<1.0> : vector<4xf32>) : !llvm<"<4 x float>">. We have SplatElementsAttr to broadcast constants in a vector, which is printed like that.

[akuegel]

Handle vector types.

[akuegel]

Sorry for not following up sooner, I was busy with other stuff yesterday. Today, with some help by Stephan I was able to make this work. I copied the handling with unrolling multi-dimensional vectors from the generic elementwise op handling. This could potentially be refactored, especially I wonder whether Tanh should also use this.

[ftynse]

I copied the handling with unrolling multi-dimensional vectors from the generic elementwise op handling. This could potentially be refactored, especially I wonder whether Tanh should also use this.

I think I've seen a commit that removed the default tanh lowering recently. Anyway, a refactoring is welcome if you think it makes sense, in a separate commit.

[akuegel]

Landed as revision 91acb5b3e1c3
Unforunately I forgot to reset the link to this review, after I had accidentally created https://reviews.llvm.org/D75353

[rriddle]

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

1676

Don't use dyn_cast_or_null here, the result of getType is never null. Use dyn_cast if you know the input is non-null.

[akuegel]

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

1676

Thanks for the suggestion :)
I made that part of a refactoring CL that I was planning to do, anyway: https://reviews.llvm.org/D75733

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

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/ADT/TypeSwitch.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/Functional.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Type.h"
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 resultType = *op->result_type_begin();
+    auto floatType = getElementTypeOrSelf(resultType).cast<FloatType>();
+    auto floatOne = rewriter.getFloatAttr(floatType, 1.0);
+
+    if (!operandType.isArrayTy()) {
+      LLVM::ConstantOp one;
+      if (operandType.isVectorTy()) {
+        one = rewriter.create<LLVM::ConstantOp>(
+            loc, operandType,
+            SplatElementsAttr::get(resultType.cast<ShapedType>(), floatOne));
+      } else {
+        one = rewriter.create<LLVM::ConstantOp>(loc, operandType, floatOne);
+      }
+      auto sqrt = rewriter.create<LLVM::SqrtOp>(loc, transformed.operand());
+      rewriter.replaceOpWithNewOp<LLVM::FDivOp>(op, operandType, one, sqrt);
+      return matchSuccess();
+    }
+
+    auto vectorType = resultType.dyn_cast<VectorType>();
+    if (!vectorType)
+      return this->matchFailure();
+
+    auto vectorTypeInfo =
+        extractNDVectorTypeInfo(vectorType, this->typeConverter);
+    auto llvmVectorTy = vectorTypeInfo.llvmVectorTy;
+    if (!llvmVectorTy || operandType != vectorTypeInfo.llvmArrayTy)
+      return this->matchFailure();
+
+    Value desc = rewriter.create<LLVM::UndefOp>(loc, operandType);
+    nDVectorIterate(vectorTypeInfo, rewriter, [&](ArrayAttr position) {
+      // For this unrolled `position` corresponding to the `linearIndex`^th
+      // element, extract operand vectors
+      auto extractedOperand = rewriter.create<LLVM::ExtractValueOp>(
+          loc, llvmVectorTy, operands[0], position);
+      auto splatAttr = SplatElementsAttr::get(
+          mlir::VectorType::get(
+              {llvmVectorTy.getUnderlyingType()->getVectorNumElements()},
+              floatType),
+          floatOne);
+      auto one =
+          rewriter.create<LLVM::ConstantOp>(loc, llvmVectorTy, splatAttr);
+      auto sqrt =
+          rewriter.create<LLVM::SqrtOp>(loc, llvmVectorTy, extractedOperand);
+      auto div = rewriter.create<LLVM::FDivOp>(loc, llvmVectorTy, one, sqrt);
+      desc = rewriter.create<LLVM::InsertValueOp>(loc, operandType, desc, div,
+                                                  position);
+    });
+    rewriter.replaceOp(op, desc);
+
+    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_double(
+// CHECK-SAME: !llvm.double
+func @rsqrt_double(%arg0 : f64) {
+  // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1.000000e+00 : f64) : !llvm.double
+  // CHECK: %[[SQRT:.*]] = "llvm.intr.sqrt"(%arg0) : (!llvm.double) -> !llvm.double
+  // CHECK: %[[DIV:.*]] = llvm.fdiv %[[ONE]], %[[SQRT]] : !llvm.double
+  %0 = rsqrt %arg0 : f64
+  std.return
+}
+
+// -----
+
+// CHECK-LABEL: func @rsqrt_vector(
+// CHECK-SAME: !llvm<"<4 x float>">
+func @rsqrt_vector(%arg0 : vector<4xf32>) {
+  // CHECK: %[[ONE:.*]] = llvm.mlir.constant(dense<1.000000e+00> : vector<4xf32>) : !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
+}
+
+// -----
+
+// CHECK-LABEL: func @rsqrt_multidim_vector(
+// CHECK-SAME: !llvm<"[4 x <3 x float>]">
+func @rsqrt_multidim_vector(%arg0 : vector<4x3xf32>) {
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractvalue %arg0[0] : !llvm<"[4 x <3 x float>]">
+  // CHECK: %[[ONE:.*]] = llvm.mlir.constant(dense<1.000000e+00> : vector<3xf32>) : !llvm<"<3 x float>">
+  // CHECK: %[[SQRT:.*]] = "llvm.intr.sqrt"(%[[EXTRACT]]) : (!llvm<"<3 x float>">) -> !llvm<"<3 x float>">
+  // CHECK: %[[DIV:.*]] = llvm.fdiv %[[ONE]], %[[SQRT]] : !llvm<"<3 x float>">
+  // CHECK: %[[INSERT:.*]] = llvm.insertvalue %[[DIV]], %0[0] : !llvm<"[4 x <3 x float>]">
+  %0 = rsqrt %arg0 : vector<4x3xf32>
+  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):
   %142 = sqrt %t : tensor<4x4x?xf32>
 
   // CHECK: = fpext {{.*}} : vector<4xf32> to vector<4xf64>
   %143 = fpext %vcf32 : vector<4xf32> to vector<4xf64>
 
   // CHECK: = fptrunc {{.*}} : vector<4xf32> to vector<4xf16>
   %144 = fptrunc %vcf32 : vector<4xf32> to vector<4xf16>
 
+  // CHECK: %{{[0-9]+}} = rsqrt %arg1 : f32
+  %145 = 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