perf(poly): use MontMul for degree inverse in INTT

Author: batzor

benchmark/ntt/ntt_benchmark_test.cc

@@ -140,12 +140,12 @@ BENCHMARK(BM_intt_mont_benchmark)->Iterations(1)->Unit(::benchmark::kSecond);
 //   L1 Data 64 KiB
 //   L1 Instruction 128 KiB
 //   L2 Unified 4096 KiB (x14)
-// Load Average: 7.37, 7.53, 7.10
+// Load Average: 3.97, 3.11, 2.96
 // ------------------------------------------------------------------------------
 // Benchmark                                    Time             CPU   Iterations
 // ------------------------------------------------------------------------------
-// BM_ntt_benchmark                          10.4 s          10.3 s             1
-// BM_intt_benchmark/iterations:1            12.1 s          11.4 s             1
-// BM_ntt_mont_benchmark                    0.201 s         0.197 s             3
-// BM_intt_mont_benchmark/iterations:1       1.38 s          1.35 s             1
+// BM_ntt_benchmark                          10.2 s          10.1 s             1
+// BM_intt_benchmark/iterations:1            11.1 s          11.1 s             1
+// BM_ntt_mont_benchmark                    0.190 s         0.190 s             3
+// BM_intt_mont_benchmark/iterations:1      0.316 s         0.304 s             1
 // NOLINTEND()

zkir/Dialect/Poly/Conversions/PolyToField/PolyToField.cpp

@@ -494,7 +494,12 @@ static Value fastNTT(ImplicitLocOpBuilder &b, PrimitiveRootAttr rootAttr,
         b.create<arith::ConstantOp>(rootAttr.getInvDegree().getValue());
     auto degreeInvTensor = b.create<tensor::SplatOp>(invDegreeConst, rootsType);
     auto fieldTensor = b.create<field::EncapsulateOp>(modType, degreeInvTensor);
-    result = b.create<field::MulOp>(result, fieldTensor);
+    if (rootAttr.getMontgomery() != mod_arith::MontgomeryAttr()) {
+      result = b.create<field::MontMulOp>(result, fieldTensor,
+                                          rootAttr.getMontgomery());
+    } else {
+      result = b.create<field::MulOp>(result, fieldTensor);
+    }
   }
 
   return result;

zkir/Dialect/Poly/IR/PolyAttributes.cpp

@@ -94,23 +94,28 @@ PrimitiveRootAttrStorage *PrimitiveRootAttrStorage::construct(
   IntegerAttr degree = std::get<1>(key);
   zkir::mod_arith::MontgomeryAttr montgomery = std::get<2>(key);
 
+  std::optional<IntegerAttr> montgomeryR;
+  if (montgomery != zkir::mod_arith::MontgomeryAttr()) {
+    montgomeryR = montgomery.getR();
+  }
+
   APInt mod = root.getType().getModulus().getValue();
   APInt rootVal = root.getValue().getValue();
   APInt invRootVal = multiplicativeInverse(rootVal, mod);
   APInt invDegreeVal = multiplicativeInverse(
       degree.getValue().zextOrTrunc(mod.getBitWidth()), mod);
 
+  if (montgomeryR) {
+    invDegreeVal = mulMod(invDegreeVal, montgomeryR->getValue(), mod);
+  }
+
   field::PrimeFieldAttr invDegree =
       field::PrimeFieldAttr::get(root.getType(), invDegreeVal);
   zkir::field::PrimeFieldAttr invRoot =
       zkir::field::PrimeFieldAttr::get(root.getType(), invRootVal);
 
   // Compute the exponent table.
   SmallVector<APInt> roots, invRoots;
-  std::optional<IntegerAttr> montgomeryR;
-  if (montgomery != zkir::mod_arith::MontgomeryAttr()) {
-    montgomeryR = montgomery.getR();
-  }
   precomputeRoots(rootVal, mod, degree.getInt(), roots, invRoots, montgomeryR);
   // Create a ranked tensor type for the exponents attribute.
   auto tensorType = RankedTensorType::get(