[torch] Add Canonicalize Pattern for embedding op

Converts PrimConvertOp followed by Embedding -> Embedding followed by
PrimConvertOp. We don't need to cast the entire matrix; just the output
of the embedding op.

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

@@ -10106,6 +10106,7 @@ def Torch_AtenEmbeddingOp : Torch_Op<"aten.embedding", [
       printDefaultTorchOp(printer, *this, 5, 1);
     }
   }];
+  let hasCanonicalizer = 1;
 }
 
 def Torch_AtenEmbeddingBagPaddingIdxOp : Torch_Op<"aten.embedding_bag.padding_idx", [

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -547,6 +547,46 @@ void PrimIfOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
     return success();
   });
 }
+//===----------------------------------------------------------------------===//
+// AtenEmbeddingOp
+//===----------------------------------------------------------------------===//
+//
+// Converts PrimConvertElementTypeOp followed by AtenEmbeddingOp to
+// AtenEmbeddingOp followed by PrimConvertElementTypeOp.
+void AtenEmbeddingOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                                  MLIRContext *context) {
+  patterns.add(+[](AtenEmbeddingOp op, PatternRewriter &rewriter) {
+    //
+    auto checkWdthAGtWdthB = [](BaseTensorType aType, BaseTensorType bType) {
+      if (!(aType.hasDtype() && bType.hasDtype()))
+        return false;
+      auto aWidth = aType.getDtype().getIntOrFloatBitWidth();
+      auto bWidth = bType.getDtype().getIntOrFloatBitWidth();
+      return aWidth > bWidth;
+    };
+
+    auto convertOp =
+        op.getWeight().getDefiningOp<Torch::PrimsConvertElementTypeOp>();
+    if (!convertOp)
+      return failure();
+    Value getWeight = convertOp.getA();
+    auto weightType = getWeight.getType().cast<BaseTensorType>();
+    auto opType = op.getType().cast<BaseTensorType>();
+
+    if (!checkWdthAGtWdthB(opType, weightType))
+      return failure();
+
+    auto updateType =
+        opType.getWithSizesAndDtype(opType.getSizes(), weightType.getDtype());
+    Value newEmbedding = rewriter.create<AtenEmbeddingOp>(
+        op.getLoc(), updateType, getWeight, op.getIndices(), op.getPaddingIdx(),
+        op.getScaleGradByFreq(), op.getSparse());
+    Value updateRes = rewriter.create<Torch::PrimsConvertElementTypeOp>(
+        op.getLoc(), op.getType(), newEmbedding, convertOp.getDtype());
+    rewriter.replaceOp(op, updateRes);
+    return success();
+  });
+}
 
 //===----------------------------------------------------------------------===//
 // AtenDotOp

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py

@@ -777,7 +777,10 @@ def emit_with_mutating_variants(key, **kwargs):
     emit("aten::detach : (Tensor) -> (Tensor)", has_folder=True)
     emit("aten::device.with_index : (str, int) -> (Device)", has_canonicalizer=True)
     emit("aten::cuda : (Tensor) -> (Tensor)", has_canonicalizer=True)
-    emit("aten::embedding : (Tensor, Tensor, int, bool, bool) -> (Tensor)")
+    emit(
+        "aten::embedding : (Tensor, Tensor, int, bool, bool) -> (Tensor)",
+        has_canonicalizer=True,
+    )
     emit(
         "aten::embedding_bag.padding_idx : (Tensor, Tensor, Tensor, bool, int, bool, Tensor?, bool, int?) -> (Tensor, Tensor, Tensor, Tensor)"
     )

test/Dialect/Torch/canonicalize.mlir

@@ -3026,3 +3026,40 @@ func.func @torch.aten.clone$no_fold(%arg0: !torch.vtensor<[1,2,50,4],f32>) -> (!
   %1 = torch.copy.to_tensor %0 : !torch.tensor
   return %1 : !torch.tensor
 }
+
+// -----
+
+// CHECK-LABEL:   func.func @torch.prims.convert_element_type$fold(
+// CHECK-SAME:                                                     %[[ARG0:.*]]: !torch.vtensor<[64,128],f16>,
+// CHECK-SAME:                                                     %[[ARG1:.*]]: !torch.vtensor<[4],si64>) -> !torch.vtensor<[4,128],f32> {
+// CHECK:           %[[INT32:.*]] = torch.constant.int 6
+// CHECK:           %[[FALSE:.*]] = torch.constant.bool false
+// CHECK:           %[[NEGONE:.*]] = torch.constant.int -1
+// CHECK:           %[[EMBD:.*]] = torch.aten.embedding %[[ARG0]], %[[ARG1]], %[[NEGONE]], %[[FALSE]], %[[FALSE]] : !torch.vtensor<[64,128],f16>, !torch.vtensor<[4],si64>, !torch.int, !torch.bool, !torch.bool -> !torch.vtensor<[4,128],f16>
+// CHECK:           %[[DTYPE:.*]] = torch.prims.convert_element_type %[[EMBD]], %[[INT32]] : !torch.vtensor<[4,128],f16>, !torch.int -> !torch.vtensor<[4,128],f32>
+// CHECK:           return %[[DTYPE]] : !torch.vtensor<[4,128],f32>
+// CHECK:         }
+func.func @torch.prims.convert_element_type$fold(%arg0: !torch.vtensor<[64,128],f16>, %arg1: !torch.vtensor<[4],si64>) -> !torch.vtensor<[4,128],f32> {
+  %int6 = torch.constant.int 6
+  %false = torch.constant.bool false
+  %int-1 = torch.constant.int -1
+  %0 = torch.prims.convert_element_type %arg0, %int6 : !torch.vtensor<[64,128],f16>, !torch.int -> !torch.vtensor<[64,128],f32>
+  %1 = torch.aten.embedding %0, %arg1, %int-1, %false, %false : !torch.vtensor<[64,128],f32>, !torch.vtensor<[4],si64>, !torch.int, !torch.bool, !torch.bool -> !torch.vtensor<[4,128],f32>
+  return %1 : !torch.vtensor<[4,128],f32>
+}
+
+// -----
+
+// CHECK-LABEL:   func.func @torch.prims.convert_element_type$not_fold(
+// CHECK-SAME:                                                     %[[ARG0:.*]]: !torch.vtensor<[64,128],f32>,
+// CHECK-SAME:                                                     %[[ARG1:.*]]: !torch.vtensor<[4],si64>) -> !torch.vtensor<[4,128],f16> {
+// CHECK:           %{{.*}} = torch.prims.convert_element_type
+// CHECK:           %{{.*}} = torch.aten.embedding
+func.func @torch.prims.convert_element_type$not_fold(%arg0: !torch.vtensor<[64,128],f32>, %arg1: !torch.vtensor<[4],si64>) -> !torch.vtensor<[4,128],f16> {
+  %int5 = torch.constant.int 5
+  %false = torch.constant.bool false
+  %int-1 = torch.constant.int -1
+  %0 = torch.prims.convert_element_type %arg0, %int5 : !torch.vtensor<[64,128],f32>, !torch.int -> !torch.vtensor<[64,128],f16>
+  %1 = torch.aten.embedding %0, %arg1, %int-1, %false, %false : !torch.vtensor<[64,128],f16>, !torch.vtensor<[4],si64>, !torch.int, !torch.bool, !torch.bool -> !torch.vtensor<[4,128],f16>
+  return %1 : !torch.vtensor<[4,128],f16>
+}