lowerd Sqrt to linalg

reused clang-format, as changes got deleted

e2e_testing/torchscript/elementwise.py

@@ -360,3 +360,21 @@ class ElementwiseLogModule(torch.nn.Module):
 @register_test_case(module_factory=lambda: ElementwiseLogModule())
 def ElementwiseLogModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 4))
+
+
+class ElementwiseSqrtModule(torch.nn.Module):
+    def __init__(self): 
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.float32, True),
+    ])
+
+    def forward(self, a):
+        return torch.sqrt(a)
+
+@register_test_case(module_factory=lambda: ElementwiseSqrtModule())
+def ElementwiseSqrtModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 4))

include/torch-mlir/Dialect/Torch/IR/GeneratedAtenOps.td

@@ -2700,3 +2700,17 @@ def Torch_AtenEqDeviceOp : Torch_Op<"aten.eq.device", [
   let assemblyFormat = "$a `,` $b attr-dict `:` type($a) `,` type($b) `->` type($result)";
 }
 
+def Torch_AtenSqrtOp : Torch_Op<"aten.sqrt", [
+    AllowsTypeRefinement,
+    HasValueSemantics
+  ]> {
+  let summary = "Generated op for `aten::sqrt : (Tensor) -> (Tensor)`";
+  let arguments = (ins 
+    AnyTorchTensorType:$a
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$a `,` attr-dict `:` type($a) `->` type($result)";
+}
+

lib/Conversion/TorchToLinalg/TorchToLinalg.cpp

@@ -1280,6 +1280,8 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
     return b.create<math::ExpOp>(loc, payloadArgs[0]);
   if (isa<AtenLogOp>(op))
     return b.create<math::LogOp>(loc, payloadArgs[0]);
+  if (isa<AtenSqrtOp>(op))
+    return b.create<math::SqrtOp>(loc, payloadArgs[0]);
   if (isa<AtenSigmoidOp>(op)) {
     Type elementType = payloadArgs[0].getType();
     auto one = b.create<arith::ConstantOp>(loc, FloatAttr::get(elementType, 1));
@@ -1663,7 +1665,8 @@ struct ConvertElementwiseOp : ConversionPattern {
     if (!isa<AtenTanhOp, AtenReluOp, AtenGeluOp, AtenAddTensorOp,
              AtenMulTensorOp, AtenDivTensorOp, AtenSubTensorOp,
              AtenLerpTensorOp, AtenSigmoidOp, AtenExpOp, AtenMinimumOp,
-             AtenMaximumOp, AtenClampOp, AtenRsubScalarOp, AtenLogOp>(op))
+             AtenMaximumOp, AtenClampOp, AtenRsubScalarOp, AtenLogOp,
+             AtenSqrtOp>(op))
       return rewriter.notifyMatchFailure(op, "not a supported elementwise op");
 
     if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
@@ -2799,7 +2802,8 @@ public:
     target.addIllegalOp<AtenTanhOp, AtenReluOp, AtenGeluOp, AtenAddTensorOp,
                         AtenMulTensorOp, AtenDivTensorOp, AtenSubTensorOp,
                         AtenLerpTensorOp, AtenSigmoidOp, AtenMinimumOp,
-                        AtenMaximumOp, AtenClampOp, AtenRsubScalarOp, AtenLogOp>();
+                        AtenMaximumOp, AtenClampOp, AtenRsubScalarOp, AtenLogOp,
+                        AtenSqrtOp>();
     patterns.add<ConvertElementwiseOp>(typeConverter, context);
     target.addIllegalOp<AtenUnsqueezeOp>();
     patterns.add<ConvertAtenUnsqueezeOp>(typeConverter, context);

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -230,7 +230,8 @@ public:
             DerefineOp, AtenToPrimDeviceOp, AtenCpuOp, AtenContiguousOp,
             AtenFill_ScalarOp, AtenDetachOp, AtenMaskedFill_ScalarOp,
             AtenCopy_Op, AtenIndexPut_Op, AtenCopy_Op, AtenCumsumOp,
-            AtenLayerNormOp, AtenClampOp, AtenRsubScalarOp, AtenLogOp>(op)) {
+            AtenLayerNormOp, AtenClampOp, AtenRsubScalarOp, AtenLogOp,
+            AtenSqrtOp>(op)) {
       return getLatticeElement(op->getResult(0)).join(*operands[0]);
     }
 

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -508,6 +508,7 @@ def emit_aten_ops(torch_ir_dir: str, registry: Registry):
         emit("aten::logsumexp : (Tensor, int[], bool) -> (Tensor)")
         emit("aten::mean.dim : (Tensor, int[], bool, int?) -> (Tensor)")
         emit("aten::__and__.Tensor : (Tensor, Tensor) -> (Tensor)")
+        emit("aten::sqrt : (Tensor) -> (Tensor)")
 
         # Misc tensor ops.
         emit("aten::unsqueeze : (Tensor, int) -> (Tensor)")