[ONNX][TORCH] Add Onnx->Linalg lowering for RotaryEmbedding Op (#4002)

This commit adds the Onnx->Linalg lowering for Onnx's RotaryEmbedding op
(ref:
https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#commicrosoftrotaryembedding)
by registering a customized torch op named
`OnnxVariantAtenRotaryEmbeddingOp`. This is done so that the Onnx's
RotaryEmbedding op can be lowered to this op and this op can be lowered
from Torch->Linalg.

The lowering has been adopted from the OnnxRuntime. Files for
references:
1.)
https://github.com/microsoft/onnxruntime/blob/e1e3f623f61816008e79dddc91a51ffe7f0ff5cf/onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc#L47-L93
2.)
https://github.com/microsoft/onnxruntime/blob/94c69f55d480cb4a8dcbc161d29ef3acca9392a7/onnxruntime/contrib_ops/cpu/bert/rotary_embedding_helper.h

---------

Signed-off-by: Vivek Khandelwal <[email protected]>
Co-authored-by: zjgarvey <[email protected]>

Author: vivekkhandelwal1

include/torch-mlir/Conversion/TorchOnnxToTorch/Patterns.h

@@ -472,6 +472,7 @@ class OnnxCustomOpConversionPattern
 
 // Patterns are split into chunks to speed compile time and reduce some
 // contention on the same source files.
+void populateComMicrosoftDomain(OnnxCustomOpConversionPattern &patterns);
 void populateDefaultDomainAtoF(OnnxCustomOpConversionPattern &patterns);
 void populateDefaultDomainGtoP(OnnxCustomOpConversionPattern &patterns);
 void populateDefaultDomainQtoZ(OnnxCustomOpConversionPattern &patterns);

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

@@ -1410,4 +1410,36 @@ def Torch_BindSymbolicShapeOp : Torch_Op<"bind_symbolic_shape", []> {
   let hasVerifier = 1;
 }
 
+// This op is corresponding to the Onnx's RotaryEmbedding operator.
+// Ref: https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#commicrosoftrotaryembedding
+def Torch_OnnxVariantRotaryEmbeddingOp: Torch_Op<"onnx.rotary_embedding", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "`rotary_embedding op : (Tensor, Tensor, Tensor, Tensor, int, int, int, int, float) -> (Tensor)`";
+  let description = [{
+    The `torch.onnx.rotary_embedding` operation is an op which is used
+    specifically for supporting the Onnx's Rotary Embedding op. The
+    reason for this is that the Onnx ops can't be directly lowered to
+    Linalg and we have to map them to a legal Torch Dialect op, hence
+    this op is used for that purpose.
+  }];
+  let arguments = (ins
+    AnyTorchTensorType:$input,
+    AnyTorchTensorType:$position_ids,
+    AnyTorchTensorType:$cos_cache,
+    AnyTorchTensorType:$sin_cache,
+    Torch_IntType:$interleaved,
+    Torch_IntType:$is_packed_batching,
+    Torch_IntType:$num_heads,
+    Torch_IntType:$rotary_embedding_dim,
+    Torch_FloatType:$scale
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+}
+
 #endif // TORCH_OPS

lib/Conversion/TorchOnnxToTorch/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_mlir_conversion_library(TorchMLIRTorchOnnxToTorch
+    ComMicrosoftDomain.cpp
     DefaultDomainAtoF.cpp
     DefaultDomainGtoP.cpp
     DefaultDomainQtoZ.cpp

lib/Conversion/TorchOnnxToTorch/ComMicrosoftDomain.cpp

@@ -0,0 +1,66 @@
+//===------------------------------------------------------------*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// Also available under a BSD-style license. See LICENSE.
+//
+//===----------------------------------------------------------------------===//
+
+#include "torch-mlir/Conversion/TorchOnnxToTorch/Patterns.h"
+#include "torch-mlir/Conversion/TorchOnnxToTorch/Utils.h"
+#include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
+#include "torch-mlir/Dialect/Torch/Utils/Utils.h"
+#include <numeric>
+
+using namespace mlir;
+using namespace mlir::torch;
+using namespace mlir::torch::onnx_c;
+
+void mlir::torch::onnx_c::populateComMicrosoftDomain(
+    OnnxCustomOpConversionPattern &patterns) {
+  patterns.onOp(
+      "RotaryEmbedding", 1,
+      [](OpBinder binder, ConversionPatternRewriter &rewriter) {
+        Location loc = binder.getLoc();
+        int64_t interleaved, isPackedBatching, numHeads, rotaryEmbeddingDim;
+        float scale;
+        Value input, positionIds, cosCache, sinCache;
+        if (binder.tensorOperandAtIndex(input, 0) ||
+            binder.tensorOperandAtIndex(positionIds, 1) ||
+            binder.tensorOperandAtIndex(cosCache, 2) ||
+            binder.tensorOperandAtIndex(sinCache, 3) ||
+            binder.s64IntegerAttr(interleaved, "interleaved", 0) ||
+            binder.s64IntegerAttr(isPackedBatching, "is_packed_batching", 0) ||
+            binder.s64IntegerAttr(numHeads, "num_heads", 0) ||
+            binder.s64IntegerAttr(rotaryEmbeddingDim, "rotary_embedding_dim",
+                                  0) ||
+            binder.f32FloatAttr(scale, "scale", 1.0)) {
+          return rewriter.notifyMatchFailure(binder.op,
+                                             "Failed to get required inputs");
+        }
+
+        Torch::ValueTensorType resultType;
+        if (binder.tensorResultType(resultType)) {
+          return rewriter.notifyMatchFailure(binder.op,
+                                             "result type bind failure");
+        }
+
+        Value cstInterleaved = rewriter.create<Torch::ConstantIntOp>(
+            loc, rewriter.getI64IntegerAttr(interleaved));
+        Value cstIsPackedBatching = rewriter.create<Torch::ConstantIntOp>(
+            loc, rewriter.getI64IntegerAttr(isPackedBatching));
+        Value cstNumHeads = rewriter.create<Torch::ConstantIntOp>(
+            loc, rewriter.getI64IntegerAttr(numHeads));
+        Value cstRotaryEmbeddingDim = rewriter.create<Torch::ConstantIntOp>(
+            loc, rewriter.getI64IntegerAttr(rotaryEmbeddingDim));
+        Value cstScale = rewriter.create<Torch::ConstantFloatOp>(
+            loc, rewriter.getF64FloatAttr(scale));
+
+        rewriter.replaceOpWithNewOp<Torch::OnnxVariantRotaryEmbeddingOp>(
+            binder.op, resultType, input, positionIds, cosCache, sinCache,
+            cstInterleaved, cstIsPackedBatching, cstNumHeads,
+            cstRotaryEmbeddingDim, cstScale);
+        return success();
+      });
+}

lib/Conversion/TorchOnnxToTorch/TorchOnnxToTorch.cpp

@@ -56,6 +56,7 @@ class ConvertTorchOnnxToTorch
         std::make_unique<OnnxCustomOpConversionPattern>(
             context, "onnx.",
             /*domainVersion=*/defaultOpsetVersion);
+    populateComMicrosoftDomain(*defaultDomainPatterns);
     populateDefaultDomainAtoF(*defaultDomainPatterns);
     populateDefaultDomainGtoP(*defaultDomainPatterns);
     populateDefaultDomainQtoZ(*defaultDomainPatterns);