ZIR-194: Fix and test BigInt type inference

zirgen/Dialect/BigInt/IR/BUILD.bazel

@@ -82,6 +82,7 @@ cc_library(
  "Dialect.cpp",
  "Eval.cpp",
  "Ops.cpp",
+ "Types.cpp",
  ],
  hdrs = [
  "BigInt.h",

zirgen/Dialect/BigInt/IR/Eval.cpp

@@ -110,7 +110,7 @@ BytePoly nondetRem(const BytePoly& lhs, const BytePoly& rhs, size_t coeffs) {
  return fromAPInt(rem, coeffs);
 }
 
-BytePoly nondetInvMod(const BytePoly& lhs, const BytePoly& rhs, size_t coeffs) {
+BytePoly nondetInv(const BytePoly& lhs, const BytePoly& rhs, size_t coeffs) {
  // Uses the formula n^(p-2) * n = 1 (mod p) to invert `lhs` (mod `rhs`)
  // (via the square and multiply technique)
  auto lhsInt = toAPInt(lhs);
@@ -226,9 +226,9 @@ EvalOutput eval(func::FuncOp inFunc, ArrayRef<APInt> witnessValues) {
  polys[op.getOut()] = poly;
  ret.privateWitness.push_back(poly);
  })
- .Case<NondetInvModOp>([&](auto op) {
+ .Case<NondetInvOp>([&](auto op) {
  uint32_t coeffs = op.getOut().getType().getCoeffs();
- auto poly = nondetInvMod(polys[op.getLhs()], polys[op.getRhs()], coeffs);
+ auto poly = nondetInv(polys[op.getLhs()], polys[op.getRhs()], coeffs);
  polys[op.getOut()] = poly;
  ret.privateWitness.push_back(poly);
  })

zirgen/Dialect/BigInt/IR/Ops.cpp

@@ -25,6 +25,9 @@
 using namespace mlir;
 using risc0::ceilDiv;
 
+// Additional comments on how type inference works for the BigInt dialect can be found in
+// `test/type_infer.mlir`, including descriptions at the beginning of each op's suite of tests.
+
 namespace zirgen::BigInt {
 
 // Type inference
@@ -66,9 +69,10 @@ LogicalResult AddOp::inferReturnTypes(MLIRContext* ctx,
  auto lhsType = adaptor.getLhs().getType().cast<BigIntType>();
  auto rhsType = adaptor.getRhs().getType().cast<BigIntType>();
  size_t maxCoeffs = std::max(lhsType.getCoeffs(), rhsType.getCoeffs());
- size_t maxPos = std::max(lhsType.getMaxPos(), rhsType.getMaxPos());
- size_t maxNeg = std::max(lhsType.getMaxNeg(), rhsType.getMaxNeg());
- size_t minBits = std::max(lhsType.getMinBits(), rhsType.getMinBits());
+ size_t maxPos = lhsType.getMaxPos() + rhsType.getMaxPos();
+ size_t maxNeg = lhsType.getMaxNeg() + rhsType.getMaxNeg();
+ // TODO: We could be more clever on minBits, but probably doesn't matter
+ size_t minBits = maxNeg > 0 ? 0 : std::max(lhsType.getMinBits(), rhsType.getMinBits());
  out.push_back(BigIntType::get(ctx, maxCoeffs, maxPos, maxNeg, minBits));
  return success();
 }
@@ -80,8 +84,8 @@ LogicalResult SubOp::inferReturnTypes(MLIRContext* ctx,
  auto lhsType = adaptor.getLhs().getType().cast<BigIntType>();
  auto rhsType = adaptor.getRhs().getType().cast<BigIntType>();
  size_t maxCoeffs = std::max(lhsType.getCoeffs(), rhsType.getCoeffs());
- size_t maxPos = std::max(lhsType.getMaxPos(), rhsType.getMaxNeg());
- size_t maxNeg = std::max(lhsType.getMaxNeg(), rhsType.getMaxPos());
+ size_t maxPos = lhsType.getMaxPos() + rhsType.getMaxNeg();
+ size_t maxNeg = lhsType.getMaxNeg() + rhsType.getMaxPos();
  // TODO: We could be more clever on minBits, but probably doesn't matter
  out.push_back(BigIntType::get(ctx, maxCoeffs, maxPos, maxNeg, 0));
  return success();
@@ -93,30 +97,56 @@ LogicalResult MulOp::inferReturnTypes(MLIRContext* ctx,
  SmallVectorImpl<Type>& out) {
  auto lhsType = adaptor.getLhs().getType().cast<BigIntType>();
  auto rhsType = adaptor.getRhs().getType().cast<BigIntType>();
- size_t maxCoeffs = std::max(lhsType.getCoeffs(), rhsType.getCoeffs());
- size_t totCoeffs = lhsType.getCoeffs() + rhsType.getCoeffs();
- size_t maxPos = std::max(lhsType.getMaxPos() * rhsType.getMaxPos(),
- lhsType.getMaxNeg() * rhsType.getMaxNeg()) *
- maxCoeffs;
- size_t maxNeg = std::max(lhsType.getMaxPos() * rhsType.getMaxNeg(),
- lhsType.getMaxNeg() * rhsType.getMaxPos()) *
- maxCoeffs;
+ size_t coeffs = lhsType.getCoeffs() + rhsType.getCoeffs() - 1;
+ // The maximum number of coefficient pairs from the inputs used to calculate an output coefficient
+ size_t maxCoeffs = std::min(lhsType.getCoeffs(), rhsType.getCoeffs());
+ // This calculation could overflow if size_t is 32 bits, so cast to 64 bits
+ uint64_t maxPos = std::max((uint64_t)lhsType.getMaxPos() * rhsType.getMaxPos(),
+ (uint64_t)lhsType.getMaxNeg() * rhsType.getMaxNeg());
+ // The next step can potentially overflow even 64 bits; but if we're already above 32 bits we'll
+ // fail validation anyway. Therefore, skip this if we're above 32 bits
+ if (maxPos < (uint64_t)1 << 32) {
+ maxPos *= maxCoeffs;
+ }
+ // Clamp to size_t
+ if (maxPos > std::numeric_limits<size_t>::max()) {
+ maxPos = std::numeric_limits<size_t>::max();
+ }
+ // As with maxPos, this could overflow if size_t is 32 bits, so cast to 64 bits
+ uint64_t maxNeg = std::max((uint64_t)lhsType.getMaxPos() * rhsType.getMaxNeg(),
+ (uint64_t)lhsType.getMaxNeg() * rhsType.getMaxPos());
+ // The next step can potentially overflow even 64 bits; but if we're already above 32 bits we'll
+ // fail validation anyway. Therefore, skip this if we're above 32 bits
+ if (maxNeg < (uint64_t)1 << 32) {
+ maxNeg *= maxCoeffs;
+ }
+ // Clamp to size_t
+ if (maxNeg > std::numeric_limits<size_t>::max()) {
+ maxNeg = std::numeric_limits<size_t>::max();
+ }
  size_t minBits;
  if (lhsType.getMinBits() == 0 || rhsType.getMinBits() == 0) {
+ // Note that this catches _both_ cases where the input might be zero _and_ cases where the input
+ // might be negative, as type verification enforces that when minBits is zero, so is maxNeg.
  minBits = 0;
  } else {
  minBits = lhsType.getMinBits() + rhsType.getMinBits() - 1;
  }
- out.push_back(BigIntType::get(ctx, totCoeffs, maxPos, maxNeg, minBits));
+ out.push_back(BigIntType::get(ctx, coeffs, maxPos, maxNeg, minBits));
  return success();
 }
 
 LogicalResult NondetRemOp::inferReturnTypes(MLIRContext* ctx,
  std::optional<Location> loc,
  Adaptor adaptor,
  SmallVectorImpl<Type>& out) {
+ auto lhsType = adaptor.getLhs().getType().cast<BigIntType>();
  auto rhsType = adaptor.getRhs().getType().cast<BigIntType>();
- size_t coeffsWidth = ceilDiv(rhsType.getMaxBits(), kBitsPerCoeff);
+ auto outBits = lhsType.getMaxPosBits();
+ if (rhsType.getMaxPosBits() < outBits) {
+ outBits = rhsType.getMaxPosBits();
+ }
+ size_t coeffsWidth = ceilDiv(outBits, kBitsPerCoeff);
  out.push_back(BigIntType::get(ctx,
  /*coeffs=*/coeffsWidth,
  /*maxPos=*/(1 << kBitsPerCoeff) - 1,
@@ -131,26 +161,26 @@ LogicalResult NondetQuotOp::inferReturnTypes(MLIRContext* ctx,
  SmallVectorImpl<Type>& out) {
  auto lhsType = adaptor.getLhs().getType().cast<BigIntType>();
  auto rhsType = adaptor.getRhs().getType().cast<BigIntType>();
- size_t outBits = lhsType.getMaxBits();
+ size_t outBits = lhsType.getMaxPosBits();
  if (rhsType.getMinBits() > 0) {
  outBits -= rhsType.getMinBits() - 1;
  }
  size_t coeffsWidth = ceilDiv(outBits, kBitsPerCoeff);
+ // TODO: We could be more clever on minBits, but probably doesn't matter
  out.push_back(BigIntType::get(ctx,
  /*coeffs=*/coeffsWidth,
  /*maxPos=*/(1 << kBitsPerCoeff) - 1,
  /*maxNeg=*/0,
- /*minBits=*/0 /*TODO: maybe better bound? */
- ));
+ /*minBits=*/0));
  return success();
 }
 
-LogicalResult NondetInvModOp::inferReturnTypes(MLIRContext* ctx,
-  std::optional<Location> loc,
-  Adaptor adaptor,
-  SmallVectorImpl<Type>& out) {
+LogicalResult NondetInvOp::inferReturnTypes(MLIRContext* ctx,
+ std::optional<Location> loc,
+ Adaptor adaptor,
+ SmallVectorImpl<Type>& out) {
  auto rhsType = adaptor.getRhs().getType().cast<BigIntType>();
- size_t coeffsWidth = ceilDiv(rhsType.getMaxBits(), kBitsPerCoeff);
+ size_t coeffsWidth = ceilDiv(rhsType.getMaxPosBits(), kBitsPerCoeff);
  out.push_back(BigIntType::get(ctx,
  /*coeffs=*/coeffsWidth,
  /*maxPos=*/(1 << kBitsPerCoeff) - 1,
@@ -159,12 +189,12 @@ LogicalResult NondetInvModOp::inferReturnTypes(MLIRContext* ctx,
  return success();
 }
 
-LogicalResult ModularInvOp::inferReturnTypes(MLIRContext* ctx,
-  std::optional<Location> loc,
-  Adaptor adaptor,
-  SmallVectorImpl<Type>& out) {
+LogicalResult InvOp::inferReturnTypes(MLIRContext* ctx,
+ std::optional<Location> loc,
+ Adaptor adaptor,
+ SmallVectorImpl<Type>& out) {
  auto rhsType = adaptor.getRhs().getType().cast<BigIntType>();
- size_t coeffsWidth = ceilDiv(rhsType.getMaxBits(), kBitsPerCoeff);
+ size_t coeffsWidth = ceilDiv(rhsType.getMaxPosBits(), kBitsPerCoeff);
  out.push_back(BigIntType::get(ctx,
  /*coeffs=*/coeffsWidth,
  /*maxPos=*/(1 << kBitsPerCoeff) - 1,
@@ -177,8 +207,13 @@ LogicalResult ReduceOp::inferReturnTypes(MLIRContext* ctx,
  std::optional<Location> loc,
  Adaptor adaptor,
  SmallVectorImpl<Type>& out) {
+ auto lhsType = adaptor.getLhs().getType().cast<BigIntType>();
  auto rhsType = adaptor.getRhs().getType().cast<BigIntType>();
- size_t coeffsWidth = ceilDiv(rhsType.getMaxBits(), kBitsPerCoeff);
+ auto outBits = lhsType.getMaxPosBits();
+ if (rhsType.getMaxPosBits() < outBits) {
+ outBits = rhsType.getMaxPosBits();
+ }
+ size_t coeffsWidth = ceilDiv(outBits, kBitsPerCoeff);
  out.push_back(BigIntType::get(ctx,
  /*coeffs=*/coeffsWidth,
  /*maxPos=*/(1 << kBitsPerCoeff) - 1,
@@ -254,7 +289,7 @@ void NondetQuotOp::emitExpr(codegen::CodegenEmitter& cg) {
  {getLhs(), getRhs(), toConstantValue(cg, getContext(), getType().getCoeffs())});
 }
 
-void NondetInvModOp::emitExpr(codegen::CodegenEmitter& cg) {
+void NondetInvOp::emitExpr(codegen::CodegenEmitter& cg) {
  cg.emitInvokeMacro(
  cg.getStringAttr("bigint_nondet_inv"),
  /*contextArgs=*/{"ctx"},

zirgen/Dialect/BigInt/IR/Ops.td

@@ -50,8 +50,8 @@ def SubOp : BinaryOp<"sub", [Pure, DeclareOpInterfaceMethods<CodegenExprOpInterf
 def MulOp : BinaryOp<"mul", [Pure, Commutative, DeclareOpInterfaceMethods<CodegenExprOpInterface>]> {}
 def NondetRemOp : BinaryOp<"nondet_rem", [DeclareOpInterfaceMethods<CodegenExprOpInterface>]> {}
 def NondetQuotOp : BinaryOp<"nondet_quot", [DeclareOpInterfaceMethods<CodegenExprOpInterface>]> {}
-def NondetInvModOp : BinaryOp<"nondet_invmod", [DeclareOpInterfaceMethods<CodegenExprOpInterface>]> {}
-def ModularInvOp : BinaryOp<"inv", []> {}
+def NondetInvOp : BinaryOp<"nondet_inv", [DeclareOpInterfaceMethods<CodegenExprOpInterface>]> {}
+def InvOp : BinaryOp<"inv", []> {}
 def ReduceOp : BinaryOp<"reduce", []> {}
 
 def EqualZeroOp : BigIntOp<"eqz", [DeclareOpInterfaceMethods<CodegenExprOpInterface>]> {

zirgen/Dialect/BigInt/IR/Types.cpp

@@ -0,0 +1,41 @@
+// Copyright 2024 RISC Zero, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "risc0/fp/fp.h"
+#include "zirgen/Dialect/BigInt/IR/BigInt.h"
+#include "zirgen/Dialect/BigInt/IR/Types.h.inc"
+
+using namespace mlir;
+
+namespace zirgen::BigInt {
+
+LogicalResult BigIntType::verify(function_ref<InFlightDiagnostic()> emitError,
+ size_t coeffs,
+ size_t maxPos,
+ size_t maxNeg,
+ size_t minBits) {
+ if (maxNeg > 0 && minBits > 0) {
+ return emitError() << "BigInts with positive minBits must be positive: maxNeg: " << maxNeg
+ << ", minBits: " << minBits;
+ }
+ // TODO: Think through whether maxPos / maxNeg can ever overflow their attribute type, which would
+ // cause problems here
+ if (maxPos + maxNeg >= risc0::Fp::P) {
+ return emitError() << "Cannot create BigInt with coefficients overflowing BabyBear: maxPos: "
+ << maxPos << " + maxNeg: " << maxNeg;
+ }
+ return success();
+}
+
+} // namespace zirgen::BigInt

zirgen/Dialect/BigInt/IR/Types.td

@@ -25,17 +25,37 @@ def BigInt : BigIntType<"BigInt", "bigint", [
  DeclareTypeInterfaceMethods<CodegenTypeInterface, ["getTypeName", "allowDuplicateTypeNames", "emitTypeDefinition"]>,
  CodegenNeedsCloneType
 ]> {
- let summary = "A big interger value represented as a polynomial";
+ let summary = "A big integer value represented as a polynomial";
  let parameters = (ins
  "size_t": $coeffs, // Number of polynomial coefficents
  "size_t": $maxPos, // Maximum positive coefficient value
  "size_t": $maxNeg, // Maximum negative coefficient value
  "size_t": $minBits // If minBits == 0, no constraint, otherwise N >= 2^(minBits - 1)
  );
  let assemblyFormat = "`<` $coeffs `,` $maxPos `,` $maxNeg `,` $minBits `>`";
+ let genVerifyDecl = 1;
  let extraClassDeclaration = [{
- size_t getMaxBits() {
- size_t extraBits = risc0::log2Ceil(getMaxPos() / (1 << kBitsPerCoeff));
+ size_t getMaxPosBits() {
+ // Because 2^k requires k+1 bits to represent, we add 1 to getMaxPos before log2Ceil
+ size_t extraBits = risc0::log2Ceil(getMaxPos() + 1);
+ if (extraBits <= kBitsPerCoeff) {
+ // When maxPos fits in a coeff, no extra bits are needed
+ extraBits = 0;
+ } else {
+ extraBits -= kBitsPerCoeff;
+ // Carries can sometimes lead to 1 extra bit so add 1
+ // Specifically, we know that the max value is
+ // getMaxPos() * sum over i in [0, getCoeffs()) of (2^kBitsPerCoeff)^i
+ // which is bounded above by
+ // getMaxPos() * 2^((kBitsPerCoeff * (getCoeffs() - 1)) + 1)
+ // which has bitwidth maxPosBitwidth + kBitsPerCoeff * (getCoeffs() - 1) + 1 which is
+ // kBitsPerCoeff * getCoeffs + extraBits
+ // where
+ // extraBits = maxPosBitwidth - kBitsPerCoeff + 1
+ if (getCoeffs() > 1) {
+ extraBits += 1;
+ }
+ }
  return kBitsPerCoeff * getCoeffs() + extraBits;
  }
  size_t getCarryOffset() {

zirgen/Dialect/BigInt/IR/test/BUILD.bazel

@@ -2,6 +2,10 @@ package(
  default_visibility = ["//visibility:public"],
 )
 
+load("//bazel/rules/lit:defs.bzl", "glob_lit_tests")
+
+glob_lit_tests()
+
 cc_binary(
  name = "test",
  srcs = [