Skip to content

Commit 795c240

Browse files
oriziclaude
andauthored
refactor(sqrt): Make sqrt libfuncs runtime-bindings based. (#1644)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent ddcc617 commit 795c240

5 files changed

Lines changed: 235 additions & 483 deletions

File tree

src/libfuncs/int.rs

Lines changed: 10 additions & 145 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,14 @@ use crate::{
33
error::{panic::ToNativeAssertError, Result},
44
execution_result::BITWISE_BUILTIN_SIZE,
55
libfuncs::{increment_builtin_counter, increment_builtin_counter_by},
6-
metadata::MetadataStorage,
7-
native_panic,
6+
metadata::{runtime_bindings::RuntimeBindingsMeta, MetadataStorage},
87
types::TypeBuilder,
98
utils::{ProgramRegistryExt, PRIME},
109
};
1110
use cairo_lang_sierra::{
1211
extensions::{
1312
bounded_int::BoundedIntDivRemAlgorithm,
14-
core::{CoreLibfunc, CoreType, CoreTypeConcrete},
13+
core::{CoreLibfunc, CoreType},
1514
int::{
1615
signed::{SintConcrete, SintTraits},
1716
signed128::Sint128Concrete,
@@ -35,8 +34,8 @@ use melior::{
3534
},
3635
helpers::{ArithBlockExt, BuiltinBlockExt, LlvmBlockExt},
3736
ir::{
38-
attribute::IntegerAttribute, operation::OperationBuilder, r#type::IntegerType, Block,
39-
BlockLike, Location, Region, ValueLike,
37+
operation::OperationBuilder, r#type::IntegerType, Block, BlockLike, Location, Region,
38+
ValueLike,
4039
},
4140
Context,
4241
};
@@ -689,157 +688,23 @@ fn build_operation<'ctx, 'this>(
689688

690689
fn build_square_root<'ctx, 'this>(
691690
context: &'ctx Context,
692-
registry: &ProgramRegistry<CoreType, CoreLibfunc>,
691+
_registry: &ProgramRegistry<CoreType, CoreLibfunc>,
693692
entry: &'this Block<'ctx>,
694693
location: Location<'ctx>,
695694
helper: &LibfuncHelper<'ctx, 'this>,
696-
_metadata: &mut MetadataStorage,
697-
info: &SignatureOnlyConcreteLibfunc,
695+
metadata: &mut MetadataStorage,
696+
_info: &SignatureOnlyConcreteLibfunc,
698697
) -> Result<()> {
699698
// The sierra-to-casm compiler uses the range_check builtin 4 times.
700699
// https://github.com/starkware-libs/cairo/blob/v2.12.0-dev.1/crates/cairo-lang-sierra-to-casm/src/invocations/int/unsigned.rs#L73
701700
let range_check =
702701
super::increment_builtin_counter_by(context, entry, location, entry.arg(0)?, 4)?;
703702

704703
let input = entry.arg(1)?;
705-
let (input_bits, value_bits) =
706-
match registry.get_type(&info.signature.param_signatures[1].ty)? {
707-
CoreTypeConcrete::Uint8(_) => (8, 8),
708-
CoreTypeConcrete::Uint16(_) => (16, 8),
709-
CoreTypeConcrete::Uint32(_) => (32, 16),
710-
CoreTypeConcrete::Uint64(_) => (64, 32),
711-
CoreTypeConcrete::Uint128(_) => (128, 64),
712-
_ => native_panic!("invalid value type in int square root"),
713-
};
714-
715-
let k1 = entry.const_int(context, location, 1, input_bits)?;
716-
let is_small = entry.cmpi(context, CmpiPredicate::Ule, input, k1, location)?;
717-
718-
let value = entry.append_op_result(scf::r#if(
719-
is_small,
720-
&[IntegerType::new(context, value_bits).into()],
721-
{
722-
let region = Region::new();
723-
let block = region.append_block(Block::new(&[]));
724-
725-
let value = block.trunci(
726-
input,
727-
IntegerType::new(context, value_bits).into(),
728-
location,
729-
)?;
730-
731-
block.append_operation(scf::r#yield(&[value], location));
732-
region
733-
},
734-
{
735-
let region = Region::new();
736-
let block = region.append_block(Block::new(&[]));
737704

738-
let leading_zeros = block.append_op_result(
739-
ods::llvm::intr_ctlz(
740-
context,
741-
IntegerType::new(context, input_bits).into(),
742-
input,
743-
IntegerAttribute::new(IntegerType::new(context, 1).into(), 1),
744-
location,
745-
)
746-
.into(),
747-
)?;
748-
749-
let k_bits = block.const_int(context, location, input_bits, input_bits)?;
750-
let num_bits = block.append_op_result(arith::subi(k_bits, leading_zeros, location))?;
751-
let shift_amount = block.addi(num_bits, k1, location)?;
752-
753-
let parity_mask = block.const_int(context, location, -2, input_bits)?;
754-
let shift_amount =
755-
block.append_op_result(arith::andi(shift_amount, parity_mask, location))?;
756-
757-
let k0 = block.const_int(context, location, 0, input_bits)?;
758-
let value = block.append_op_result(scf::r#while(
759-
&[k0, shift_amount],
760-
&[
761-
IntegerType::new(context, input_bits).into(),
762-
IntegerType::new(context, input_bits).into(),
763-
],
764-
{
765-
let region = Region::new();
766-
let block = region.append_block(Block::new(&[
767-
(IntegerType::new(context, input_bits).into(), location),
768-
(IntegerType::new(context, input_bits).into(), location),
769-
]));
770-
771-
let value = block.shli(block.arg(0)?, k1, location)?;
772-
let large_candidate =
773-
block.append_op_result(arith::xori(value, k1, location))?;
774-
let large_candidate_squared =
775-
block.muli(large_candidate, large_candidate, location)?;
776-
777-
let threshold = block.shrui(input, block.arg(1)?, location)?;
778-
let threshold_is_poison =
779-
block.cmpi(context, CmpiPredicate::Eq, block.arg(1)?, k_bits, location)?;
780-
let threshold = block.append_op_result(arith::select(
781-
threshold_is_poison,
782-
k0,
783-
threshold,
784-
location,
785-
))?;
786-
787-
let is_in_range = block.cmpi(
788-
context,
789-
CmpiPredicate::Ule,
790-
large_candidate_squared,
791-
threshold,
792-
location,
793-
)?;
794-
let value = block.append_op_result(arith::select(
795-
is_in_range,
796-
large_candidate,
797-
value,
798-
location,
799-
))?;
800-
801-
let k2 = block.const_int(context, location, 2, input_bits)?;
802-
let shift_amount =
803-
block.append_op_result(arith::subi(block.arg(1)?, k2, location))?;
804-
805-
let should_continue =
806-
block.cmpi(context, CmpiPredicate::Sge, shift_amount, k0, location)?;
807-
block.append_operation(scf::condition(
808-
should_continue,
809-
&[value, shift_amount],
810-
location,
811-
));
812-
813-
region
814-
},
815-
{
816-
let region = Region::new();
817-
let block = region.append_block(Block::new(&[
818-
(IntegerType::new(context, input_bits).into(), location),
819-
(IntegerType::new(context, input_bits).into(), location),
820-
]));
821-
822-
block.append_operation(scf::r#yield(&[block.arg(0)?, block.arg(1)?], location));
823-
region
824-
},
825-
location,
826-
))?;
827-
828-
let value = if input_bits == value_bits {
829-
value
830-
} else {
831-
block.trunci(
832-
value,
833-
IntegerType::new(context, value_bits).into(),
834-
location,
835-
)?
836-
};
837-
838-
block.append_operation(scf::r#yield(&[value], location));
839-
region
840-
},
841-
location,
842-
))?;
705+
let runtime_bindings = metadata.get_or_insert_with(RuntimeBindingsMeta::default);
706+
let value =
707+
runtime_bindings.integer_square_root(context, helper.module, entry, location, input)?;
843708

844709
helper.br(entry, 0, &[range_check, value], location)
845710
}

0 commit comments

Comments
 (0)