feat(elliptic_curve): lower MSMOp

Author: ashjeong

tests/Dialect/EllipticCurve/elliptic_curve_to_field.mlir

@@ -192,3 +192,14 @@ func.func @test_point_set() {
   %doubled = elliptic_curve.double %point1 : !affine -> !jacobian
   return
 }
+
+func.func @test_msm() {
+  %var1 = field.pf.constant 1 : !PF
+  %var5 = field.pf.constant 5 : !PF
+
+  %scalars = tensor.from_elements %var1, %var5, %var5 : tensor<3x!PF>
+  %affine1 = elliptic_curve.point %var1, %var5 : !PF -> !affine
+  %points = elliptic_curve.point_set.from_elements %affine1, %affine1, %affine1 : tensor<3x!affine>
+  %msm_result = elliptic_curve.msm %scalars, %points : tensor<3x!PF>, tensor<3x!affine> -> !jacobian
+  return
+}

zkir/Dialect/EllipticCurve/Conversions/EllipticCurveToField/EllipticCurveToField.cpp

@@ -613,6 +613,74 @@ struct ConvertScalarMul : public OpConversionPattern<ScalarMulOp> {
   }
 };
 
+struct ConvertMSM : public OpConversionPattern<MSMOp> {
+  explicit ConvertMSM(mlir::MLIRContext *context)
+      : OpConversionPattern<MSMOp>(context) {}
+
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      MSMOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+
+    // 2D tensor of prime field values, e.g.:
+    // |x1, y1, z1|
+    // |x2, y2, z2|
+    Value loweredPointSet = adaptor.getPoints();
+    // 1d tensor of PF, e.g.:
+    // | s1 , s2 |
+    Value scalars = op.getScalars();
+    RankedTensorType loweredPointSetType =
+        cast<RankedTensorType>(loweredPointSet.getType());
+    field::PrimeFieldType baseFieldType =
+        cast<field::PrimeFieldType>(loweredPointSetType.getElementType());
+    unsigned numScalarMuls = loweredPointSetType.getShape()[0];
+    unsigned numCoords = loweredPointSetType.getShape()[1];
+
+    Type inputPointType =
+        cast<RankedTensorType>(op.getPoints().getType()).getElementType();
+    Type outputPointType = op.getOutput().getType();
+    RankedTensorType loweredOutputPointType =
+        RankedTensorType::get({numCoords}, baseFieldType);
+
+    Value accumulator;
+    auto zero = b.create<arith::ConstantIndexOp>(0);
+    auto one = b.create<arith::ConstantIndexOp>(1);
+    auto sz = b.create<arith::ConstantIndexOp>(numCoords);
+    SmallVector<Value> sizes{one, sz};
+    SmallVector<Value> strides{one, one};
+    for (size_t i = 0; i < numScalarMuls; ++i) {
+      auto idx = b.create<arith::ConstantIndexOp>(i);
+
+      // scalar
+      auto scalar = b.create<tensor::ExtractOp>(scalars, ValueRange{idx});
+
+      // point
+      //  - extract point = tensor<2x3x!PF> -> tensor<1x3x!PF>
+      SmallVector<Value> offsets{idx, zero};
+      auto higherRankedPoint = b.create<tensor::ExtractSliceOp>(
+          loweredPointSet, offsets, sizes, strides);
+      // - reshape point = tensor<1x3x!PF> -> tensor<3x!PF>
+      SmallVector<Value> _outputShape{sz};
+      auto outputShape = b.create<tensor::FromElementsOp>(_outputShape);
+      auto point = b.create<tensor::ReshapeOp>(loweredOutputPointType,
+                                               higherRankedPoint, outputShape);
+
+      Value adder = convertScalarMulImpl(point, scalar, inputPointType,
+                                         outputPointType, b);
+      if (i != 0) {
+        accumulator = convertAddImpl(accumulator, adder, outputPointType,
+                                     outputPointType, outputPointType, b);
+      } else {
+        accumulator = adder;
+      }
+    }
+    rewriter.replaceOp(op, accumulator);
+    return success();
+  }
+};
+
 namespace rewrites {
 // In an inner namespace to avoid conflicts with canonicalization patterns
 #include "zkir/Dialect/EllipticCurve/Conversions/EllipticCurveToField/EllipticCurveToField.cpp.inc"
@@ -636,10 +704,11 @@ void EllipticCurveToField::runOnOperation() {
 
   RewritePatternSet patterns(context);
   rewrites::populateWithGenerated(patterns);
-  patterns.add<ConvertPoint, ConvertPointSet, ConvertPointSetExtract,
-               ConvertExtract, ConvertConvertPointType, ConvertAdd,
-               ConvertDouble, ConvertNegate, ConvertSub, ConvertScalarMul>(
-      typeConverter, context);
+  patterns
+      .add<ConvertPoint, ConvertPointSet, ConvertPointSetExtract,
+           ConvertExtract, ConvertConvertPointType, ConvertAdd, ConvertDouble,
+           ConvertNegate, ConvertSub, ConvertScalarMul, ConvertMSM>(
+          typeConverter, context);
 
   addStructuralConversionPatterns(typeConverter, patterns, target);