1use std::cmp::Ordering;
4
5use runmat_accelerate_api::{GpuTensorHandle, GpuTensorStorage};
6use runmat_value::{
7 ComplexTensor, IntValue, IntegerStorage, LogicalArray, NumericScalar, Tensor, Value,
8};
9
10use crate::builtins::common::broadcast::{broadcast_shapes, BroadcastPlan};
11use crate::builtins::common::gpu_helpers;
12
13#[derive(Clone, Copy)]
14pub enum IntegerComparisonOp {
15 Eq,
16 Ne,
17 Lt,
18 Le,
19 Gt,
20 Ge,
21}
22
23#[derive(Debug)]
24pub enum IntegerComparisonError {
25 SizeMismatch,
26 Internal,
27}
28
29pub(crate) async fn try_gpu_equality_comparison(
30 lhs: &GpuTensorHandle,
31 rhs: &GpuTensorHandle,
32 operation: IntegerComparisonOp,
33) -> Option<crate::BuiltinResult<Value>> {
34 if lhs.device_id != rhs.device_id {
35 return None;
36 }
37 let provider = resolved_actual_owner(lhs)?;
38 let rhs_owner = resolved_actual_owner(rhs)?;
39 if !std::ptr::eq(provider, rhs_owner)
40 || runmat_accelerate_api::handle_precision(lhs)
41 != runmat_accelerate_api::handle_precision(rhs)
42 {
43 return None;
44 }
45 let result = match operation {
46 IntegerComparisonOp::Eq => provider.elem_eq(lhs, rhs).await,
47 IntegerComparisonOp::Ne => provider.elem_ne(lhs, rhs).await,
48 _ => unreachable!("equality helper only supports eq and ne"),
49 };
50 match result {
51 Ok(mut handle) if valid_equality_output(&handle, lhs, rhs, provider) => {
52 let provenance = [lhs, rhs]
53 .into_iter()
54 .filter_map(runmat_accelerate_api::handle_provenance)
55 .find(|provenance| {
56 *provenance == runmat_accelerate_api::GpuHandleProvenance::Explicit
57 })
58 .unwrap_or(runmat_accelerate_api::GpuHandleProvenance::Automatic);
59 runmat_accelerate_api::set_handle_provenance(&mut handle, provenance);
60 Some(Ok(gpu_helpers::logical_gpu_value(handle)))
61 }
62 Ok(handle) => {
63 free_rejected_gpu_handle(&handle, &[lhs, rhs]);
64 None
65 }
66 Err(_) => None,
67 }
68}
69
70pub(crate) fn select_comparison_output_source(
71 lhs: &Value,
72 rhs: &Value,
73 builtin: &str,
74) -> crate::BuiltinResult<Option<GpuTensorHandle>> {
75 gpu_helpers::select_resident_output_source(
76 [lhs, rhs].into_iter().filter_map(|value| match value {
77 Value::GpuTensor(handle) => Some(handle.clone()),
78 _ => None,
79 }),
80 builtin,
81 )
82}
83
84pub(crate) fn restore_explicit_comparison_result(
85 value: Value,
86 source: Option<&GpuTensorHandle>,
87 builtin: &str,
88) -> crate::BuiltinResult<Value> {
89 let Some(source) = source else {
90 return Ok(value);
91 };
92 let value = match value {
93 Value::Bool(bit) => Value::LogicalArray(
94 LogicalArray::new(vec![u8::from(bit)], vec![1, 1]).map_err(|error| {
95 crate::build_runtime_error(format!(
96 "{builtin}: invalid scalar logical result: {error}"
97 ))
98 .with_builtin(builtin)
99 .build()
100 })?,
101 ),
102 value => value,
103 };
104 let restored = gpu_helpers::restore_class_preserving_value(source, value, builtin)?;
105 if runmat_accelerate_api::handle_is_explicit(source) && !matches!(restored, Value::GpuTensor(_))
106 {
107 return Err(crate::build_runtime_error(format!(
108 "{builtin}: provider cannot preserve explicit gpuArray output residency"
109 ))
110 .with_builtin(builtin)
111 .with_identifier(format!("RunMat:{builtin}:GpuUploadFailed"))
112 .with_gpu_gather_retry(crate::GpuGatherRetry::Never)
113 .build());
114 }
115 Ok(restored)
116}
117
118pub(crate) async fn try_gpu_ordering_comparison(
121 lhs: &GpuTensorHandle,
122 rhs: &GpuTensorHandle,
123 operation: IntegerComparisonOp,
124) -> Option<crate::BuiltinResult<Value>> {
125 if lhs.device_id != rhs.device_id {
126 return None;
127 }
128 let provider = resolved_actual_owner(lhs)?;
129 let rhs_owner = resolved_actual_owner(rhs)?;
130 if !std::ptr::eq(provider, rhs_owner) {
131 return None;
132 }
133 let mut temporary_lhs = None;
134 let mut temporary_rhs = None;
135 let lhs_real =
136 if runmat_accelerate_api::handle_storage(lhs) == GpuTensorStorage::ComplexInterleaved {
137 match provider.unary_real(lhs).await {
138 Ok(handle) if valid_real_projection(&handle, lhs, provider) => {
139 temporary_lhs = Some(handle);
140 temporary_lhs.as_ref().expect("temporary lhs")
141 }
142 Ok(handle) => {
143 free_rejected_gpu_handle(&handle, &[lhs, rhs]);
144 return None;
145 }
146 Err(_) => return None,
147 }
148 } else {
149 lhs
150 };
151 let rhs_real =
152 if runmat_accelerate_api::handle_storage(rhs) == GpuTensorStorage::ComplexInterleaved {
153 match provider.unary_real(rhs).await {
154 Ok(handle) if valid_real_projection(&handle, rhs, provider) => {
155 temporary_rhs = Some(handle);
156 temporary_rhs.as_ref().expect("temporary rhs")
157 }
158 Ok(handle) => {
159 free_rejected_gpu_handle(&handle, &[lhs, rhs, lhs_real]);
160 if let Some(handle) = temporary_lhs.as_ref() {
161 let _ = provider.free(handle);
162 }
163 return None;
164 }
165 Err(_) => {
166 if let Some(handle) = temporary_lhs.as_ref() {
167 let _ = provider.free(handle);
168 }
169 return None;
170 }
171 }
172 } else {
173 rhs
174 };
175 if runmat_accelerate_api::handle_precision(lhs_real)
176 != runmat_accelerate_api::handle_precision(rhs_real)
177 {
178 if let Some(handle) = temporary_lhs.as_ref() {
179 let _ = provider.free(handle);
180 }
181 if let Some(handle) = temporary_rhs.as_ref() {
182 let _ = provider.free(handle);
183 }
184 return None;
185 }
186 let result = match operation {
187 IntegerComparisonOp::Lt => provider.elem_lt(lhs_real, rhs_real).await,
188 IntegerComparisonOp::Le => provider.elem_le(lhs_real, rhs_real).await,
189 IntegerComparisonOp::Gt => provider.elem_gt(lhs_real, rhs_real).await,
190 IntegerComparisonOp::Ge => provider.elem_ge(lhs_real, rhs_real).await,
191 IntegerComparisonOp::Eq | IntegerComparisonOp::Ne => {
192 unreachable!("resident complex ordering helper only supports lt/le/gt/ge")
193 }
194 };
195 let result = match result {
196 Ok(mut handle) if valid_ordering_output(&handle, lhs_real, rhs_real, provider) => {
197 let provenance = [lhs, rhs]
198 .into_iter()
199 .filter_map(runmat_accelerate_api::handle_provenance)
200 .find(|provenance| {
201 *provenance == runmat_accelerate_api::GpuHandleProvenance::Explicit
202 })
203 .unwrap_or(runmat_accelerate_api::GpuHandleProvenance::Automatic);
204 runmat_accelerate_api::set_handle_provenance(&mut handle, provenance);
205 Some(gpu_helpers::logical_gpu_value(handle))
206 }
207 Ok(handle) => {
208 free_rejected_gpu_handle(&handle, &[lhs, rhs, lhs_real, rhs_real]);
209 None
210 }
211 Err(_) => None,
212 };
213 if let Some(handle) = temporary_lhs.as_ref() {
214 let _ = provider.free(handle);
215 }
216 if let Some(handle) = temporary_rhs.as_ref() {
217 let _ = provider.free(handle);
218 }
219 result.map(Ok)
220}
221
222fn resolved_actual_owner(
223 handle: &GpuTensorHandle,
224) -> Option<&'static dyn runmat_accelerate_api::AccelProvider> {
225 runmat_accelerate_api::provider_for_handle(handle)
226 .filter(|owner| owner.device_id() == handle.device_id)
227}
228
229fn gpu_handles_alias(lhs: &GpuTensorHandle, rhs: &GpuTensorHandle) -> bool {
230 lhs.device_id == rhs.device_id && lhs.buffer_id == rhs.buffer_id
231}
232
233fn valid_real_projection(
234 output: &GpuTensorHandle,
235 input: &GpuTensorHandle,
236 provider: &'static dyn runmat_accelerate_api::AccelProvider,
237) -> bool {
238 output.shape == input.shape
239 && output.device_id == input.device_id
240 && !gpu_handles_alias(output, input)
241 && runmat_accelerate_api::handle_storage(output) == GpuTensorStorage::Real
242 && runmat_accelerate_api::handle_precision(output)
243 == runmat_accelerate_api::handle_precision(input)
244 && runmat_accelerate_api::handle_integer_type(output).is_none()
245 && !runmat_accelerate_api::handle_is_logical(output)
246 && resolved_actual_owner(output).is_some_and(|owner| std::ptr::eq(owner, provider))
247}
248
249fn valid_ordering_output(
250 output: &GpuTensorHandle,
251 lhs: &GpuTensorHandle,
252 rhs: &GpuTensorHandle,
253 provider: &'static dyn runmat_accelerate_api::AccelProvider,
254) -> bool {
255 let expected_shape = broadcast_shapes("comparison", &lhs.shape, &rhs.shape).ok();
256 expected_shape.as_deref() == Some(output.shape.as_slice())
257 && output.device_id == lhs.device_id
258 && !gpu_handles_alias(output, lhs)
259 && !gpu_handles_alias(output, rhs)
260 && runmat_accelerate_api::handle_storage(output) == GpuTensorStorage::Real
261 && runmat_accelerate_api::handle_precision(output)
262 == runmat_accelerate_api::handle_precision(lhs)
263 && runmat_accelerate_api::handle_integer_type(output).is_none()
264 && resolved_actual_owner(output).is_some_and(|owner| std::ptr::eq(owner, provider))
265}
266
267fn valid_equality_output(
268 output: &GpuTensorHandle,
269 lhs: &GpuTensorHandle,
270 rhs: &GpuTensorHandle,
271 provider: &'static dyn runmat_accelerate_api::AccelProvider,
272) -> bool {
273 let expected_shape = broadcast_shapes("comparison", &lhs.shape, &rhs.shape).ok();
274 expected_shape.as_deref() == Some(output.shape.as_slice())
275 && output.device_id == lhs.device_id
276 && !gpu_handles_alias(output, lhs)
277 && !gpu_handles_alias(output, rhs)
278 && runmat_accelerate_api::handle_storage(output) == GpuTensorStorage::Real
279 && runmat_accelerate_api::handle_precision(output)
280 == runmat_accelerate_api::handle_precision(lhs)
281 && runmat_accelerate_api::handle_integer_type(output).is_none()
282 && resolved_actual_owner(output).is_some_and(|owner| std::ptr::eq(owner, provider))
283}
284
285fn free_rejected_gpu_handle(handle: &GpuTensorHandle, protected: &[&GpuTensorHandle]) {
286 if protected
287 .iter()
288 .any(|protected| gpu_handles_alias(handle, protected))
289 {
290 return;
291 }
292 if let Some(owner) = resolved_actual_owner(handle) {
293 let _ = owner.free(handle);
294 }
295}
296
297pub fn try_integer_comparison(
301 lhs: &Value,
302 rhs: &Value,
303 operation: IntegerComparisonOp,
304) -> Result<Option<Value>, IntegerComparisonError> {
305 let lhs_integer = integer_operand(lhs);
306 let rhs_integer = integer_operand(rhs);
307 let result = match (lhs_integer, rhs_integer) {
308 (None, None) => return Ok(None),
309 (Some(lhs), Some(rhs)) => compare_integer_operands(&lhs, &rhs, operation)?,
310 (Some(lhs), None) => {
311 let Some(rhs) = numeric_operand(rhs) else {
312 return Ok(None);
313 };
314 compare_integer_numeric(&lhs, &rhs, true, operation)?
315 }
316 (None, Some(rhs)) => {
317 let Some(lhs) = numeric_operand(lhs) else {
318 return Ok(None);
319 };
320 compare_integer_numeric(&rhs, &lhs, false, operation)?
321 }
322 };
323 Ok(Some(result))
324}
325
326pub fn try_real_ordering_comparison(
330 lhs: &Value,
331 rhs: &Value,
332 operation: IntegerComparisonOp,
333) -> Result<Option<Value>, IntegerComparisonError> {
334 let Some(lhs) = real_operand(lhs) else {
335 return Ok(None);
336 };
337 let Some(rhs) = real_operand(rhs) else {
338 return Ok(None);
339 };
340 let plan = BroadcastPlan::new(&lhs.shape, &rhs.shape)
341 .map_err(|_| IntegerComparisonError::SizeMismatch)?;
342 let mut data = Vec::with_capacity(plan.len());
343 for (_, lhs_index, rhs_index) in plan.iter() {
344 let ordering = compare_real_values(lhs.value_at(lhs_index), rhs.value_at(rhs_index));
345 data.push(matches_optional_relation(ordering, operation) as u8);
346 }
347 logical_result(data, plan.output_shape().to_vec()).map(Some)
348}
349
350pub(crate) fn try_complex_integer_equality_comparison(
353 lhs: &Value,
354 rhs: &Value,
355 operation: IntegerComparisonOp,
356) -> Result<Option<Value>, IntegerComparisonError> {
357 debug_assert!(matches!(
358 operation,
359 IntegerComparisonOp::Eq | IntegerComparisonOp::Ne
360 ));
361 let lhs_complex = complex_operand(lhs);
362 let rhs_complex = complex_operand(rhs);
363 let result = match (lhs_complex, rhs_complex) {
364 (Some(lhs), Some(rhs)) if lhs.has_integer_storage() || rhs.has_integer_storage() => {
365 compare_complex_operands(&lhs, &rhs, operation)?
366 }
367 (Some(lhs), None) => {
368 let Some(rhs) = real_operand(rhs) else {
369 return Ok(None);
370 };
371 if !lhs.has_integer_storage() && !rhs.has_integer_storage() {
372 return Ok(None);
373 }
374 compare_complex_real(&lhs, &rhs, operation)?
375 }
376 (None, Some(rhs)) => {
377 let Some(lhs) = real_operand(lhs) else {
378 return Ok(None);
379 };
380 if !rhs.has_integer_storage() && !lhs.has_integer_storage() {
381 return Ok(None);
382 }
383 compare_complex_real(&rhs, &lhs, operation)?
384 }
385 _ => return Ok(None),
386 };
387 Ok(Some(result))
388}
389
390pub fn try_complex_ordering_comparison(
394 lhs: &Value,
395 rhs: &Value,
396 operation: IntegerComparisonOp,
397) -> Result<Option<Value>, IntegerComparisonError> {
398 debug_assert!(matches!(
399 operation,
400 IntegerComparisonOp::Lt
401 | IntegerComparisonOp::Le
402 | IntegerComparisonOp::Gt
403 | IntegerComparisonOp::Ge
404 ));
405 let lhs_complex = complex_operand(lhs);
406 let rhs_complex = complex_operand(rhs);
407 let result = match (lhs_complex, rhs_complex) {
408 (Some(lhs), Some(rhs)) => compare_complex_real_components(&lhs, &rhs, operation)?,
409 (Some(lhs), None) => {
410 let Some(rhs) = real_operand(rhs) else {
411 return Ok(None);
412 };
413 compare_complex_real_ordering(&lhs, &rhs, true, operation)?
414 }
415 (None, Some(rhs)) => {
416 let Some(lhs) = real_operand(lhs) else {
417 return Ok(None);
418 };
419 compare_complex_real_ordering(&rhs, &lhs, false, operation)?
420 }
421 (None, None) => return Ok(None),
422 };
423 Ok(Some(result))
424}
425
426fn compare_integer_operands(
427 lhs: &IntegerOperand<'_>,
428 rhs: &IntegerOperand<'_>,
429 operation: IntegerComparisonOp,
430) -> Result<Value, IntegerComparisonError> {
431 let plan = BroadcastPlan::new(&lhs.shape, &rhs.shape)
432 .map_err(|_| IntegerComparisonError::SizeMismatch)?;
433 let mut data = Vec::with_capacity(plan.len());
434 for (_, lhs_index, rhs_index) in plan.iter() {
435 let ordering = compare_integer_values(lhs.value_at(lhs_index), rhs.value_at(rhs_index));
436 data.push(matches_relation(ordering, operation) as u8);
437 }
438 logical_result(data, plan.output_shape().to_vec())
439}
440
441fn compare_integer_numeric(
442 integer: &IntegerOperand<'_>,
443 numeric: &NumericOperand<'_>,
444 integer_is_left: bool,
445 operation: IntegerComparisonOp,
446) -> Result<Value, IntegerComparisonError> {
447 let plan = BroadcastPlan::new(&integer.shape, numeric.shape())
448 .map_err(|_| IntegerComparisonError::SizeMismatch)?;
449 let mut data = Vec::with_capacity(plan.len());
450 for (_, integer_index, numeric_index) in plan.iter() {
451 let ordering = integer_f64_order(
452 integer.value_at(integer_index),
453 numeric.value_at(numeric_index),
454 );
455 let ordering = if integer_is_left {
456 ordering
457 } else {
458 ordering.map(Ordering::reverse)
459 };
460 data.push(matches_optional_relation(ordering, operation) as u8);
461 }
462 logical_result(data, plan.output_shape().to_vec())
463}
464
465fn logical_result(data: Vec<u8>, shape: Vec<usize>) -> Result<Value, IntegerComparisonError> {
466 if data.len() == 1 {
467 return Ok(Value::Bool(data[0] != 0));
468 }
469 Ok(Value::LogicalArray(
470 LogicalArray::new(data, shape).map_err(|_| IntegerComparisonError::Internal)?,
471 ))
472}
473
474struct ComplexOperand<'a> {
475 source: ComplexSource<'a>,
476 shape: Vec<usize>,
477}
478
479impl ComplexOperand<'_> {
480 fn has_integer_storage(&self) -> bool {
481 match self.source {
482 ComplexSource::Scalar(_, _) => false,
483 ComplexSource::Dense(tensor) => tensor.integer_storage().is_some(),
484 }
485 }
486
487 fn real_imag_at(&self, index: usize) -> ComplexValue {
488 match self.source {
489 ComplexSource::Scalar(real, imag) => ComplexValue::Float(real, imag),
490 ComplexSource::Dense(tensor) => complex_value_from_scalars(
491 tensor
492 .numeric_value_at(index)
493 .expect("complex tensor storage must match shape"),
494 ),
495 }
496 }
497}
498
499enum ComplexSource<'a> {
500 Scalar(f64, f64),
501 Dense(&'a ComplexTensor),
502}
503
504enum ComplexValue {
505 Integer(IntValue, IntValue),
506 Float(f64, f64),
507}
508
509enum RealValue {
510 Integer(IntValue),
511 Float(f64),
512}
513
514fn real_value_from_scalar(value: NumericScalar) -> RealValue {
515 match value {
516 NumericScalar::F64(value) => RealValue::Float(value),
517 NumericScalar::F32(value) => RealValue::Float(f64::from(value)),
518 integer => RealValue::Integer(
519 integer
520 .into_int_value()
521 .expect("non-floating numeric scalar must be integer"),
522 ),
523 }
524}
525
526fn complex_value_from_scalars((real, imag): (NumericScalar, NumericScalar)) -> ComplexValue {
527 match (real.into_int_value(), imag.into_int_value()) {
528 (Some(real), Some(imag)) => ComplexValue::Integer(real, imag),
529 (None, None) => {
530 ComplexValue::Float(floating_scalar_to_f64(real), floating_scalar_to_f64(imag))
531 }
532 _ => unreachable!("complex storage components must use the same numeric domain"),
533 }
534}
535
536fn floating_scalar_to_f64(value: NumericScalar) -> f64 {
537 match value {
538 NumericScalar::F64(value) => value,
539 NumericScalar::F32(value) => f64::from(value),
540 _ => unreachable!("expected floating numeric scalar"),
541 }
542}
543
544struct RealOperand<'a> {
545 source: RealSource<'a>,
546 shape: Vec<usize>,
547}
548
549impl RealOperand<'_> {
550 fn has_integer_storage(&self) -> bool {
551 match self.source {
552 RealSource::ScalarInteger(_) => true,
553 RealSource::Dense(tensor) => tensor.integer_storage().is_some(),
554 RealSource::ScalarFloat(_) | RealSource::Logical { .. } | RealSource::Char(_) => false,
555 }
556 }
557
558 fn value_at(&self, index: usize) -> RealValue {
559 match self.source {
560 RealSource::ScalarInteger(ref value) => RealValue::Integer(value.clone()),
561 RealSource::ScalarFloat(value) => RealValue::Float(value),
562 RealSource::Dense(tensor) => real_value_from_scalar(
563 tensor
564 .numeric_value_at(index)
565 .expect("tensor storage must match shape"),
566 ),
567 RealSource::Logical { data } => RealValue::Float(f64::from(data[index] != 0)),
568 RealSource::Char(array) => {
569 let row = index % array.rows;
570 let column = index / array.rows;
571 RealValue::Float(f64::from(array.data[row * array.cols + column] as u32))
572 }
573 }
574 }
575}
576
577enum RealSource<'a> {
578 ScalarInteger(IntValue),
579 ScalarFloat(f64),
580 Dense(&'a Tensor),
581 Logical { data: &'a [u8] },
582 Char(&'a runmat_value::CharArray),
583}
584
585fn compare_complex_operands(
586 lhs: &ComplexOperand<'_>,
587 rhs: &ComplexOperand<'_>,
588 operation: IntegerComparisonOp,
589) -> Result<Value, IntegerComparisonError> {
590 let plan = BroadcastPlan::new(&lhs.shape, &rhs.shape)
591 .map_err(|_| IntegerComparisonError::SizeMismatch)?;
592 let mut data = Vec::with_capacity(plan.len());
593 for (_, lhs_index, rhs_index) in plan.iter() {
594 let matches =
595 complex_values_equal(lhs.real_imag_at(lhs_index), rhs.real_imag_at(rhs_index));
596 data.push(matches_relation_bool(matches, operation) as u8);
597 }
598 logical_result(data, plan.output_shape().to_vec())
599}
600
601fn compare_complex_real_components(
602 lhs: &ComplexOperand<'_>,
603 rhs: &ComplexOperand<'_>,
604 operation: IntegerComparisonOp,
605) -> Result<Value, IntegerComparisonError> {
606 let plan = BroadcastPlan::new(&lhs.shape, &rhs.shape)
607 .map_err(|_| IntegerComparisonError::SizeMismatch)?;
608 let mut data = Vec::with_capacity(plan.len());
609 for (_, lhs_index, rhs_index) in plan.iter() {
610 let ordering = compare_real_values(
611 complex_real_component(lhs.real_imag_at(lhs_index)),
612 complex_real_component(rhs.real_imag_at(rhs_index)),
613 );
614 data.push(matches_optional_relation(ordering, operation) as u8);
615 }
616 logical_result(data, plan.output_shape().to_vec())
617}
618
619fn compare_complex_real_ordering(
620 complex: &ComplexOperand<'_>,
621 real: &RealOperand<'_>,
622 complex_is_left: bool,
623 operation: IntegerComparisonOp,
624) -> Result<Value, IntegerComparisonError> {
625 let plan = BroadcastPlan::new(&complex.shape, &real.shape)
626 .map_err(|_| IntegerComparisonError::SizeMismatch)?;
627 let mut data = Vec::with_capacity(plan.len());
628 for (_, complex_index, real_index) in plan.iter() {
629 let ordering = compare_real_values(
630 complex_real_component(complex.real_imag_at(complex_index)),
631 real.value_at(real_index),
632 );
633 let ordering = if complex_is_left {
634 ordering
635 } else {
636 ordering.map(Ordering::reverse)
637 };
638 data.push(matches_optional_relation(ordering, operation) as u8);
639 }
640 logical_result(data, plan.output_shape().to_vec())
641}
642
643fn complex_real_component(value: ComplexValue) -> RealValue {
644 match value {
645 ComplexValue::Integer(real, _) => RealValue::Integer(real),
646 ComplexValue::Float(real, _) => RealValue::Float(real),
647 }
648}
649
650fn compare_real_values(lhs: RealValue, rhs: RealValue) -> Option<Ordering> {
651 match (lhs, rhs) {
652 (RealValue::Integer(lhs), RealValue::Integer(rhs)) => {
653 Some(compare_integer_values(lhs, rhs))
654 }
655 (RealValue::Integer(lhs), RealValue::Float(rhs)) => integer_f64_order(lhs, rhs),
656 (RealValue::Float(lhs), RealValue::Integer(rhs)) => {
657 integer_f64_order(rhs, lhs).map(Ordering::reverse)
658 }
659 (RealValue::Float(lhs), RealValue::Float(rhs)) => lhs.partial_cmp(&rhs),
660 }
661}
662
663pub(crate) fn compare_numeric_scalars_exact(
668 lhs: NumericScalar,
669 rhs: NumericScalar,
670) -> Option<Ordering> {
671 fn real(value: NumericScalar) -> RealValue {
672 match value {
673 NumericScalar::F64(value) => RealValue::Float(value),
674 NumericScalar::F32(value) => RealValue::Float(f64::from(value)),
675 value => RealValue::Integer(
676 value
677 .into_int_value()
678 .expect("nonfloating NumericScalar must contain an integer"),
679 ),
680 }
681 }
682 compare_real_values(real(lhs), real(rhs))
683}
684
685fn compare_complex_real(
686 complex: &ComplexOperand<'_>,
687 real: &RealOperand<'_>,
688 operation: IntegerComparisonOp,
689) -> Result<Value, IntegerComparisonError> {
690 let plan = BroadcastPlan::new(&complex.shape, &real.shape)
691 .map_err(|_| IntegerComparisonError::SizeMismatch)?;
692 let mut data = Vec::with_capacity(plan.len());
693 for (_, complex_index, real_index) in plan.iter() {
694 let matches = complex_value_equals_real(
695 complex.real_imag_at(complex_index),
696 real.value_at(real_index),
697 );
698 data.push(matches_relation_bool(matches, operation) as u8);
699 }
700 logical_result(data, plan.output_shape().to_vec())
701}
702
703fn complex_values_equal(lhs: ComplexValue, rhs: ComplexValue) -> bool {
704 match (lhs, rhs) {
705 (ComplexValue::Integer(lhs_real, lhs_imag), ComplexValue::Integer(rhs_real, rhs_imag)) => {
706 compare_integer_values(lhs_real, rhs_real) == Ordering::Equal
707 && compare_integer_values(lhs_imag, rhs_imag) == Ordering::Equal
708 }
709 (ComplexValue::Integer(real, imag), ComplexValue::Float(rhs_real, rhs_imag)) => {
710 integer_value_equals_f64(real, rhs_real) && integer_value_equals_f64(imag, rhs_imag)
711 }
712 (ComplexValue::Float(lhs_real, lhs_imag), ComplexValue::Integer(real, imag)) => {
713 integer_value_equals_f64(real, lhs_real) && integer_value_equals_f64(imag, lhs_imag)
714 }
715 (ComplexValue::Float(lhs_real, lhs_imag), ComplexValue::Float(rhs_real, rhs_imag)) => {
716 lhs_real == rhs_real && lhs_imag == rhs_imag
717 }
718 }
719}
720
721fn complex_value_equals_real(complex: ComplexValue, real: RealValue) -> bool {
722 match (complex, real) {
723 (ComplexValue::Integer(complex_real, complex_imag), RealValue::Integer(real)) => {
724 complex_imag.is_zero() && compare_integer_values(complex_real, real) == Ordering::Equal
725 }
726 (ComplexValue::Integer(complex_real, complex_imag), RealValue::Float(real)) => {
727 complex_imag.is_zero() && integer_value_equals_f64(complex_real, real)
728 }
729 (ComplexValue::Float(complex_real, complex_imag), RealValue::Integer(real)) => {
730 complex_imag == 0.0 && integer_value_equals_f64(real, complex_real)
731 }
732 (ComplexValue::Float(complex_real, complex_imag), RealValue::Float(real)) => {
733 complex_imag == 0.0 && complex_real == real
734 }
735 }
736}
737
738fn integer_value_equals_f64(integer: IntValue, float: f64) -> bool {
739 integer_f64_order(integer, float) == Some(Ordering::Equal)
740}
741
742fn matches_relation_bool(matches: bool, operation: IntegerComparisonOp) -> bool {
743 match operation {
744 IntegerComparisonOp::Eq => matches,
745 IntegerComparisonOp::Ne => !matches,
746 IntegerComparisonOp::Lt
747 | IntegerComparisonOp::Le
748 | IntegerComparisonOp::Gt
749 | IntegerComparisonOp::Ge => {
750 unreachable!("complex integer equality helper only supports eq/ne")
751 }
752 }
753}
754
755fn complex_operand(value: &Value) -> Option<ComplexOperand<'_>> {
756 match value {
757 Value::Complex(real, imag) => Some(ComplexOperand {
758 source: ComplexSource::Scalar(*real, *imag),
759 shape: vec![1, 1],
760 }),
761 Value::ComplexTensor(tensor) => Some(complex_tensor_operand(tensor)),
762 _ => None,
763 }
764}
765
766fn complex_tensor_operand(tensor: &ComplexTensor) -> ComplexOperand<'_> {
767 ComplexOperand {
768 source: ComplexSource::Dense(tensor),
769 shape: tensor.shape.clone(),
770 }
771}
772
773struct IntegerOperand<'a> {
774 storage: IntegerStorageRef<'a>,
775 shape: Vec<usize>,
776}
777
778impl IntegerOperand<'_> {
779 fn value_at(&self, index: usize) -> IntValue {
780 self.storage.value_at(index)
781 }
782}
783
784enum IntegerStorageRef<'a> {
785 Scalar(&'a IntValue),
786 Array(&'a IntegerStorage),
787}
788
789impl IntegerStorageRef<'_> {
790 fn value_at(&self, index: usize) -> IntValue {
791 match self {
792 Self::Scalar(value) => (*value).clone(),
793 Self::Array(storage) => storage_value(storage, index),
794 }
795 }
796}
797
798fn integer_operand(value: &Value) -> Option<IntegerOperand<'_>> {
799 match value {
800 Value::Int(value) => Some(IntegerOperand {
801 storage: IntegerStorageRef::Scalar(value),
802 shape: vec![1, 1],
803 }),
804 Value::Tensor(tensor) => tensor.integer_storage().map(|storage| IntegerOperand {
805 storage: IntegerStorageRef::Array(storage),
806 shape: tensor.shape.clone(),
807 }),
808 _ => None,
809 }
810}
811
812fn real_operand(value: &Value) -> Option<RealOperand<'_>> {
813 match value {
814 Value::Int(value) => Some(RealOperand {
815 source: RealSource::ScalarInteger(value.clone()),
816 shape: vec![1, 1],
817 }),
818 Value::Num(value) => Some(RealOperand {
819 source: RealSource::ScalarFloat(*value),
820 shape: vec![1, 1],
821 }),
822 Value::Bool(value) => Some(RealOperand {
823 source: RealSource::ScalarFloat(if *value { 1.0 } else { 0.0 }),
824 shape: vec![1, 1],
825 }),
826 Value::Tensor(tensor) => Some(RealOperand {
827 source: RealSource::Dense(tensor),
828 shape: tensor.shape.clone(),
829 }),
830 Value::LogicalArray(array) => Some(RealOperand {
831 source: RealSource::Logical { data: &array.data },
832 shape: array.shape.clone(),
833 }),
834 Value::CharArray(array) => Some(RealOperand {
835 source: RealSource::Char(array),
836 shape: vec![array.rows, array.cols],
837 }),
838 _ => None,
839 }
840}
841
842pub(crate) fn compare_integer_values(lhs: IntValue, rhs: IntValue) -> Ordering {
843 match (signed_value(&lhs), signed_value(&rhs)) {
844 (Some(lhs), Some(rhs)) => lhs.cmp(&rhs),
845 (None, None) => unsigned_value(&lhs).cmp(&unsigned_value(&rhs)),
846 (Some(lhs), None) => {
847 if lhs < 0 {
848 Ordering::Less
849 } else {
850 (lhs as u64).cmp(&unsigned_value(&rhs))
851 }
852 }
853 (None, Some(rhs)) => {
854 if rhs < 0 {
855 Ordering::Greater
856 } else {
857 unsigned_value(&lhs).cmp(&(rhs as u64))
858 }
859 }
860 }
861}
862
863pub(crate) fn integer_f64_order(integer: IntValue, float: f64) -> Option<Ordering> {
864 if float.is_nan() {
865 return None;
866 }
867 if float == f64::INFINITY {
868 return Some(Ordering::Less);
869 }
870 if float == f64::NEG_INFINITY {
871 return Some(Ordering::Greater);
872 }
873
874 const MIN_I64: f64 = -9_223_372_036_854_775_808.0;
875 const U64_EXCLUSIVE_UPPER: f64 = 18_446_744_073_709_551_616.0;
876 if float < MIN_I64 {
877 return Some(Ordering::Greater);
878 }
879 if float >= U64_EXCLUSIVE_UPPER {
880 return Some(Ordering::Less);
881 }
882
883 let integer = integer_as_i128(&integer);
884 let truncated = float as i128;
885 let ordering = integer.cmp(&truncated);
886 if float.fract() == 0.0 {
887 return Some(ordering);
888 }
889 Some(if float.is_sign_positive() {
890 if ordering == Ordering::Greater {
891 Ordering::Greater
892 } else {
893 Ordering::Less
894 }
895 } else if ordering == Ordering::Less {
896 Ordering::Less
897 } else {
898 Ordering::Greater
899 })
900}
901
902fn integer_as_i128(value: &IntValue) -> i128 {
903 match value {
904 IntValue::I8(value) => i128::from(*value),
905 IntValue::I16(value) => i128::from(*value),
906 IntValue::I32(value) => i128::from(*value),
907 IntValue::I64(value) => i128::from(*value),
908 IntValue::U8(value) => i128::from(*value),
909 IntValue::U16(value) => i128::from(*value),
910 IntValue::U32(value) => i128::from(*value),
911 IntValue::U64(value) => i128::from(*value),
912 }
913}
914
915enum NumericOperand<'a> {
916 Scalar(f64),
917 Dense(&'a Tensor),
918 Logical(&'a [u8], &'a [usize]),
919}
920
921impl NumericOperand<'_> {
922 fn shape(&self) -> &[usize] {
923 match self {
924 Self::Scalar(_) => &[1, 1],
925 Self::Dense(tensor) => &tensor.shape,
926 Self::Logical(_, shape) => shape,
927 }
928 }
929
930 fn value_at(&self, index: usize) -> f64 {
931 match self {
932 Self::Scalar(value) => *value,
933 Self::Dense(tensor) => match tensor
934 .numeric_value_at(index)
935 .expect("tensor storage must match shape")
936 {
937 NumericScalar::F64(value) => value,
938 NumericScalar::F32(value) => f64::from(value),
939 _ => unreachable!("integer tensors use the exact integer operand path"),
940 },
941 Self::Logical(data, _) => f64::from(data[index] != 0),
942 }
943 }
944}
945
946fn numeric_operand(value: &Value) -> Option<NumericOperand<'_>> {
947 match value {
948 Value::Num(value) => Some(NumericOperand::Scalar(*value)),
949 Value::Bool(value) => Some(NumericOperand::Scalar(if *value { 1.0 } else { 0.0 })),
950 Value::Tensor(tensor) if tensor.integer_storage().is_none() => {
951 Some(NumericOperand::Dense(tensor))
952 }
953 Value::LogicalArray(array) => Some(NumericOperand::Logical(&array.data, &array.shape)),
954 _ => None,
955 }
956}
957
958fn signed_value(value: &IntValue) -> Option<i64> {
959 match value {
960 IntValue::I8(value) => Some(*value as i64),
961 IntValue::I16(value) => Some(*value as i64),
962 IntValue::I32(value) => Some(*value as i64),
963 IntValue::I64(value) => Some(*value),
964 IntValue::U8(_) | IntValue::U16(_) | IntValue::U32(_) | IntValue::U64(_) => None,
965 }
966}
967
968fn unsigned_value(value: &IntValue) -> u64 {
969 match value {
970 IntValue::U8(value) => *value as u64,
971 IntValue::U16(value) => *value as u64,
972 IntValue::U32(value) => *value as u64,
973 IntValue::U64(value) => *value,
974 IntValue::I8(_) | IntValue::I16(_) | IntValue::I32(_) | IntValue::I64(_) => {
975 unreachable!("unsigned conversion is only used for unsigned integer values")
976 }
977 }
978}
979
980pub(crate) fn storage_value(storage: &IntegerStorage, index: usize) -> IntValue {
981 match storage {
982 IntegerStorage::I8(values) => IntValue::I8(values[index]),
983 IntegerStorage::I16(values) => IntValue::I16(values[index]),
984 IntegerStorage::I32(values) => IntValue::I32(values[index]),
985 IntegerStorage::I64(values) => IntValue::I64(values[index]),
986 IntegerStorage::U8(values) => IntValue::U8(values[index]),
987 IntegerStorage::U16(values) => IntValue::U16(values[index]),
988 IntegerStorage::U32(values) => IntValue::U32(values[index]),
989 IntegerStorage::U64(values) => IntValue::U64(values[index]),
990 }
991}
992
993pub(crate) fn matches_relation(ordering: Ordering, operation: IntegerComparisonOp) -> bool {
994 match operation {
995 IntegerComparisonOp::Eq => ordering == Ordering::Equal,
996 IntegerComparisonOp::Ne => ordering != Ordering::Equal,
997 IntegerComparisonOp::Lt => ordering == Ordering::Less,
998 IntegerComparisonOp::Le => ordering != Ordering::Greater,
999 IntegerComparisonOp::Gt => ordering == Ordering::Greater,
1000 IntegerComparisonOp::Ge => ordering != Ordering::Less,
1001 }
1002}
1003
1004pub(crate) fn matches_optional_relation(
1005 ordering: Option<Ordering>,
1006 operation: IntegerComparisonOp,
1007) -> bool {
1008 match ordering {
1009 Some(ordering) => matches_relation(ordering, operation),
1010 None => matches!(operation, IntegerComparisonOp::Ne),
1011 }
1012}
1013
1014#[cfg(test)]
1015mod tests {
1016 use super::*;
1017 use crate::builtins::common::test_support;
1018 use runmat_value::ComplexStorage;
1019
1020 fn array(storage: IntegerStorage, shape: Vec<usize>) -> Value {
1021 Value::Tensor(runmat_value::Tensor::new_integer(storage, shape).expect("integer tensor"))
1022 }
1023
1024 #[test]
1025 fn compares_signed_unsigned_and_uint64_exactly() {
1026 let lhs = Value::Int(IntValue::U64(u64::MAX));
1027 let rhs = Value::Int(IntValue::I64(i64::MAX));
1028 assert_eq!(
1029 try_integer_comparison(&lhs, &rhs, IntegerComparisonOp::Gt).expect("comparison"),
1030 Some(Value::Bool(true))
1031 );
1032 assert_eq!(
1033 try_integer_comparison(
1034 &Value::Int(IntValue::I8(-1)),
1035 &Value::Int(IntValue::U8(0)),
1036 IntegerComparisonOp::Lt,
1037 )
1038 .expect("comparison"),
1039 Some(Value::Bool(true))
1040 );
1041 }
1042
1043 #[test]
1044 fn broadcasts_exact_integer_arrays_for_all_relations() {
1045 let lhs = array(IntegerStorage::U64(vec![0, u64::MAX]), vec![2, 1]);
1046 let rhs = array(IntegerStorage::I64(vec![0, 1, i64::MAX]), vec![1, 3]);
1047 let result = try_integer_comparison(&lhs, &rhs, IntegerComparisonOp::Ge)
1048 .expect("comparison")
1049 .expect("integer path");
1050 assert_eq!(
1051 result,
1052 Value::LogicalArray(
1053 LogicalArray::new(vec![1, 1, 0, 1, 0, 1], vec![2, 3]).expect("logical result")
1054 )
1055 );
1056 }
1057
1058 #[test]
1059 fn compares_integer_storage_to_scalar_double_without_64_bit_loss() {
1060 let exact = Value::Int(IntValue::U64((1_u64 << 53) + 1));
1061 let rounded = Value::Num((1_u64 << 53) as f64);
1062 assert_eq!(
1063 try_integer_comparison(&exact, &rounded, IntegerComparisonOp::Eq).expect("comparison"),
1064 Some(Value::Bool(false))
1065 );
1066 assert_eq!(
1067 try_integer_comparison(&exact, &rounded, IntegerComparisonOp::Gt).expect("comparison"),
1068 Some(Value::Bool(true))
1069 );
1070
1071 let tensor = array(IntegerStorage::U64(vec![0, (1_u64 << 53) + 1]), vec![1, 2]);
1072 assert_eq!(
1073 try_integer_comparison(&tensor, &rounded, IntegerComparisonOp::Ne).expect("comparison"),
1074 Some(Value::LogicalArray(
1075 LogicalArray::new(vec![1, 1], vec![1, 2]).expect("logical result")
1076 ))
1077 );
1078 }
1079
1080 #[test]
1081 fn compares_integer_storage_to_broadcast_float_arrays_without_64_bit_loss() {
1082 let integer = array(
1083 IntegerStorage::U64(vec![1_u64 << 53, (1_u64 << 53) + 1]),
1084 vec![2, 1],
1085 );
1086 let float = Value::Tensor(
1087 runmat_value::Tensor::new(
1088 vec![(1_u64 << 53) as f64, 0.0, (1_u64 << 53) as f64],
1089 vec![1, 3],
1090 )
1091 .expect("float tensor"),
1092 );
1093 let result = try_integer_comparison(&integer, &float, IntegerComparisonOp::Eq)
1094 .expect("comparison")
1095 .expect("integer path");
1096 assert_eq!(
1097 result,
1098 Value::LogicalArray(
1099 LogicalArray::new(vec![1, 0, 0, 0, 1, 0], vec![2, 3]).expect("logical result")
1100 )
1101 );
1102
1103 let result = try_integer_comparison(&float, &integer, IntegerComparisonOp::Lt)
1104 .expect("comparison")
1105 .expect("integer path");
1106 assert_eq!(
1107 result,
1108 Value::LogicalArray(
1109 LogicalArray::new(vec![0, 1, 1, 1, 0, 1], vec![2, 3]).expect("logical result")
1110 )
1111 );
1112 }
1113
1114 #[test]
1115 fn compares_integer_storage_to_logical_arrays() {
1116 let integer = array(IntegerStorage::I8(vec![0, 1]), vec![1, 2]);
1117 let logical =
1118 Value::LogicalArray(LogicalArray::new(vec![0, 1], vec![1, 2]).expect("logical array"));
1119 assert_eq!(
1120 try_integer_comparison(&integer, &logical, IntegerComparisonOp::Eq)
1121 .expect("comparison"),
1122 Some(Value::LogicalArray(
1123 LogicalArray::new(vec![1, 1], vec![1, 2]).expect("logical result")
1124 ))
1125 );
1126 }
1127
1128 #[test]
1129 fn compares_all_integer_storage_classes_to_complex_tensors_exactly() {
1130 let cases = [
1131 (
1132 IntegerStorage::I8(vec![-7, 5]),
1133 vec![(-7.0, 0.0), (0.0, 0.0)],
1134 vec![1, 0],
1135 ),
1136 (
1137 IntegerStorage::I16(vec![-300, 5]),
1138 vec![(-300.0, 0.0), (0.0, 0.0)],
1139 vec![1, 0],
1140 ),
1141 (
1142 IntegerStorage::I32(vec![-70_000, 5]),
1143 vec![(-70_000.0, 0.0), (0.0, 0.0)],
1144 vec![1, 0],
1145 ),
1146 (
1147 IntegerStorage::I64(vec![i64::MAX, -9_007_199_254_740_991]),
1148 vec![(i64::MAX as f64, 0.0), (-9_007_199_254_740_991.0, 0.0)],
1149 vec![0, 1],
1150 ),
1151 (
1152 IntegerStorage::U8(vec![7, 5]),
1153 vec![(7.0, 0.0), (0.0, 0.0)],
1154 vec![1, 0],
1155 ),
1156 (
1157 IntegerStorage::U16(vec![300, 5]),
1158 vec![(300.0, 0.0), (0.0, 0.0)],
1159 vec![1, 0],
1160 ),
1161 (
1162 IntegerStorage::U32(vec![70_000, 5]),
1163 vec![(70_000.0, 0.0), (0.0, 0.0)],
1164 vec![1, 0],
1165 ),
1166 (
1167 IntegerStorage::U64(vec![(1_u64 << 53) + 1, u64::MAX]),
1168 vec![((1_u64 << 53) as f64, 0.0), (u64::MAX as f64, 0.0)],
1169 vec![0, 0],
1170 ),
1171 ];
1172
1173 for (storage, complex_data, expected_eq) in cases {
1174 let integer =
1175 runmat_value::Tensor::new_integer(storage, vec![1, 2]).expect("integer tensor");
1176 let complex = Value::ComplexTensor(
1177 ComplexTensor::new(complex_data, vec![1, 2]).expect("complex tensor"),
1178 );
1179 let integer = Value::Tensor(integer);
1180
1181 assert_eq!(
1182 try_complex_integer_equality_comparison(
1183 &integer,
1184 &complex,
1185 IntegerComparisonOp::Eq,
1186 )
1187 .expect("comparison"),
1188 Some(Value::LogicalArray(
1189 LogicalArray::new(expected_eq.clone(), vec![1, 2]).expect("logical result")
1190 ))
1191 );
1192 assert_eq!(
1193 try_complex_integer_equality_comparison(
1194 &complex,
1195 &integer,
1196 IntegerComparisonOp::Ne,
1197 )
1198 .expect("comparison"),
1199 Some(Value::LogicalArray(
1200 LogicalArray::new(
1201 expected_eq.iter().map(|value| *value ^ 1).collect(),
1202 vec![1, 2],
1203 )
1204 .expect("logical result")
1205 ))
1206 );
1207 }
1208 }
1209
1210 #[test]
1211 fn complex_ordering_uses_only_real_components_for_all_relations() {
1212 let lhs = Value::Complex(2.0, f64::NAN);
1213 let rhs = Value::Complex(2.0, f64::INFINITY);
1214 for (operation, expected) in [
1215 (IntegerComparisonOp::Lt, false),
1216 (IntegerComparisonOp::Le, true),
1217 (IntegerComparisonOp::Gt, false),
1218 (IntegerComparisonOp::Ge, true),
1219 ] {
1220 assert_eq!(
1221 try_complex_ordering_comparison(&lhs, &rhs, operation).expect("comparison"),
1222 Some(Value::Bool(expected))
1223 );
1224 }
1225
1226 let nan_real = Value::Complex(f64::NAN, 0.0);
1227 for operation in [
1228 IntegerComparisonOp::Lt,
1229 IntegerComparisonOp::Le,
1230 IntegerComparisonOp::Gt,
1231 IntegerComparisonOp::Ge,
1232 ] {
1233 assert_eq!(
1234 try_complex_ordering_comparison(&nan_real, &Value::Num(0.0), operation)
1235 .expect("NaN comparison"),
1236 Some(Value::Bool(false))
1237 );
1238 }
1239 assert_eq!(
1240 try_complex_ordering_comparison(
1241 &Value::Complex(1.0, 0.0),
1242 &Value::Num(2.0),
1243 IntegerComparisonOp::Lt,
1244 )
1245 .expect("structurally complex zero-imaginary comparison"),
1246 Some(Value::Bool(true))
1247 );
1248 }
1249
1250 #[test]
1251 fn complex_integer_ordering_preserves_wide_real_components_and_broadcasts() {
1252 let storage = runmat_value::IntegerComplexStorage::new(
1253 IntegerStorage::U64(vec![(1_u64 << 53) + 1, u64::MAX]),
1254 IntegerStorage::U64(vec![u64::MAX, 0]),
1255 )
1256 .expect("complex integer storage");
1257 let complex = Value::ComplexTensor(
1258 ComplexTensor::new_integer(storage, vec![2, 1]).expect("complex integer tensor"),
1259 );
1260 let rounded = Value::Tensor(
1261 Tensor::new(vec![(1_u64 << 53) as f64, u64::MAX as f64], vec![1, 2])
1262 .expect("double tensor"),
1263 );
1264
1265 assert_eq!(
1266 try_complex_ordering_comparison(&complex, &rounded, IntegerComparisonOp::Gt)
1267 .expect("complex/double comparison"),
1268 Some(Value::LogicalArray(
1269 LogicalArray::new(vec![1, 1, 0, 0], vec![2, 2]).expect("logical result")
1270 ))
1271 );
1272 assert_eq!(
1273 try_complex_ordering_comparison(&rounded, &complex, IntegerComparisonOp::Lt)
1274 .expect("double/complex comparison"),
1275 Some(Value::LogicalArray(
1276 LogicalArray::new(vec![1, 1, 0, 0], vec![2, 2]).expect("logical result")
1277 ))
1278 );
1279
1280 let rounded_complex = Value::ComplexTensor(
1281 ComplexTensor::new(
1282 vec![((1_u64 << 53) as f64, -1.0), (u64::MAX as f64, 1.0)],
1283 vec![1, 2],
1284 )
1285 .expect("floating complex tensor"),
1286 );
1287 assert_eq!(
1288 try_complex_ordering_comparison(&complex, &rounded_complex, IntegerComparisonOp::Gt,)
1289 .expect("integer-complex/floating-complex comparison"),
1290 Some(Value::LogicalArray(
1291 LogicalArray::new(vec![1, 1, 0, 0], vec![2, 2]).expect("logical result")
1292 ))
1293 );
1294 }
1295
1296 #[test]
1297 fn complex_single_and_double_ordering_broadcasts_without_using_imaginary_components() {
1298 let lhs = Value::ComplexTensor(
1299 ComplexTensor::from_complex_storage(
1300 ComplexStorage::F32(vec![(1.0, f32::NAN), (3.0, f32::INFINITY)]),
1301 vec![2, 1],
1302 )
1303 .expect("complex single"),
1304 );
1305 let rhs = Value::Tensor(Tensor::new(vec![2.0, 3.0], vec![1, 2]).expect("double tensor"));
1306 let expected = LogicalArray::new(vec![1, 0, 1, 1], vec![2, 2]).expect("logical result");
1307 assert_eq!(
1308 try_complex_ordering_comparison(&lhs, &rhs, IntegerComparisonOp::Le)
1309 .expect("complex-single/double comparison"),
1310 Some(Value::LogicalArray(expected.clone()))
1311 );
1312 assert_eq!(
1313 try_complex_ordering_comparison(&rhs, &lhs, IntegerComparisonOp::Ge)
1314 .expect("double/complex-single comparison"),
1315 Some(Value::LogicalArray(expected))
1316 );
1317 }
1318
1319 #[test]
1320 fn complex_ordering_accepts_logical_and_character_numeric_operands() {
1321 let logical =
1322 Value::LogicalArray(LogicalArray::new(vec![0, 1], vec![1, 2]).expect("logical"));
1323 assert_eq!(
1324 try_complex_ordering_comparison(
1325 &Value::Complex(0.5, 99.0),
1326 &logical,
1327 IntegerComparisonOp::Gt,
1328 )
1329 .expect("complex/logical comparison"),
1330 Some(Value::LogicalArray(
1331 LogicalArray::new(vec![1, 0], vec![1, 2]).expect("logical result")
1332 ))
1333 );
1334
1335 let chars = Value::CharArray(
1336 runmat_value::CharArray::new(vec!['A', 'C', 'B', 'D'], 2, 2).expect("character array"),
1337 );
1338 assert_eq!(
1339 try_complex_ordering_comparison(
1340 &Value::Complex(66.0, -99.0),
1341 &chars,
1342 IntegerComparisonOp::Lt,
1343 )
1344 .expect("complex/character comparison"),
1345 Some(Value::LogicalArray(
1346 LogicalArray::new(vec![0, 0, 1, 1], vec![2, 2]).expect("logical result")
1347 ))
1348 );
1349 }
1350
1351 fn assert_resident_complex_ordering(provider: &dyn runmat_accelerate_api::AccelProvider) {
1352 let complex =
1353 ComplexTensor::new(vec![(1.0, 99.0), (3.0, -99.0)], vec![1, 2]).expect("complex");
1354 let complex_rhs =
1355 ComplexTensor::new(vec![(2.0, -7.0), (2.0, 7.0)], vec![1, 2]).expect("complex rhs");
1356 let real = Tensor::new(vec![2.0, 2.0], vec![1, 2]).expect("real");
1357 let complex_handle = gpu_helpers::upload_complex_tensor(provider, &complex).unwrap();
1358 let complex_rhs_handle =
1359 gpu_helpers::upload_complex_tensor(provider, &complex_rhs).unwrap();
1360 let real_handle = gpu_helpers::upload_tensor(provider, &real).unwrap();
1361 for (builtin, expected, reverse_expected) in [
1362 ("lt", vec![1.0, 0.0], vec![0.0, 1.0]),
1363 ("le", vec![1.0, 0.0], vec![0.0, 1.0]),
1364 ("gt", vec![0.0, 1.0], vec![1.0, 0.0]),
1365 ("ge", vec![0.0, 1.0], vec![1.0, 0.0]),
1366 ] {
1367 for rhs in [&real_handle, &complex_rhs_handle] {
1368 let result = crate::call_builtin(
1369 builtin,
1370 &[
1371 Value::GpuTensor(complex_handle.clone()),
1372 Value::GpuTensor(rhs.clone()),
1373 ],
1374 )
1375 .expect("resident complex ordering");
1376 assert!(
1377 matches!(&result, Value::GpuTensor(handle) if runmat_accelerate_api::handle_is_logical(handle)),
1378 "{builtin} did not preserve logical residency for rhs storage {:?}: {result:?}",
1379 runmat_accelerate_api::handle_storage(rhs)
1380 );
1381 let gathered = test_support::gather(result).expect("gather logical result");
1382 assert_eq!(gathered.shape, vec![1, 2]);
1383 assert_eq!(gathered.materialize_f64(), expected);
1384 }
1385 let result = crate::call_builtin(
1386 builtin,
1387 &[
1388 Value::GpuTensor(real_handle.clone()),
1389 Value::GpuTensor(complex_handle.clone()),
1390 ],
1391 )
1392 .expect("reverse resident complex ordering");
1393 let gathered = test_support::gather(result).expect("gather reverse logical result");
1394 assert_eq!(gathered.materialize_f64(), reverse_expected);
1395 }
1396 let _ = provider.free(&complex_handle);
1397 let _ = provider.free(&complex_rhs_handle);
1398 let _ = provider.free(&real_handle);
1399 }
1400
1401 #[test]
1402 fn resident_complex_ordering_uses_provider_real_component_paths() {
1403 test_support::with_test_provider(assert_resident_complex_ordering);
1404 }
1405
1406 #[cfg(feature = "wgpu")]
1407 #[test]
1408 fn wgpu_complex_ordering_uses_provider_real_component_paths() {
1409 if runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
1410 runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
1411 )
1412 .is_err()
1413 {
1414 return;
1415 }
1416 let provider = runmat_accelerate_api::provider().expect("WGPU provider");
1417 assert_resident_complex_ordering(provider);
1418 }
1419}