convert com.microsoft to onnx RotaryEmbedding

Author: p-lanza

src/Dialect/ONNX/ONNXOps/NN/RotaryEmbedding.cpp

@@ -90,7 +90,7 @@ LogicalResult ONNXRotaryEmbeddingOp::verify() {
           *this->getOperation(), cosCache, lastIndex, cosCacheShape[lastIndex],
           std::to_string(rotaryEmbeddingDim / 2));
     lastIndex = sinCacheShape.size() - 1;
-    if (sinCacheShape[lastIndex] == rotaryEmbeddingDim / 2)
+    if (sinCacheShape[lastIndex] != rotaryEmbeddingDim / 2)
       return onnx_mlir::Diagnostic::emitDimensionHasUnexpectedValueError(
           *this->getOperation(), sinCache, lastIndex, sinCacheShape[lastIndex],
           std::to_string(rotaryEmbeddingDim / 2));

src/Dialect/ONNX/Transforms/Decompose.cpp

@@ -3458,6 +3458,62 @@ struct MicrosoftGroupQueryAttention : public CustomOpToOnnxOps {
   };
 };
 
+struct MicrosoftRotaryEmbedding : public CustomOpToOnnxOps {
+  MicrosoftRotaryEmbedding(MLIRContext *ctx, PatternBenefit b = 1)
+      : CustomOpToOnnxOps(ctx, MicrosoftDomainName, "RotaryEmbedding", b) {}
+
+  LogicalResult matchAndRewriteImpl(
+      ONNXCustomOp customOp, PatternRewriter &rewriter) const final {
+
+    const Location loc = customOp.getLoc();
+    const int64_t numIn = customOp.getNumOperands();
+    assert((numIn == 4) && "expects 4 inputs");
+    const int64_t numOut = customOp.getNumResults();
+    assert((numOut == 1) && "expects 1 outputs");
+
+    Value input = customOp.getOperand(0);
+    Value position_ids = customOp.getOperand(1);
+    Value cos_cache = customOp.getOperand(2);
+    Value sin_cache = customOp.getOperand(3);
+
+    if (customOp->hasAttrOfType<IntegerAttr>("is_packed_batching") &&
+        customOp->getAttrOfType<IntegerAttr>("is_packed_batching").getSInt() !=
+            0)
+      return rewriter.notifyMatchFailure(
+          customOp, "attribute 'is_packed_batching' not supported by "
+                    "onnx.RotaryEmbedding");
+    if (customOp->hasAttrOfType<IntegerAttr>("scale") &&
+        customOp->getAttrOfType<FloatAttr>("scale").getValueAsDouble() != 1.0f)
+      return rewriter.notifyMatchFailure(
+          customOp, "attribute 'scale' not supported by onnx.RotaryEmbedding");
+
+    auto rotaryEmbedding =
+        rewriter.create<ONNXRotaryEmbeddingOp>(loc, customOp->getResultTypes(),
+            ValueRange{input, cos_cache, sin_cache, position_ids});
+
+    if (customOp->hasAttrOfType<IntegerAttr>("num_heads"))
+      rotaryEmbedding->setAttr(
+          "num_heads", customOp->getAttrOfType<IntegerAttr>("num_heads"));
+
+    if (customOp->hasAttrOfType<IntegerAttr>("interleaved"))
+      rotaryEmbedding->setAttr(
+          "interleaved", customOp->getAttrOfType<IntegerAttr>("interleaved"));
+
+    if (customOp->hasAttrOfType<IntegerAttr>("rotary_embedding_dim"))
+      rotaryEmbedding->setAttr("rotary_embedding_dim",
+          customOp->getAttrOfType<IntegerAttr>("rotary_embedding_dim"));
+
+    if (failed(verifyOpsErasingOnError({rotaryEmbedding}, rewriter))) {
+      return rewriter.notifyMatchFailure(
+          customOp, "Decomposition failed verification");
+    }
+
+    rewriter.replaceOp(customOp, rotaryEmbedding);
+
+    return success();
+  };
+};
+
 template <typename OpToCreate>
 struct CustomOpMicrosoftToSingleOnnxOp : public CustomOpToOnnxOps {
   CustomOpMicrosoftToSingleOnnxOp(MLIRContext *context,
@@ -3946,6 +4002,7 @@ void onnx_mlir::getDecomposeONNXToONNXPatterns(
   patterns.insert<SimplifiedLayerNorm>(context);
   patterns.insert<MicrosoftSkipSimplifiedLayerNorm>(context);
   patterns.insert<MicrosoftGroupQueryAttention>(context);
+  patterns.insert<MicrosoftRotaryEmbedding>(context);
   patterns.insert<DecomposeSlicePadPattern>(context);
   patterns.insert<DecomposeScatterNDPattern>(context);
   patterns.insert<SoftmaxCrossEntropyPattern>(context);

test/mlir/onnx/onnx_decompose_customop.mlir

@@ -866,3 +866,48 @@ func.func @gqa_with_scale_softcap_and_qk_output_2(
 // CHECK-SAME:          : (tensor<1x128x3072xf32>, tensor<1x128x1536xf32>, tensor<1x128x1536xf32>, none, tensor<1x16x256x96xf32>, tensor<1x16x256x96xf32>) -> (tensor<1x128x3072xf32>, tensor<1x16x384x96xf32>, tensor<1x16x384x96xf32>, tensor<1x32x128x256xf32>)
 // CHECK:           return %[[VAL_7]], %[[VAL_8]], %[[VAL_9]], %[[VAL_10]] : tensor<1x128x3072xf32>, tensor<1x16x384x96xf32>, tensor<1x16x384x96xf32>, tensor<1x32x128x256xf32>
 // CHECK:         }
+
+// -----
+
+func.func @rotary_embedding_4d_interleaved_rotdim_16(%data: tensor<1x32x128x96xf32>, %pos_ids: tensor<1x128xi64>, %cos_cache: tensor<4096x8xf32>, %sin_cache: tensor<4096x8xf32>) -> tensor<1x32x128x96xf32> {
+  %0 = "onnx.Custom"(%data, %pos_ids, %cos_cache, %sin_cache) {
+    domain_name = "com.microsoft",
+    function_name = "RotaryEmbedding",
+    interleaved = 1 : si64,
+    rotary_embedding_dim = 16 : si64
+  }: (tensor<1x32x128x96xf32>, tensor<1x128xi64>, tensor<4096x8xf32>, tensor<4096x8xf32>) -> tensor<1x32x128x96xf32>
+  return %0 : tensor<1x32x128x96xf32>
+}
+
+// CHECK-LABEL:   func.func @rotary_embedding_4d_interleaved_rotdim_16(
+// CHECK-SAME:                                                     %[[VAL_0:.*]]: tensor<1x32x128x96xf32>,
+// CHECK-SAME:                                                     %[[VAL_1:.*]]: tensor<1x128xi64>,
+// CHECK-SAME:                                                     %[[VAL_2:.*]]: tensor<4096x8xf32>,
+// CHECK-SAME:                                                     %[[VAL_3:.*]]: tensor<4096x8xf32>) -> tensor<1x32x128x96xf32> {
+// CHECK:           %[[VAL_4:.*]] = "onnx.RotaryEmbedding"(%[[VAL_0]], %[[VAL_2]], %[[VAL_3]], %[[VAL_1]]) 
+// CHECK-SAME:          {interleaved = 1 : si64, rotary_embedding_dim = 16 : si64} 
+// CHECK-SAME:          : (tensor<1x32x128x96xf32>, tensor<4096x8xf32>, tensor<4096x8xf32>, tensor<1x128xi64>) -> tensor<1x32x128x96xf32>
+// CHECK:           return %[[VAL_4]] : tensor<1x32x128x96xf32>
+// CHECK:         }
+
+// -----
+
+func.func @test_rotary_embedding_3d(%data: tensor<1x128x3072xf32>, %pos_ids: tensor<1x128xi64>, %cos_cache: tensor<4096x48xf32>, %sin_cache: tensor<4096x48xf32>) -> tensor<1x128x3072xf32> {
+  %0 = "onnx.Custom"(%data, %pos_ids, %cos_cache, %sin_cache) {
+    domain_name = "com.microsoft",
+    function_name = "RotaryEmbedding",
+    num_heads = 32: si64
+  } : (tensor<1x128x3072xf32>, tensor<1x128xi64>, tensor<4096x48xf32>, tensor<4096x48xf32>) -> tensor<1x128x3072xf32>
+  return %0 : tensor<1x128x3072xf32>
+}
+
+// CHECK-LABEL:   func.func @test_rotary_embedding_3d(
+// CHECK-SAME:                                        %[[VAL_0:.*]]: tensor<1x128x3072xf32>,
+// CHECK-SAME:                                        %[[VAL_1:.*]]: tensor<1x128xi64>,
+// CHECK-SAME:                                        %[[VAL_2:.*]]: tensor<4096x48xf32>,
+// CHECK-SAME:                                        %[[VAL_3:.*]]: tensor<4096x48xf32>) -> tensor<1x128x3072xf32> {
+// CHECK:           %[[VAL_4:.*]] = "onnx.RotaryEmbedding"(%[[VAL_0]], %[[VAL_2]], %[[VAL_3]], %[[VAL_1]]) 
+// CHECK-SAME:          {interleaved = 0 : si64, num_heads = 32 : si64, rotary_embedding_dim = 0 : si64} 
+// CHECK-SAME:          : (tensor<1x128x3072xf32>, tensor<4096x48xf32>, tensor<4096x48xf32>, tensor<1x128xi64>) -> tensor<1x128x3072xf32>
+// CHECK:           return %[[VAL_4]] : tensor<1x128x3072xf32>
+// CHECK:         }