[MLIR] Add the sqrt operation to mlir.

[llitchev]

Summary

Add and pipe through the sqrt operation for Standard and LLVM dialects.

[merge_guards_bot]

Unit tests: pass. 62248 tests passed, 0 failed and 816 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[ftynse]

mlir/docs/Dialects/Standard.md

606

f8 isn't a valid type :)

609

This description should go into ODS. We have an open bug to fully autogenerate this doc from there.

[llitchev]

Addressed some CR comments re: comments - fixed a type to match an existing type and moved an operation description to a OSD file.

[llitchev]

mlir/docs/Dialects/Standard.md

606

Changed the type to f32.

609

Moved the description to the op def.

[merge_guards_bot]

Unit tests: pass. 62248 tests passed, 0 failed and 816 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[ftynse]

One more nit, thanks!

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

445

Can we rather take an f64 argument to this function than add another op (constant) irrelevant to the test?

Orthogonally, we should not pattern-match SSA value names (%21). I see that the rest of this file does and that should be fixed independently (patches in that direction appreciated), but let's not introduce any new such matches. See https://mlir.llvm.org/getting_started/TestingGuide/ for more details.

[llitchev]

Addressed CR re: test matching rules and use of a temp f64 const.

[llitchev]

Clean up one more ordinal used for f32 sqrt test.

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

445

Yeah, my thought was to stick with the general file style - the tanh test above uses const and yes, the file uses the ordinal names for the vars...

[merge_guards_bot]

Unit tests: fail. 62247 tests passed, 1 failed and 816 were skipped.

failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_timedmutex_requirements/thread_timedmutex_recursive/try_lock_for.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: pass. 62248 tests passed, 0 failed and 816 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[ftynse]

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

445

I must have missed that when reviewing that commit. Thanks!

[llitchev]

Resolved a merge conflict in LLVMOps.td due to changes coming from master.

[merge_guards_bot]

Unit tests: pass. 62318 tests passed, 0 failed and 838 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[ftynse]

Please LMK if you need somebody to commit this for you

[flaub]

Thanks @ftynse, I'll land this shortly.

mlir/docs/Dialects/Standard.md

 %x = negf %y : tensor<4x?xf8>
 ```
 
 The `negf` operation computes the negation of a given value. 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.
 
+### 'sqrt' operation
+
+Syntax:
+
+```
+operation ::= ssa-id `=` `sqrt` ssa-use `:` type
+```
+
+Examples:
+
+```mlir
+// Scalar square root value.
+%a = sqrt %b : f64
+// SIMD vector element-wise square root value.
+%f = sqrt %g : vector<4xf32>
+// Tensor element-wise square root value.
+%x = sqrt %y : tensor<4x?xf32>
+```
+
 ### 'tanh' operation
 
 Syntax:
 
 ```
 operation ::= ssa-id `=` `tanh` ssa-use `:` type
 ```
 
 Examples:

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 }
 
 def LLVM_ExpOp : LLVM_UnaryIntrinsicOp<"exp">;
 def LLVM_FAbsOp : LLVM_UnaryIntrinsicOp<"fabs">;
 def LLVM_FCeilOp : LLVM_UnaryIntrinsicOp<"ceil">;
 def LLVM_CosOp : LLVM_UnaryIntrinsicOp<"cos">;
 def LLVM_CopySignOp : LLVM_BinarySameArgsIntrinsicOp<"copysign">;
 def LLVM_FMulAddOp : LLVM_TernarySameArgsIntrinsicOp<"fmuladd">;
+def LLVM_SqrtOp : LLVM_UnaryIntrinsicOp<"sqrt">;
 
 def LLVM_LogOp : LLVM_Op<"intr.log", [NoSideEffect]>,
                    Arguments<(ins LLVM_Type:$in)>,
                    Results<(outs LLVM_Type:$res)> {
   let llvmBuilder = [{
     llvm::Module *module = builder.GetInsertBlock()->getModule();
     llvm::Function *fn = llvm::Intrinsic::getDeclaration(
         module, llvm::Intrinsic::log, {$in->getType()});

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

struct Range {
       Value offset, size, stride;
     };
     SmallVector<Range, 8> getRanges();
   }];
 
   let hasCanonicalizer = 1;
 }
 
+def SqrtOp : FloatUnaryOp<"sqrt"> {
+  let summary = "sqrt of the specified value";
+  let description = [{
+    The `sqrt` operation computes 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 TanhOp : FloatUnaryOp<"tanh"> {
   let summary = "hyperbolic tangent of the specified value";
   let description = [{
     The `tanh` operation computes the hyperbolic tangent. 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.
   }];

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

 };
 struct MulIOpLowering : public BinaryOpLLVMOpLowering<MulIOp, LLVM::MulOp> {
   using Super::Super;
 };
 struct SignedDivIOpLowering
     : public BinaryOpLLVMOpLowering<SignedDivIOp, LLVM::SDivOp> {
   using Super::Super;
 };
+struct SqrtOpLowering : public UnaryOpLLVMOpLowering<SqrtOp, LLVM::SqrtOp> {
+  using Super::Super;
+};
 struct UnsignedDivIOpLowering
     : public BinaryOpLLVMOpLowering<UnsignedDivIOp, LLVM::UDivOp> {
   using Super::Super;
 };
 struct SignedRemIOpLowering
     : public BinaryOpLLVMOpLowering<SignedRemIOp, LLVM::SRemOp> {
   using Super::Super;
 };
patterns.insert<
       SelectOpLowering,
       ShiftLeftOpLowering,
       SignExtendIOpLowering,
       SignedDivIOpLowering,
       SignedRemIOpLowering,
       SignedShiftRightOpLowering,
       SplatOpLowering,
       SplatNdOpLowering,
+      SqrtOpLowering,
       SubFOpLowering,
       SubIOpLowering,
       TanhOpLowering,
       TruncateIOpLowering,
       UnsignedDivIOpLowering,
       UnsignedRemIOpLowering,
       UnsignedShiftRightOpLowering,
       XOrOpLowering,

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

 // CHECK-NEXT:  %12 = llvm.ashr %arg2, %arg2 : !llvm<"<4 x i64>">
   %13 = shift_right_signed %arg2, %arg2 : vector<4xi64>
 // CHECK-NEXT:  %13 = llvm.lshr %arg2, %arg2 : !llvm<"<4 x i64>">
   %14 = shift_right_unsigned %arg2, %arg2 : vector<4xi64>
   return %1 : vector<4xf32>
 }
 
 // CHECK-LABEL: @ops
-func @ops(f32, f32, i32, i32) -> (f32, i32) {
-^bb0(%arg0: f32, %arg1: f32, %arg2: i32, %arg3: i32):
+func @ops(f32, f32, i32, i32, f64) -> (f32, i32) {
+^bb0(%arg0: f32, %arg1: f32, %arg2: i32, %arg3: i32, %arg4: f64):
 // CHECK-NEXT:  %0 = llvm.fsub %arg0, %arg1 : !llvm.float
   %0 = subf %arg0, %arg1: f32
 // CHECK-NEXT:  %1 = llvm.sub %arg2, %arg3 : !llvm.i32
   %1 = subi %arg2, %arg3: i32
 // CHECK-NEXT:  %2 = llvm.icmp "slt" %arg2, %1 : !llvm.i32
   %2 = cmpi "slt", %arg2, %1 : i32
 // CHECK-NEXT:  %3 = llvm.sdiv %arg2, %arg3 : !llvm.i32
   %4 = divi_signed %arg2, %arg3 : i32

 // CHECK-NEXT: %16 = llvm.call @tanh(%15) : (!llvm.double) -> !llvm.double
   %17 = std.tanh %16 : f64
 // CHECK-NEXT: %17 = llvm.shl %arg2, %arg3 : !llvm.i32
   %18 = shift_left %arg2, %arg3 : i32
 // CHECK-NEXT: %18 = llvm.ashr %arg2, %arg3 : !llvm.i32
   %19 = shift_right_signed %arg2, %arg3 : i32
 // CHECK-NEXT: %19 = llvm.lshr %arg2, %arg3 : !llvm.i32
   %20 = shift_right_unsigned %arg2, %arg3 : i32
-
+// CHECK-NEXT: %{{[0-9]+}} = "llvm.intr.sqrt"(%arg0) : (!llvm.float) -> !llvm.float
+  %21 = std.sqrt %arg0 : f32
+// CHECK-NEXT: %{{[0-9]+}} = "llvm.intr.sqrt"(%arg4) : (!llvm.double) -> !llvm.double
+  %22 = std.sqrt %arg4 : f64
   return %0, %4 : f32, i32
 }
 
 // Checking conversion of index types to integers using i1, assuming no target
 // system would have a 1-bit address space.  Otherwise, we would have had to
 // make this test dependent on the pointer size on the target system.
 // CHECK-LABEL: @index_cast
 func @index_cast(%arg0: index, %arg1: i1) {

mlir/test/IR/core-ops.mlir

^bb42(%t: tensor<4x4x?xf32>, %f: f32, %i: i32, %idx : index, %j: i64, %half: f16):
   %136 = shift_right_unsigned %idx, %idx : index
 
   // CHECK: %{{[0-9]+}} = shift_right_unsigned %cst_5, %cst_5 : vector<42xi32>
   %137 = shift_right_unsigned %vci32, %vci32 : vector<42 x i32>
 
   // CHECK: %{{[0-9]+}} = shift_right_unsigned %cst_4, %cst_4 : tensor<42xi32>
   %138 = shift_right_unsigned %tci32, %tci32 : tensor<42 x i32>
 
+  // CHECK: %{{[0-9]+}} = sqrt %arg1 : f32
+  %139 = "std.sqrt"(%f) : (f32) -> f32
+
+  // CHECK: %{{[0-9]+}} = sqrt %arg1 : f32
+  %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>
+
   return
 }
 
 // CHECK-LABEL: func @affine_apply() {
 func @affine_apply() {
   %i = "std.constant"() {value = 0: index} : () -> index
   %j = "std.constant"() {value = 1: index} : () -> index
 

mlir/test/Target/llvmir-intrinsics.mlir

 llvm.func @fabs_test(%arg0: !llvm.float, %arg1: !llvm<"<8 x float>">) {
   // CHECK: call float @llvm.fabs.f32
   "llvm.intr.fabs"(%arg0) : (!llvm.float) -> !llvm.float
   // CHECK: call <8 x float> @llvm.fabs.v8f32
   "llvm.intr.fabs"(%arg1) : (!llvm<"<8 x float>">) -> !llvm<"<8 x float>">
   llvm.return
 }
 
+// CHECK-LABEL: @sqrt_test
+llvm.func @sqrt_test(%arg0: !llvm.float, %arg1: !llvm<"<8 x float>">) {
+  // CHECK: call float @llvm.sqrt.f32
+  "llvm.intr.sqrt"(%arg0) : (!llvm.float) -> !llvm.float
+  // CHECK: call <8 x float> @llvm.sqrt.v8f32
+  "llvm.intr.sqrt"(%arg1) : (!llvm<"<8 x float>">) -> !llvm<"<8 x float>">
+  llvm.return
+}
+
 // CHECK-LABEL: @ceil_test
 llvm.func @ceil_test(%arg0: !llvm.float, %arg1: !llvm<"<8 x float>">) {
   // CHECK: call float @llvm.ceil.f32
   "llvm.intr.ceil"(%arg0) : (!llvm.float) -> !llvm.float
   // CHECK: call <8 x float> @llvm.ceil.v8f32
   "llvm.intr.ceil"(%arg1) : (!llvm<"<8 x float>">) -> !llvm<"<8 x float>">
   llvm.return
 }

 // CHECK: declare float @llvm.log.f32(float)
 // CHECK: declare <8 x float> @llvm.log.v8f32(<8 x float>) #0
 // CHECK: declare float @llvm.log10.f32(float)
 // CHECK: declare <8 x float> @llvm.log10.v8f32(<8 x float>) #0
 // CHECK: declare float @llvm.log2.f32(float)
 // CHECK: declare <8 x float> @llvm.log2.v8f32(<8 x float>) #0
 // CHECK: declare float @llvm.fabs.f32(float)
 // CHECK: declare <8 x float> @llvm.fabs.v8f32(<8 x float>) #0
+// CHECK: declare float @llvm.sqrt.f32(float)
+// CHECK: declare <8 x float> @llvm.sqrt.v8f32(<8 x float>) #0
 // CHECK: declare float @llvm.ceil.f32(float)
 // CHECK: declare <8 x float> @llvm.ceil.v8f32(<8 x float>) #0
 // CHECK: declare float @llvm.cos.f32(float)
 // CHECK: declare <8 x float> @llvm.cos.v8f32(<8 x float>) #0
 // CHECK: declare float @llvm.copysign.f32(float, float)