refactor(metal_mont_backend): adapted from https://github.com/geometryxyz/msl-secp256k1 #13

moven0831 · 2024-11-08T02:33:29Z

Overview

This PR introduces a well-organized structure for further MSM implementation with Metal and adapts a more efficient EC metal library from Zprize 2023 winner, Tal and Koh's work of MSM with WebGPU.

Related Issue

Implement Optimized Montgomery Multiplication for Metal #6

Key Changes

Core Implementation

Refactored Metal MSM Module
- Added new metal_msm module with host, shader, and utils submodules
  - The metal_msm aims to replace the previous metal directory as a refactored version. The metal directory will be deleted once the metal_msm is implemented.
- Implemented basic Metal state management and GPU operations
Better EC backend adapted from Zprize 2023 winner
- BigInt Operations
  - New BigInt and BigIntWide structs with arithmetic operations
  - New Metal shaders for BigInt addition and subtraction
- Field Operations
  - New field arithmetic operations (addition, subtraction)
  - Implemented BN254 field operations with updated modulus
  - Added Montgomery multiplication implementations (CIOS, modified, optimized)
- Elliptic Curve Operations
  - Implemented ECPoint class with Jacobian coordinate operations
  - Added functions for point addition and doubling

Testing & Benchmarking

Test Suites
- BigInt operations tests
- Field arithmetic tests
- Montgomery multiplication tests
- Curve operation tests
Benchmarking
- Added benchmarking modules for Montgomery multiplication
- Implemented performance testing for different multiplication variants

For the BN254 curve we're targeting, the modified mont_mul with a 15-bit limb size appears to work best. This makes sense intuitively: with a 15-bit limb size, there are 17 limbs, resulting in 15 * 17 = 255 bits, which is closest to the 254-bit base field modulus of BN254.

Benchmark result:

running 1 test
=== benchmarking mont_mul_modified ===
benchmark for 11-bit limbs took 306ms
benchmark for 12-bit limbs took 174ms
benchmark for 13-bit limbs took 139ms
benchmark for 14-bit limbs took 144ms
benchmark for 15-bit limbs took 128ms
benchmark for 16-bit limbs: Benchmark failed: results do not match expected values

=== benchmarking mont_mul_optimised ===
benchmark for 11-bit limbs took 281ms
benchmark for 12-bit limbs took 172ms
benchmark for 13-bit limbs took 137ms
benchmark for 14-bit limbs took 137ms
benchmark for 15-bit limbs: Benchmark failed: results do not match expected values
benchmark for 16-bit limbs: Benchmark failed: results do not match expected values

=== benchmarking mont_mul_cios ===
benchmark for 11-bit limbs: Benchmark failed: results do not match expected values
benchmark for 12-bit limbs: Benchmark failed: results do not match expected values
benchmark for 13-bit limbs took 368ms
benchmark for 14-bit limbs took 354ms
benchmark for 15-bit limbs took 259ms
benchmark for 16-bit limbs took 129ms
test msm::metal_msm::tests::mont_backend::mont_benchmarks::all_benchmarks ... ok

Utility Functions

Added limb conversion utilities from BigInt to arbitrary limb size format in Vec
Implemented Montgomery parameter calculations

Dependencies

Updated various dependencies including:

num: 0.4.0
num-bigint: 0.4.3

Next Steps

Review and merge into main branch
Build new MSM algorithms with metal on top of this ec backend
Consider additional optimizations based on benchmark results

Reference

…-secp256k1

…is is completed

…each cases

…sl-secp256k1

…tation

…ryxyz/msl-secp256k1

…://github.com/geometers/multiprecision

…ec<u32>

…rbitrary limbs of Vec<u32>

…eometryxyz/msl-secp256k1

yaroslavyaroslav

LGTM

Haven't reviewed any metal file, bc I believe they're just moved rather modified, so I consider them out of scope this particular PR.

These build.rs although being nice to have not a critical part of PR, so in case of tight deadlines if any we can put in a separate PR, I can even take it on my own afterwards.

yaroslavyaroslav · 2024-11-08T23:07:12Z

mopro-msm/src/msm/metal/abstraction/limbs_conversion.rs

+        let mut limb_idx = 0;
+
+        for &byte in bytes.iter() {
+            val |= (byte as u32) << bits;


I'd consider to operate by a whole word here rather than a single byte. However I think this optimisation should be a part of a separate PR.

yaroslavyaroslav · 2024-11-09T13:36:42Z

mopro-msm/src/msm/metal_msm/host/errors.rs

+pub enum MetalError {
+    #[error("Couldn't find a system default device for Metal")]
+    DeviceNotFound(),
+    #[error("Couldn't create a new Metal library: {0}")]


Seems like this one is used within metalcompiller.rs, so I believe it could be removed due to reasoning we have discussed in DM

yaroslavyaroslav · 2024-11-09T13:50:26Z

mopro-msm/src/msm/metal_msm/host/shader.rs

+use std::string::String;
+
+pub fn compile_metal(path_from_cargo_manifest_dir: &str, input_filename: &str) -> String {
+    let input_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"))


Just saving here all the main points of the talk we have had in DM recently.

The metal building code is pretty similar to what I've made for my internal use in #14 so it's perfectly fine to keep it as it is except for two points:

Move the code in build.rs as will enable it building during the library itself build stage, allowing cargo to handle incremental builds of metal shaders effortlessly.

It worth it to keep built library in target directory with env::var("OUT_DIR") env variable, which could be referenced in the main code base as env!("OUT_DIR")

It's worth it to conditionally append debug symbols generation attribute for profiling as it's made here: https://github.com/zkmopro/gpu-acceleration/pull/14/files#diff-1fb60ca492a52fb05b37909bd795664d0a6c31e77d34f28bc1baa3b7d91093a7R35-R40

When it comes to build.rs you have to specify the list of foreign language files to follow, which is done here https://github.com/zkmopro/gpu-acceleration/pull/14/files#diff-1fb60ca492a52fb05b37909bd795664d0a6c31e77d34f28bc1baa3b7d91093a7R103-R106

moven0831 added 21 commits November 8, 2024 08:53

feat: import mont_mul backend from https://github.com/geometryxyz/msl…

c2c42bb

…-secp256k1

chore: add refactored metal_msm, will remove the previous one once th…

318189b

…is is completed

chore: migrate prev utils mod to refactored metal msm

c45993b

feat: conversion between bigint and arbitrary limb size

100badd

lint

777eafe

chore: ignore all metal ir and lib

161c0ac

test(bigint): adapt from https://github.com/geometryxyz/msl-secp256k1

94849f4

chore: ignore all constants file since it's been automatically generated

d41be85

refactor: add overflow detection and correct suitable bigint val for …

bc90c55

…each cases

test(field): adapt the ff tests from https://github.com/geometryxyz/m…

e9600ad

…sl-secp256k1

test(field): add check to scalarfield modulus match

7f054f4

chore: correct the docs, use new refactored code for correct implemen…

c820867

…tation

chore: update path for contants.metal

f5375a5

test(mont_mul): adapted mont mul tests from https://github.com/geomet…

d36a9ec

…ryxyz/msl-secp256k1

feat(mont_mul): adapted utils function related to mont_mul from https…

155389c

…://github.com/geometers/multiprecision

feat: add limbs conversion from ark_ff bigint to arbitrary limbs in V…

fe96fa7

…ec<u32>

chore: update path for utils for limb conversion

9ca0419

fix(mont_mul): correct the conversion from arkworks' scalarfield to a…

1578e15

…rbitrary limbs of Vec<u32>

test(mont_mul): adapted benchmark functions from https://github.com/g…

2b99568

…eometryxyz/msl-secp256k1

feat(metal): add mont_mul cios

57bed4d

test(mont_mul): add cios mont_mul test and benchmark

d6c81b2

moven0831 requested a review from FoodChain1028 November 8, 2024 02:33

moven0831 self-assigned this Nov 8, 2024

moven0831 linked an issue Nov 8, 2024 that may be closed by this pull request

Implement Optimized Montgomery Multiplication for Metal #6

Open

yaroslavyaroslav approved these changes Nov 9, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

refactor(metal_mont_backend): adapted from https://github.com/geometryxyz/msl-secp256k1 #13

refactor(metal_mont_backend): adapted from https://github.com/geometryxyz/msl-secp256k1 #13

moven0831 commented Nov 8, 2024

yaroslavyaroslav left a comment

yaroslavyaroslav Nov 8, 2024

yaroslavyaroslav Nov 9, 2024

yaroslavyaroslav Nov 9, 2024

refactor(metal_mont_backend): adapted from https://github.com/geometryxyz/msl-secp256k1 #13

Are you sure you want to change the base?

refactor(metal_mont_backend): adapted from https://github.com/geometryxyz/msl-secp256k1 #13

Conversation

moven0831 commented Nov 8, 2024

Overview

Related Issue

Key Changes

Core Implementation

Testing & Benchmarking

Utility Functions

Dependencies

Next Steps

Reference

yaroslavyaroslav left a comment

Choose a reason for hiding this comment

yaroslavyaroslav Nov 8, 2024

Choose a reason for hiding this comment

yaroslavyaroslav Nov 9, 2024

Choose a reason for hiding this comment

yaroslavyaroslav Nov 9, 2024

Choose a reason for hiding this comment