nac3core/src/codegen/builtin_fns.rs

@@ -661,6 +661,90 @@ pub fn call_min<'ctx>(
     }
 }
 
+/// Invokes the `np_min` builtin function.
+pub fn call_numpy_min<'ctx, G: CodeGenerator + ?Sized>(
+    generator: &mut G,
+    ctx: &mut CodeGenContext<'ctx, '_>,
+    a: (Type, BasicValueEnum<'ctx>),
+) -> Result<BasicValueEnum<'ctx>, String> {
+    const FN_NAME: &str = "np_min";
+
+    let llvm_usize = generator.get_size_type(ctx.ctx);
+
+    let (a_ty, a) = a;
+
+    Ok(match a {
+        BasicValueEnum::IntValue(_) | BasicValueEnum::FloatValue(_) => {
+            debug_assert!([
+                ctx.primitives.bool,
+                ctx.primitives.int32,
+                ctx.primitives.uint32,
+                ctx.primitives.int64,
+                ctx.primitives.uint64,
+                ctx.primitives.float,
+            ]
+            .iter()
+            .any(|ty| ctx.unifier.unioned(a_ty, *ty)));
+
+            a
+        }
+
+        BasicValueEnum::PointerValue(n)
+            if a_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) =>
+        {
+            let (elem_ty, _) = unpack_ndarray_var_tys(&mut ctx.unifier, a_ty);
+            let llvm_ndarray_ty = ctx.get_llvm_type(generator, elem_ty);
+
+            let n = NDArrayValue::from_ptr_val(n, llvm_usize, None);
+            let n_sz = irrt::call_ndarray_calc_size(generator, ctx, &n.dim_sizes(), (None, None));
+            if ctx.registry.llvm_options.opt_level == OptimizationLevel::None {
+                let n_sz_eqz = ctx
+                    .builder
+                    .build_int_compare(IntPredicate::NE, n_sz, n_sz.get_type().const_zero(), "")
+                    .unwrap();
+
+                ctx.make_assert(
+                    generator,
+                    n_sz_eqz,
+                    "0:ValueError",
+                    "zero-size array to reduction operation minimum which has no identity",
+                    [None, None, None],
+                    ctx.current_loc,
+                );
+            }
+
+            let accumulator_addr = generator.gen_var_alloc(ctx, llvm_ndarray_ty, None)?;
+            unsafe {
+                let identity =
+                    n.data().get_unchecked(ctx, generator, &llvm_usize.const_zero(), None);
+                ctx.builder.build_store(accumulator_addr, identity).unwrap();
+            }
+
+            gen_for_callback_incrementing(
+                generator,
+                ctx,
+                llvm_usize.const_int(1, false),
+                (n_sz, false),
+                |generator, ctx, _, idx| {
+                    let elem = unsafe { n.data().get_unchecked(ctx, generator, &idx, None) };
+
+                    let accumulator = ctx.builder.build_load(accumulator_addr, "").unwrap();
+                    let result = call_min(ctx, (elem_ty, accumulator), (elem_ty, elem));
+                    ctx.builder.build_store(accumulator_addr, result).unwrap();
+
+                    Ok(())
+                },
+                llvm_usize.const_int(1, false),
+            )?;
+
+            let accumulator = ctx.builder.build_load(accumulator_addr, "").unwrap();
+            accumulator
+        }
+
+        _ => unsupported_type(ctx, FN_NAME, &[a_ty]),
+    })
+}
+
 /// Invokes the `np_minimum` builtin function.
 pub fn call_numpy_minimum<'ctx, G: CodeGenerator + ?Sized>(
     generator: &mut G,
@@ -793,21 +877,19 @@ pub fn call_max<'ctx>(
     }
 }
 
-/// Invokes the np_max, np_min, np_argmax, np_argmin functions
-/// * `fn_name`: Can be one of "np_argmin", "np_argmax", "np_max", "np_min"
-pub fn call_numpy_max_min<'ctx, G: CodeGenerator + ?Sized>(
+/// Invokes the `np_max` builtin function.
+pub fn call_numpy_max<'ctx, G: CodeGenerator + ?Sized>(
     generator: &mut G,
     ctx: &mut CodeGenContext<'ctx, '_>,
     a: (Type, BasicValueEnum<'ctx>),
-    fn_name: &str,
 ) -> Result<BasicValueEnum<'ctx>, String> {
-    debug_assert!(["np_argmin", "np_argmax", "np_max", "np_min"].iter().any(|f| *f == fn_name));
+    const FN_NAME: &str = "np_max";
 
-    let llvm_int64 = ctx.ctx.i64_type();
     let llvm_usize = generator.get_size_type(ctx.ctx);
 
     let (a_ty, a) = a;
-    Ok( match a {
+
+    Ok(match a {
         BasicValueEnum::IntValue(_) | BasicValueEnum::FloatValue(_) => {
             debug_assert!([
                 ctx.primitives.bool,
@@ -820,12 +902,9 @@ pub fn call_numpy_max_min<'ctx, G: CodeGenerator + ?Sized>(
             .iter()
             .any(|ty| ctx.unifier.unioned(a_ty, *ty)));
 
-            match fn_name {
-                "np_argmin" | "np_argmax" => llvm_int64.const_zero().into(),
-                "np_max" | "np_min" => a,
-                _ => unreachable!()
-            }
+            a
         }
+
         BasicValueEnum::PointerValue(n)
             if a_ty.obj_id(&ctx.unifier).is_some_and(|id| id == PrimDef::NDArray.id()) =>
         {
@@ -844,71 +923,41 @@ pub fn call_numpy_max_min<'ctx, G: CodeGenerator + ?Sized>(
                     generator,
                     n_sz_eqz,
                     "0:ValueError",
-                    format!("zero-size array to reduction operation {}", fn_name).as_str(),
+                    "zero-size array to reduction operation minimum which has no identity",
                     [None, None, None],
                     ctx.current_loc,
                 );
             }
 
             let accumulator_addr = generator.gen_var_alloc(ctx, llvm_ndarray_ty, None)?;
-            let res_idx = generator.gen_var_alloc(ctx, llvm_int64.into(), None)?;
-
             unsafe {
                 let identity =
                     n.data().get_unchecked(ctx, generator, &llvm_usize.const_zero(), None);
                 ctx.builder.build_store(accumulator_addr, identity).unwrap();
-                ctx.builder.build_store(res_idx, llvm_int64.const_zero()).unwrap();
             }
 
             gen_for_callback_incrementing(
                 generator,
                 ctx,
-                llvm_int64.const_int(1, false),
+                llvm_usize.const_int(1, false),
                 (n_sz, false),
-                |generator, ctx, _, idx,| {
+                |generator, ctx, _, idx| {
                     let elem = unsafe { n.data().get_unchecked(ctx, generator, &idx, None) };
+
                     let accumulator = ctx.builder.build_load(accumulator_addr, "").unwrap();
-                    let cur_idx = ctx.builder.build_load(res_idx, "").unwrap();
-
-                    let result = match fn_name {
-                        "np_argmin" | "np_min" => call_min(ctx, (elem_ty, accumulator), (elem_ty, elem)),
-                        "np_argmax" | "np_max" => call_max(ctx, (elem_ty, accumulator), (elem_ty, elem)),
-                        _ => unreachable!()
-                    };
-
-                    let updated_idx = match (accumulator, result){
-                        (BasicValueEnum::IntValue(m), BasicValueEnum::IntValue(n)) => {
-                            ctx.builder.build_select(
-                                ctx.builder.build_int_compare(IntPredicate::NE,m, n, "").unwrap(), 
-                                idx.into(), 
-                                cur_idx,
-                                "").unwrap()
-                        },
-                        (BasicValueEnum::FloatValue(m), BasicValueEnum::FloatValue(n)) => {
-                            ctx.builder.build_select(
-                                ctx.builder.build_float_compare(FloatPredicate::ONE,m, n, "").unwrap(), 
-                                idx.into(), 
-                                cur_idx,
-                                "").unwrap()
-                        },
-                        _ => unsupported_type(ctx, fn_name, &[elem_ty, elem_ty]),
-                    };
-                    ctx.builder.build_store(res_idx, updated_idx).unwrap();
+                    let result = call_max(ctx, (elem_ty, accumulator), (elem_ty, elem));
                     ctx.builder.build_store(accumulator_addr, result).unwrap();
 
                     Ok(())
                 },
-                llvm_int64.const_int(1, false),
+                llvm_usize.const_int(1, false),
             )?;
 
-            match fn_name {
-                "np_argmin" | "np_argmax" => ctx.builder.build_load(res_idx, "").unwrap(),
-                "np_max" | "np_min" => ctx.builder.build_load(accumulator_addr, "").unwrap(),
-                _ => unreachable!()
-            }
+            let accumulator = ctx.builder.build_load(accumulator_addr, "").unwrap();
+            accumulator
         }
 
-        _ => unsupported_type(ctx, fn_name, &[a_ty])
+        _ => unsupported_type(ctx, FN_NAME, &[a_ty]),
     })
 }
 

nac3core/src/toplevel/builtins.rs

@@ -510,10 +510,7 @@ impl<'a> BuiltinBuilder<'a> {
 
             PrimDef::FunMin | PrimDef::FunMax => self.build_min_max_function(prim),
 
-            PrimDef::FunNpArgmin 
-            | PrimDef::FunNpArgmax
-            | PrimDef::FunNpMin 
-            | PrimDef::FunNpMax => self.build_np_max_min_function(prim),
+            PrimDef::FunNpMin | PrimDef::FunNpMax => self.build_np_min_max_function(prim),
 
             PrimDef::FunNpMinimum | PrimDef::FunNpMaximum => {
                 self.build_np_minimum_maximum_function(prim)
@@ -1558,42 +1555,39 @@ impl<'a> BuiltinBuilder<'a> {
         }
     }
 
-    /// Build the functions `np_max()`, `np_min()`, `np_argmax()` and `np_argmin()`
-    /// Calls `call_numpy_max_min` with the function name
-    fn build_np_max_min_function(&mut self, prim: PrimDef) -> TopLevelDef {
-        debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpArgmin, PrimDef::FunNpArgmax, PrimDef::FunNpMin, PrimDef::FunNpMax]);
+    /// Build the functions `np_min()` and `np_max()`.
+    fn build_np_min_max_function(&mut self, prim: PrimDef) -> TopLevelDef {
+        debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpMin, PrimDef::FunNpMax]);
 
-        let (var_map, ret_ty) = match prim {
-            PrimDef::FunNpArgmax | PrimDef::FunNpArgmin => {
-                (self.num_or_ndarray_var_map.clone(), self.primitives.int64)
-            },
-            PrimDef::FunNpMax | PrimDef::FunNpMin => {
-                let ret_ty = self.unifier.get_fresh_var(Some("R".into()), None);
-                let var_map = self
-                    .num_or_ndarray_var_map
-                    .clone()
-                    .into_iter()
-                    .chain(once((ret_ty.id, ret_ty.ty)))
-                    .collect::<IndexMap<_, _>>();
-                (var_map, ret_ty.ty)
-            },
-            _ => unreachable!()
-        };
+        let ret_ty = self.unifier.get_fresh_var(Some("R".into()), None);
+        let var_map = self
+            .num_or_ndarray_var_map
+            .clone()
+            .into_iter()
+            .chain(once((ret_ty.id, ret_ty.ty)))
+            .collect::<IndexMap<_, _>>();
 
         create_fn_by_codegen(
             self.unifier,
             &var_map,
             prim.name(),
-            ret_ty,
-            &[(self.num_or_ndarray_ty.ty, "a")],
+            ret_ty.ty,
+            &[(self.float_or_ndarray_ty.ty, "a")],
             Box::new(move |ctx, _, fun, args, generator| {
                 let a_ty = fun.0.args[0].ty;
                 let a = args[0].1.clone().to_basic_value_enum(ctx, generator, a_ty)?;
 
-                Ok(Some(builtin_fns::call_numpy_max_min(generator, ctx, (a_ty, a), &prim.name())?))
+                let func = match prim {
+                    PrimDef::FunNpMin => builtin_fns::call_numpy_min,
+                    PrimDef::FunNpMax => builtin_fns::call_numpy_max,
+                    _ => unreachable!(),
+                };
+
+                Ok(Some(func(generator, ctx, (a_ty, a))?))
             }),
         )
     }
+
     /// Build the functions `np_minimum()` and `np_maximum()`.
     fn build_np_minimum_maximum_function(&mut self, prim: PrimDef) -> TopLevelDef {
         debug_assert_prim_is_allowed(prim, &[PrimDef::FunNpMinimum, PrimDef::FunNpMaximum]);

nac3core/src/toplevel/helper.rs

@@ -62,11 +62,9 @@ pub enum PrimDef {
     FunMin,
     FunNpMin,
     FunNpMinimum,
-    FunNpArgmin,
     FunMax,
     FunNpMax,
     FunNpMaximum,
-    FunNpArgmax,
     FunAbs,
     FunNpIsNan,
     FunNpIsInf,
@@ -218,11 +216,9 @@ impl PrimDef {
             PrimDef::FunMin => fun("min", None),
             PrimDef::FunNpMin => fun("np_min", None),
             PrimDef::FunNpMinimum => fun("np_minimum", None),
-            PrimDef::FunNpArgmin => fun("np_argmin", None),
             PrimDef::FunMax => fun("max", None),
             PrimDef::FunNpMax => fun("np_max", None),
             PrimDef::FunNpMaximum => fun("np_maximum", None),
-            PrimDef::FunNpArgmax => fun("np_argmax", None),
             PrimDef::FunAbs => fun("abs", None),
             PrimDef::FunNpIsNan => fun("np_isnan", None),
             PrimDef::FunNpIsInf => fun("np_isinf", None),

nac3standalone/demo/src/ndarray.py

@@ -867,13 +867,6 @@ def test_ndarray_minimum_broadcast_rhs_scalar():
     output_ndarray_float_2(min_x_zeros)
     output_ndarray_float_2(min_x_ones)
 
-def test_ndarray_argmin():
-    x = np_array([[1., 2.], [3., 4.]])
-    y = np_argmin(x)
-
-    output_ndarray_float_2(x)
-    output_int64(y)
-
 def test_ndarray_max():
     x = np_identity(2)
     y = np_max(x)
@@ -917,13 +910,6 @@ def test_ndarray_maximum_broadcast_rhs_scalar():
     output_ndarray_float_2(max_x_zeros)
     output_ndarray_float_2(max_x_ones)
 
-def test_ndarray_argmax():
-    x = np_array([[1., 2.], [3., 4.]])
-    y = np_argmax(x)
-
-    output_ndarray_float_2(x)
-    output_int64(y)
-
 def test_ndarray_abs():
     x = np_identity(2)
     y = abs(x)
@@ -1538,13 +1524,11 @@ def run() -> int32:
     test_ndarray_minimum_broadcast()
     test_ndarray_minimum_broadcast_lhs_scalar()
     test_ndarray_minimum_broadcast_rhs_scalar()
-    test_ndarray_argmin()
     test_ndarray_max()
     test_ndarray_maximum()
     test_ndarray_maximum_broadcast()
     test_ndarray_maximum_broadcast_lhs_scalar()
     test_ndarray_maximum_broadcast_rhs_scalar()
-    test_ndarray_argmax()
     test_ndarray_abs()
     test_ndarray_isnan()
     test_ndarray_isinf()