1use log::debug;
4use runmat_accelerate_api::{GpuTensorHandle, GpuTensorStorage, ProviderBandwidth};
5use runmat_builtins::{
6 BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
7 BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
8 ComplexTensor, LogicalArray, SparseTensor, Tensor, Value,
9};
10use runmat_macros::runtime_builtin;
11
12use crate::builtins::common::gpu_helpers;
13use crate::builtins::common::spec::{
14 BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
15 ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
16};
17use crate::builtins::math::linalg::type_resolvers::logical_scalar_type;
18use crate::{build_runtime_error, BuiltinResult, RuntimeError};
19
20const OUTPUT: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
21 name: "tf",
22 ty: BuiltinParamType::LogicalArray,
23 arity: BuiltinParamArity::Required,
24 default: None,
25 description: "True when the input matrix satisfies the requested structure predicate.",
26}];
27
28const INPUTS: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
29 name: "A",
30 ty: BuiltinParamType::Any,
31 arity: BuiltinParamArity::Required,
32 default: None,
33 description: "Input numeric, logical, complex, sparse, or gpuArray matrix.",
34}];
35
36const ISDIAG_SIGNATURES: [BuiltinSignatureDescriptor; 1] = [BuiltinSignatureDescriptor {
37 label: "tf = isdiag(A)",
38 inputs: &INPUTS,
39 outputs: &OUTPUT,
40}];
41
42const ISTRIL_SIGNATURES: [BuiltinSignatureDescriptor; 1] = [BuiltinSignatureDescriptor {
43 label: "tf = istril(A)",
44 inputs: &INPUTS,
45 outputs: &OUTPUT,
46}];
47
48const ISTRIU_SIGNATURES: [BuiltinSignatureDescriptor; 1] = [BuiltinSignatureDescriptor {
49 label: "tf = istriu(A)",
50 inputs: &INPUTS,
51 outputs: &OUTPUT,
52}];
53
54const ISDIAG_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
55 code: "RM.ISDIAG.INVALID_INPUT",
56 identifier: Some("RunMat:isdiag:InvalidInput"),
57 when: "Input cannot be interpreted as a supported numeric, logical, sparse, complex, or gpuArray matrix.",
58 message: "isdiag: invalid input",
59};
60
61const ISDIAG_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
62 code: "RM.ISDIAG.INTERNAL",
63 identifier: Some("RunMat:isdiag:Internal"),
64 when: "Runtime fails while gathering GPU input or constructing an internal matrix view.",
65 message: "isdiag: internal runtime failure",
66};
67
68const ISTRIL_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
69 code: "RM.ISTRIL.INVALID_INPUT",
70 identifier: Some("RunMat:istril:InvalidInput"),
71 when: "Input cannot be interpreted as a supported numeric, logical, sparse, complex, or gpuArray matrix.",
72 message: "istril: invalid input",
73};
74
75const ISTRIL_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
76 code: "RM.ISTRIL.INTERNAL",
77 identifier: Some("RunMat:istril:Internal"),
78 when: "Runtime fails while gathering GPU input or constructing an internal matrix view.",
79 message: "istril: internal runtime failure",
80};
81
82const ISTRIU_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
83 code: "RM.ISTRIU.INVALID_INPUT",
84 identifier: Some("RunMat:istriu:InvalidInput"),
85 when: "Input cannot be interpreted as a supported numeric, logical, sparse, complex, or gpuArray matrix.",
86 message: "istriu: invalid input",
87};
88
89const ISTRIU_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
90 code: "RM.ISTRIU.INTERNAL",
91 identifier: Some("RunMat:istriu:Internal"),
92 when: "Runtime fails while gathering GPU input or constructing an internal matrix view.",
93 message: "istriu: internal runtime failure",
94};
95
96const ISDIAG_ERRORS: [BuiltinErrorDescriptor; 2] =
97 [ISDIAG_ERROR_INVALID_INPUT, ISDIAG_ERROR_INTERNAL];
98const ISTRIL_ERRORS: [BuiltinErrorDescriptor; 2] =
99 [ISTRIL_ERROR_INVALID_INPUT, ISTRIL_ERROR_INTERNAL];
100const ISTRIU_ERRORS: [BuiltinErrorDescriptor; 2] =
101 [ISTRIU_ERROR_INVALID_INPUT, ISTRIU_ERROR_INTERNAL];
102
103pub const ISDIAG_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
104 signatures: &ISDIAG_SIGNATURES,
105 output_mode: BuiltinOutputMode::Fixed,
106 completion_policy: BuiltinCompletionPolicy::Public,
107 errors: &ISDIAG_ERRORS,
108};
109
110pub const ISTRIL_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
111 signatures: &ISTRIL_SIGNATURES,
112 output_mode: BuiltinOutputMode::Fixed,
113 completion_policy: BuiltinCompletionPolicy::Public,
114 errors: &ISTRIL_ERRORS,
115};
116
117pub const ISTRIU_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
118 signatures: &ISTRIU_SIGNATURES,
119 output_mode: BuiltinOutputMode::Fixed,
120 completion_policy: BuiltinCompletionPolicy::Public,
121 errors: &ISTRIU_ERRORS,
122};
123
124#[runmat_macros::register_gpu_spec(
125 builtin_path = "crate::builtins::math::linalg::structure::diagonal_triangular"
126)]
127pub const ISDIAG_GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
128 name: "isdiag",
129 op_kind: GpuOpKind::Custom("structure_analysis"),
130 supported_precisions: &[ScalarType::F32, ScalarType::F64],
131 broadcast: BroadcastSemantics::None,
132 provider_hooks: &[ProviderHook::Custom("bandwidth")],
133 constant_strategy: ConstantStrategy::InlineLiteral,
134 residency: ResidencyPolicy::GatherImmediately,
135 nan_mode: ReductionNaN::Include,
136 two_pass_threshold: None,
137 workgroup_size: None,
138 accepts_nan_mode: false,
139 notes: "Real and logical gpuArray inputs use the provider bandwidth hook and read back only the small [lower, upper] result. Complex-interleaved inputs and providers without bandwidth support fall back to the host path.",
140};
141
142#[runmat_macros::register_gpu_spec(
143 builtin_path = "crate::builtins::math::linalg::structure::diagonal_triangular"
144)]
145pub const ISTRIL_GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
146 name: "istril",
147 op_kind: GpuOpKind::Custom("structure_analysis"),
148 supported_precisions: &[ScalarType::F32, ScalarType::F64],
149 broadcast: BroadcastSemantics::None,
150 provider_hooks: &[ProviderHook::Custom("bandwidth")],
151 constant_strategy: ConstantStrategy::InlineLiteral,
152 residency: ResidencyPolicy::GatherImmediately,
153 nan_mode: ReductionNaN::Include,
154 two_pass_threshold: None,
155 workgroup_size: None,
156 accepts_nan_mode: false,
157 notes:
158 "Real and logical gpuArray inputs use the provider bandwidth hook and read back only the small [lower, upper] result. Complex-interleaved inputs and providers without bandwidth support fall back to the host path.",
159};
160
161#[runmat_macros::register_gpu_spec(
162 builtin_path = "crate::builtins::math::linalg::structure::diagonal_triangular"
163)]
164pub const ISTRIU_GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
165 name: "istriu",
166 op_kind: GpuOpKind::Custom("structure_analysis"),
167 supported_precisions: &[ScalarType::F32, ScalarType::F64],
168 broadcast: BroadcastSemantics::None,
169 provider_hooks: &[ProviderHook::Custom("bandwidth")],
170 constant_strategy: ConstantStrategy::InlineLiteral,
171 residency: ResidencyPolicy::GatherImmediately,
172 nan_mode: ReductionNaN::Include,
173 two_pass_threshold: None,
174 workgroup_size: None,
175 accepts_nan_mode: false,
176 notes:
177 "Real and logical gpuArray inputs use the provider bandwidth hook and read back only the small [lower, upper] result. Complex-interleaved inputs and providers without bandwidth support fall back to the host path.",
178};
179
180#[runmat_macros::register_fusion_spec(
181 builtin_path = "crate::builtins::math::linalg::structure::diagonal_triangular"
182)]
183pub const ISDIAG_FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
184 name: "isdiag",
185 shape: ShapeRequirements::Any,
186 constant_strategy: ConstantStrategy::InlineLiteral,
187 elementwise: None,
188 reduction: None,
189 emits_nan: false,
190 notes: "Returns a host logical scalar and acts as a fusion sink.",
191};
192
193#[runmat_macros::register_fusion_spec(
194 builtin_path = "crate::builtins::math::linalg::structure::diagonal_triangular"
195)]
196pub const ISTRIL_FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
197 name: "istril",
198 shape: ShapeRequirements::Any,
199 constant_strategy: ConstantStrategy::InlineLiteral,
200 elementwise: None,
201 reduction: None,
202 emits_nan: false,
203 notes: "Returns a host logical scalar and acts as a fusion sink.",
204};
205
206#[runmat_macros::register_fusion_spec(
207 builtin_path = "crate::builtins::math::linalg::structure::diagonal_triangular"
208)]
209pub const ISTRIU_FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
210 name: "istriu",
211 shape: ShapeRequirements::Any,
212 constant_strategy: ConstantStrategy::InlineLiteral,
213 elementwise: None,
214 reduction: None,
215 emits_nan: false,
216 notes: "Returns a host logical scalar and acts as a fusion sink.",
217};
218
219#[derive(Clone, Copy)]
220enum StructurePredicate {
221 Diagonal,
222 LowerTriangular,
223 UpperTriangular,
224}
225
226#[derive(Clone, Copy)]
227struct BuiltinContext {
228 name: &'static str,
229 invalid_input: &'static BuiltinErrorDescriptor,
230 internal: &'static BuiltinErrorDescriptor,
231}
232
233#[runtime_builtin(
234 name = "isdiag",
235 category = "math/linalg/structure",
236 summary = "Determine whether a matrix is diagonal.",
237 keywords = "isdiag,diagonal,matrix structure,sparse,gpu",
238 accel = "structure",
239 type_resolver(logical_scalar_type),
240 descriptor(crate::builtins::math::linalg::structure::diagonal_triangular::ISDIAG_DESCRIPTOR),
241 builtin_path = "crate::builtins::math::linalg::structure::diagonal_triangular"
242)]
243async fn isdiag_builtin(value: Value) -> BuiltinResult<Value> {
244 structure_predicate_builtin(
245 value,
246 StructurePredicate::Diagonal,
247 BuiltinContext {
248 name: "isdiag",
249 invalid_input: &ISDIAG_ERROR_INVALID_INPUT,
250 internal: &ISDIAG_ERROR_INTERNAL,
251 },
252 )
253 .await
254}
255
256#[runtime_builtin(
257 name = "istril",
258 category = "math/linalg/structure",
259 summary = "Determine whether a matrix is lower triangular.",
260 keywords = "istril,lower triangular,matrix structure,sparse,gpu",
261 accel = "structure",
262 type_resolver(logical_scalar_type),
263 descriptor(crate::builtins::math::linalg::structure::diagonal_triangular::ISTRIL_DESCRIPTOR),
264 builtin_path = "crate::builtins::math::linalg::structure::diagonal_triangular"
265)]
266async fn istril_builtin(value: Value) -> BuiltinResult<Value> {
267 structure_predicate_builtin(
268 value,
269 StructurePredicate::LowerTriangular,
270 BuiltinContext {
271 name: "istril",
272 invalid_input: &ISTRIL_ERROR_INVALID_INPUT,
273 internal: &ISTRIL_ERROR_INTERNAL,
274 },
275 )
276 .await
277}
278
279#[runtime_builtin(
280 name = "istriu",
281 category = "math/linalg/structure",
282 summary = "Determine whether a matrix is upper triangular.",
283 keywords = "istriu,upper triangular,matrix structure,sparse,gpu",
284 accel = "structure",
285 type_resolver(logical_scalar_type),
286 descriptor(crate::builtins::math::linalg::structure::diagonal_triangular::ISTRIU_DESCRIPTOR),
287 builtin_path = "crate::builtins::math::linalg::structure::diagonal_triangular"
288)]
289async fn istriu_builtin(value: Value) -> BuiltinResult<Value> {
290 structure_predicate_builtin(
291 value,
292 StructurePredicate::UpperTriangular,
293 BuiltinContext {
294 name: "istriu",
295 invalid_input: &ISTRIU_ERROR_INVALID_INPUT,
296 internal: &ISTRIU_ERROR_INTERNAL,
297 },
298 )
299 .await
300}
301
302async fn structure_predicate_builtin(
303 value: Value,
304 predicate: StructurePredicate,
305 ctx: BuiltinContext,
306) -> BuiltinResult<Value> {
307 if let Value::GpuTensor(handle) = value {
308 return Ok(Value::Bool(
309 gpu_structure_predicate(handle, predicate, ctx).await?,
310 ));
311 }
312 let matrix = MatrixInput::from_value(value, ctx).await?;
313 Ok(Value::Bool(matrix.satisfies(predicate)))
314}
315
316async fn gpu_structure_predicate(
317 handle: GpuTensorHandle,
318 predicate: StructurePredicate,
319 ctx: BuiltinContext,
320) -> BuiltinResult<bool> {
321 if runmat_accelerate_api::handle_storage(&handle) != GpuTensorStorage::ComplexInterleaved {
322 let Some((_rows, _cols)) = matrix_dims_or_false(&handle.shape) else {
323 return Ok(false);
324 };
325 if let Some(provider) = runmat_accelerate_api::provider_for_handle(&handle)
326 .or_else(runmat_accelerate_api::provider)
327 {
328 match provider.bandwidth(&handle) {
329 Ok(bandwidth) => return Ok(bandwidth_satisfies(bandwidth, predicate)),
330 Err(err) => debug!("{}: provider bandwidth fallback: {err}", ctx.name),
331 }
332 }
333 }
334
335 let matrix = matrix_from_gpu(handle, ctx).await?;
336 Ok(matrix.satisfies(predicate))
337}
338
339fn bandwidth_satisfies(bandwidth: ProviderBandwidth, predicate: StructurePredicate) -> bool {
340 match predicate {
341 StructurePredicate::Diagonal => bandwidth.lower == 0 && bandwidth.upper == 0,
342 StructurePredicate::LowerTriangular => bandwidth.upper == 0,
343 StructurePredicate::UpperTriangular => bandwidth.lower == 0,
344 }
345}
346
347enum MatrixInput {
348 DenseReal {
349 rows: usize,
350 cols: usize,
351 data: Vec<f64>,
352 },
353 DenseComplex {
354 rows: usize,
355 cols: usize,
356 data: Vec<(f64, f64)>,
357 },
358 Sparse(SparseTensor),
359 TooManyDimensions,
360}
361
362impl MatrixInput {
363 async fn from_value(value: Value, ctx: BuiltinContext) -> BuiltinResult<Self> {
364 if let Value::GpuTensor(handle) = value {
365 return matrix_from_gpu(handle, ctx).await;
366 }
367 matrix_from_host_value(value, ctx)
368 }
369
370 fn satisfies(&self, predicate: StructurePredicate) -> bool {
371 match self {
372 Self::DenseReal { rows, cols, data } => {
373 dense_real_satisfies(*rows, *cols, data, predicate)
374 }
375 Self::DenseComplex { rows, cols, data } => {
376 dense_complex_satisfies(*rows, *cols, data, predicate)
377 }
378 Self::Sparse(sparse) => sparse_satisfies(sparse, predicate),
379 Self::TooManyDimensions => false,
380 }
381 }
382}
383
384fn matrix_from_host_value(value: Value, ctx: BuiltinContext) -> BuiltinResult<MatrixInput> {
385 match value {
386 Value::Tensor(tensor) => matrix_from_real_tensor(tensor, ctx),
387 Value::LogicalArray(logical) => matrix_from_logical(logical, ctx),
388 Value::SparseTensor(sparse) => Ok(MatrixInput::Sparse(sparse)),
389 Value::ComplexTensor(tensor) => matrix_from_complex_tensor(tensor, ctx),
390 Value::Num(n) => Ok(MatrixInput::DenseReal {
391 rows: 1,
392 cols: 1,
393 data: vec![n],
394 }),
395 Value::Int(i) => Ok(MatrixInput::DenseReal {
396 rows: 1,
397 cols: 1,
398 data: vec![i.to_f64()],
399 }),
400 Value::Bool(b) => Ok(MatrixInput::DenseReal {
401 rows: 1,
402 cols: 1,
403 data: vec![if b { 1.0 } else { 0.0 }],
404 }),
405 Value::Complex(re, im) => Ok(MatrixInput::DenseComplex {
406 rows: 1,
407 cols: 1,
408 data: vec![(re, im)],
409 }),
410 other => Err(runtime_error_with_detail(
411 ctx,
412 ctx.invalid_input,
413 format!(
414 "unsupported input type {:?}; expected numeric, logical, sparse, complex, or gpuArray matrix",
415 other
416 ),
417 )),
418 }
419}
420
421async fn matrix_from_gpu(
422 handle: GpuTensorHandle,
423 ctx: BuiltinContext,
424) -> BuiltinResult<MatrixInput> {
425 let gathered = gpu_helpers::gather_value_async(&Value::GpuTensor(handle))
426 .await
427 .map_err(|err| runtime_error_with_detail(ctx, ctx.internal, err))?;
428 matrix_from_host_value(gathered, ctx)
429}
430
431fn matrix_from_real_tensor(tensor: Tensor, ctx: BuiltinContext) -> BuiltinResult<MatrixInput> {
432 let Some((rows, cols)) = matrix_dims_or_false(&tensor.shape) else {
433 return Ok(MatrixInput::TooManyDimensions);
434 };
435 if rows
436 .checked_mul(cols)
437 .is_none_or(|len| len > tensor.data.len())
438 {
439 return Err(runtime_error_with_detail(
440 ctx,
441 ctx.internal,
442 "tensor shape exceeds backing data length",
443 ));
444 }
445 Ok(MatrixInput::DenseReal {
446 rows,
447 cols,
448 data: tensor.data,
449 })
450}
451
452fn matrix_from_complex_tensor(
453 tensor: ComplexTensor,
454 ctx: BuiltinContext,
455) -> BuiltinResult<MatrixInput> {
456 let Some((rows, cols)) = matrix_dims_or_false(&tensor.shape) else {
457 return Ok(MatrixInput::TooManyDimensions);
458 };
459 if rows
460 .checked_mul(cols)
461 .is_none_or(|len| len > tensor.data.len())
462 {
463 return Err(runtime_error_with_detail(
464 ctx,
465 ctx.internal,
466 "complex tensor shape exceeds backing data length",
467 ));
468 }
469 Ok(MatrixInput::DenseComplex {
470 rows,
471 cols,
472 data: tensor.data,
473 })
474}
475
476fn matrix_from_logical(logical: LogicalArray, ctx: BuiltinContext) -> BuiltinResult<MatrixInput> {
477 let Some((rows, cols)) = matrix_dims_or_false(&logical.shape) else {
478 return Ok(MatrixInput::TooManyDimensions);
479 };
480 if rows
481 .checked_mul(cols)
482 .is_none_or(|len| len > logical.data.len())
483 {
484 return Err(runtime_error_with_detail(
485 ctx,
486 ctx.internal,
487 "logical array shape exceeds backing data length",
488 ));
489 }
490 Ok(MatrixInput::DenseReal {
491 rows,
492 cols,
493 data: logical
494 .data
495 .into_iter()
496 .map(|value| if value != 0 { 1.0 } else { 0.0 })
497 .collect(),
498 })
499}
500
501fn matrix_dims_or_false(shape: &[usize]) -> Option<(usize, usize)> {
502 match shape.len() {
503 0 => Some((1, 1)),
504 1 => Some((1, shape[0])),
505 _ if shape[2..].iter().any(|&dim| dim != 1) => None,
506 _ => Some((shape[0], shape[1])),
507 }
508}
509
510fn dense_real_satisfies(
511 rows: usize,
512 cols: usize,
513 data: &[f64],
514 predicate: StructurePredicate,
515) -> bool {
516 scan_dense(
517 rows,
518 cols,
519 |row, col| {
520 let value = data[row + col * rows];
521 value != 0.0 || value.is_nan()
522 },
523 predicate,
524 )
525}
526
527fn dense_complex_satisfies(
528 rows: usize,
529 cols: usize,
530 data: &[(f64, f64)],
531 predicate: StructurePredicate,
532) -> bool {
533 scan_dense(
534 rows,
535 cols,
536 |row, col| {
537 let (re, im) = data[row + col * rows];
538 !(re == 0.0 && im == 0.0)
539 },
540 predicate,
541 )
542}
543
544fn scan_dense(
545 rows: usize,
546 cols: usize,
547 mut is_nonzero: impl FnMut(usize, usize) -> bool,
548 predicate: StructurePredicate,
549) -> bool {
550 for col in 0..cols {
551 for row in 0..rows {
552 if is_forbidden_nonzero(row, col, predicate) && is_nonzero(row, col) {
553 return false;
554 }
555 }
556 }
557 true
558}
559
560fn sparse_satisfies(sparse: &SparseTensor, predicate: StructurePredicate) -> bool {
561 for col in 0..sparse.cols {
562 for idx in sparse.col_ptrs[col]..sparse.col_ptrs[col + 1] {
563 let row = sparse.row_indices[idx];
564 let value = sparse.values[idx];
565 if is_forbidden_nonzero(row, col, predicate) && (value != 0.0 || value.is_nan()) {
566 return false;
567 }
568 }
569 }
570 true
571}
572
573fn is_forbidden_nonzero(row: usize, col: usize, predicate: StructurePredicate) -> bool {
574 match predicate {
575 StructurePredicate::Diagonal => row != col,
576 StructurePredicate::LowerTriangular => row < col,
577 StructurePredicate::UpperTriangular => row > col,
578 }
579}
580
581fn runtime_error_with_detail(
582 ctx: BuiltinContext,
583 error: &'static BuiltinErrorDescriptor,
584 detail: impl std::fmt::Display,
585) -> RuntimeError {
586 let mut builder =
587 build_runtime_error(format!("{}: {}", error.message, detail)).with_builtin(ctx.name);
588 if let Some(identifier) = error.identifier {
589 builder = builder.with_identifier(identifier);
590 }
591 builder.build()
592}
593
594#[cfg(test)]
595mod tests {
596 use super::*;
597 use futures::executor::block_on;
598 use runmat_accelerate_api::{
599 AccelDownloadFuture, AccelProvider, HostTensorOwned, HostTensorView,
600 };
601 use runmat_builtins::{IntValue, ResolveContext, Type};
602
603 use crate::builtins::common::test_support;
604
605 fn call_isdiag(value: Value) -> BuiltinResult<Value> {
606 block_on(isdiag_builtin(value))
607 }
608
609 fn call_istril(value: Value) -> BuiltinResult<Value> {
610 block_on(istril_builtin(value))
611 }
612
613 fn call_istriu(value: Value) -> BuiltinResult<Value> {
614 block_on(istriu_builtin(value))
615 }
616
617 fn expect_bool(value: Value) -> bool {
618 match value {
619 Value::Bool(value) => value,
620 other => panic!("expected bool, got {other:?}"),
621 }
622 }
623
624 #[test]
625 fn descriptors_cover_core_forms() {
626 assert_eq!(ISDIAG_DESCRIPTOR.signatures[0].label, "tf = isdiag(A)");
627 assert_eq!(ISTRIL_DESCRIPTOR.signatures[0].label, "tf = istril(A)");
628 assert_eq!(ISTRIU_DESCRIPTOR.signatures[0].label, "tf = istriu(A)");
629 }
630
631 #[test]
632 fn type_resolver_returns_logical_scalar() {
633 let out = logical_scalar_type(
634 &[Type::Tensor {
635 shape: Some(vec![Some(3), Some(3)]),
636 }],
637 &ResolveContext::new(Vec::new()),
638 );
639 assert_eq!(out, Type::Bool);
640 }
641
642 #[test]
643 fn scalars_are_diagonal_and_triangular() {
644 assert!(expect_bool(
645 call_isdiag(Value::Int(IntValue::I32(7))).unwrap()
646 ));
647 assert!(expect_bool(call_istril(Value::Num(7.0)).unwrap()));
648 assert!(expect_bool(call_istriu(Value::Bool(true)).unwrap()));
649 assert!(expect_bool(call_isdiag(Value::Complex(0.0, 2.0)).unwrap()));
650 }
651
652 #[test]
653 fn dense_diagonal_matrix_satisfies_all_three_predicates() {
654 let tensor = Tensor::new(
655 vec![1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0],
656 vec![3, 3],
657 )
658 .unwrap();
659 assert!(expect_bool(
660 call_isdiag(Value::Tensor(tensor.clone())).unwrap()
661 ));
662 assert!(expect_bool(
663 call_istril(Value::Tensor(tensor.clone())).unwrap()
664 ));
665 assert!(expect_bool(call_istriu(Value::Tensor(tensor)).unwrap()));
666 }
667
668 #[test]
669 fn lower_and_upper_triangular_are_distinguished() {
670 let lower = Tensor::new(
671 vec![1.0, 2.0, 3.0, 0.0, 4.0, 5.0, 0.0, 0.0, 6.0],
672 vec![3, 3],
673 )
674 .unwrap();
675 assert!(!expect_bool(
676 call_isdiag(Value::Tensor(lower.clone())).unwrap()
677 ));
678 assert!(expect_bool(
679 call_istril(Value::Tensor(lower.clone())).unwrap()
680 ));
681 assert!(!expect_bool(call_istriu(Value::Tensor(lower)).unwrap()));
682
683 let upper = Tensor::new(
684 vec![1.0, 0.0, 0.0, 2.0, 3.0, 0.0, 4.0, 5.0, 6.0],
685 vec![3, 3],
686 )
687 .unwrap();
688 assert!(!expect_bool(
689 call_isdiag(Value::Tensor(upper.clone())).unwrap()
690 ));
691 assert!(!expect_bool(
692 call_istril(Value::Tensor(upper.clone())).unwrap()
693 ));
694 assert!(expect_bool(call_istriu(Value::Tensor(upper)).unwrap()));
695 }
696
697 #[test]
698 fn rectangular_diagonal_matrix_can_be_true() {
699 let tensor = Tensor::new(vec![1.0, 0.0, 0.0, 0.0, 2.0, 0.0], vec![3, 2]).unwrap();
700 assert!(expect_bool(
701 call_isdiag(Value::Tensor(tensor.clone())).unwrap()
702 ));
703 assert!(expect_bool(
704 call_istril(Value::Tensor(tensor.clone())).unwrap()
705 ));
706 assert!(expect_bool(call_istriu(Value::Tensor(tensor)).unwrap()));
707 }
708
709 #[test]
710 fn zero_and_empty_matrices_are_vacuously_true() {
711 let zeros = Tensor::new(vec![0.0; 6], vec![2, 3]).unwrap();
712 assert!(expect_bool(
713 call_isdiag(Value::Tensor(zeros.clone())).unwrap()
714 ));
715 assert!(expect_bool(
716 call_istril(Value::Tensor(zeros.clone())).unwrap()
717 ));
718 assert!(expect_bool(call_istriu(Value::Tensor(zeros)).unwrap()));
719
720 let empty = Tensor::new(Vec::new(), vec![0, 4]).unwrap();
721 assert!(expect_bool(
722 call_isdiag(Value::Tensor(empty.clone())).unwrap()
723 ));
724 assert!(expect_bool(
725 call_istril(Value::Tensor(empty.clone())).unwrap()
726 ));
727 assert!(expect_bool(call_istriu(Value::Tensor(empty)).unwrap()));
728 }
729
730 #[test]
731 fn one_dimensional_tensor_shape_uses_row_vector_semantics() {
732 let row = Tensor::new(vec![1.0, 2.0, 3.0], vec![3]).unwrap();
733 assert!(!expect_bool(
734 call_isdiag(Value::Tensor(row.clone())).unwrap()
735 ));
736 assert!(!expect_bool(
737 call_istril(Value::Tensor(row.clone())).unwrap()
738 ));
739 assert!(expect_bool(call_istriu(Value::Tensor(row)).unwrap()));
740 }
741
742 #[test]
743 fn arrays_with_nonsingleton_trailing_dimensions_return_false() {
744 let tensor = Tensor::new(vec![0.0; 8], vec![2, 2, 2]).unwrap();
745 assert!(!expect_bool(
746 call_isdiag(Value::Tensor(tensor.clone())).unwrap()
747 ));
748 assert!(!expect_bool(
749 call_istril(Value::Tensor(tensor.clone())).unwrap()
750 ));
751 assert!(!expect_bool(call_istriu(Value::Tensor(tensor)).unwrap()));
752 }
753
754 #[test]
755 fn arrays_with_trailing_zero_dimensions_return_false() {
756 let tensor = Tensor::new(Vec::new(), vec![2, 2, 0]).unwrap();
757 assert!(!expect_bool(
758 call_isdiag(Value::Tensor(tensor.clone())).unwrap()
759 ));
760 assert!(!expect_bool(
761 call_istril(Value::Tensor(tensor.clone())).unwrap()
762 ));
763 assert!(!expect_bool(call_istriu(Value::Tensor(tensor)).unwrap()));
764 }
765
766 #[test]
767 fn logical_arrays_are_supported() {
768 let logical = LogicalArray::new(vec![1, 0, 0, 1], vec![2, 2]).unwrap();
769 assert!(expect_bool(
770 call_isdiag(Value::LogicalArray(logical)).unwrap()
771 ));
772 }
773
774 #[test]
775 fn complex_off_structure_value_is_nonzero() {
776 let tensor = ComplexTensor::new(
777 vec![(1.0, 0.0), (0.0, 2.0), (0.0, 0.0), (1.0, 0.0)],
778 vec![2, 2],
779 )
780 .unwrap();
781 assert!(!expect_bool(
782 call_isdiag(Value::ComplexTensor(tensor.clone())).unwrap()
783 ));
784 assert!(expect_bool(
785 call_istril(Value::ComplexTensor(tensor.clone())).unwrap()
786 ));
787 assert!(!expect_bool(
788 call_istriu(Value::ComplexTensor(tensor)).unwrap()
789 ));
790 }
791
792 #[test]
793 fn sparse_inputs_are_scanned_without_densifying() {
794 let sparse =
795 SparseTensor::new(3, 3, vec![0, 2, 3, 3], vec![0, 2, 1], vec![1.0, 4.0, 2.0]).unwrap();
796 assert!(!expect_bool(
797 call_isdiag(Value::SparseTensor(sparse.clone())).unwrap()
798 ));
799 assert!(expect_bool(
800 call_istril(Value::SparseTensor(sparse.clone())).unwrap()
801 ));
802 assert!(!expect_bool(
803 call_istriu(Value::SparseTensor(sparse)).unwrap()
804 ));
805 }
806
807 #[test]
808 fn explicit_nan_off_structure_counts_as_nonzero() {
809 let tensor = Tensor::new(vec![1.0, f64::NAN, 0.0, 1.0], vec![2, 2]).unwrap();
810 assert!(!expect_bool(
811 call_isdiag(Value::Tensor(tensor.clone())).unwrap()
812 ));
813 assert!(expect_bool(
814 call_istril(Value::Tensor(tensor.clone())).unwrap()
815 ));
816 assert!(!expect_bool(call_istriu(Value::Tensor(tensor)).unwrap()));
817 }
818
819 #[derive(Clone, Copy)]
820 struct FixedBandwidthProvider {
821 lower: u32,
822 upper: u32,
823 device_id: u32,
824 }
825
826 impl AccelProvider for FixedBandwidthProvider {
827 fn upload(&self, host: &HostTensorView) -> anyhow::Result<GpuTensorHandle> {
828 Ok(GpuTensorHandle {
829 shape: host.shape.to_vec(),
830 device_id: self.device_id,
831 buffer_id: 1,
832 })
833 }
834
835 fn download<'a>(&'a self, _h: &'a GpuTensorHandle) -> AccelDownloadFuture<'a> {
836 Box::pin(async move { Err(anyhow::anyhow!("download should not be called")) })
837 }
838
839 fn free(&self, _h: &GpuTensorHandle) -> anyhow::Result<()> {
840 Ok(())
841 }
842
843 fn device_info(&self) -> String {
844 "fixed-bandwidth-test-provider".to_string()
845 }
846
847 fn device_id(&self) -> u32 {
848 self.device_id
849 }
850
851 fn bandwidth(&self, _matrix: &GpuTensorHandle) -> anyhow::Result<ProviderBandwidth> {
852 Ok(ProviderBandwidth {
853 lower: self.lower,
854 upper: self.upper,
855 })
856 }
857 }
858
859 static DIAGONAL_PROVIDER: FixedBandwidthProvider = FixedBandwidthProvider {
860 lower: 0,
861 upper: 0,
862 device_id: 92_001,
863 };
864 static LOWER_PROVIDER: FixedBandwidthProvider = FixedBandwidthProvider {
865 lower: 2,
866 upper: 0,
867 device_id: 92_002,
868 };
869 static UPPER_PROVIDER: FixedBandwidthProvider = FixedBandwidthProvider {
870 lower: 0,
871 upper: 2,
872 device_id: 92_003,
873 };
874
875 fn fixed_gpu_handle(provider: &FixedBandwidthProvider) -> Value {
876 Value::GpuTensor(GpuTensorHandle {
877 shape: vec![3, 3],
878 device_id: provider.device_id(),
879 buffer_id: 1,
880 })
881 }
882
883 #[derive(Clone, Copy)]
884 struct FallbackDownloadProvider;
885
886 impl AccelProvider for FallbackDownloadProvider {
887 fn upload(&self, host: &HostTensorView) -> anyhow::Result<GpuTensorHandle> {
888 Ok(GpuTensorHandle {
889 shape: host.shape.to_vec(),
890 device_id: self.device_id(),
891 buffer_id: 1,
892 })
893 }
894
895 fn download<'a>(&'a self, h: &'a GpuTensorHandle) -> AccelDownloadFuture<'a> {
896 Box::pin(async move {
897 Ok(HostTensorOwned {
898 data: vec![1.0, 3.0, 0.0, 2.0],
899 shape: h.shape.clone(),
900 storage: GpuTensorStorage::Real,
901 })
902 })
903 }
904
905 fn free(&self, _h: &GpuTensorHandle) -> anyhow::Result<()> {
906 Ok(())
907 }
908
909 fn device_info(&self) -> String {
910 "fallback-download-test-provider".to_string()
911 }
912
913 fn device_id(&self) -> u32 {
914 92_004
915 }
916
917 fn bandwidth(&self, _matrix: &GpuTensorHandle) -> anyhow::Result<ProviderBandwidth> {
918 Err(anyhow::anyhow!("bandwidth intentionally unavailable"))
919 }
920 }
921
922 static FALLBACK_PROVIDER: FallbackDownloadProvider = FallbackDownloadProvider;
923
924 #[test]
925 fn gpu_inputs_use_provider_bandwidth_without_gathering() {
926 let _guard = test_support::accel_test_lock();
927 let _thread_provider =
928 runmat_accelerate_api::ThreadProviderGuard::set(Some(&DIAGONAL_PROVIDER));
929 assert!(expect_bool(
930 call_isdiag(fixed_gpu_handle(&DIAGONAL_PROVIDER)).unwrap()
931 ));
932 assert!(expect_bool(
933 call_istril(fixed_gpu_handle(&DIAGONAL_PROVIDER)).unwrap()
934 ));
935 assert!(expect_bool(
936 call_istriu(fixed_gpu_handle(&DIAGONAL_PROVIDER)).unwrap()
937 ));
938
939 let _thread_provider =
940 runmat_accelerate_api::ThreadProviderGuard::set(Some(&LOWER_PROVIDER));
941 assert!(!expect_bool(
942 call_isdiag(fixed_gpu_handle(&LOWER_PROVIDER)).unwrap()
943 ));
944 assert!(expect_bool(
945 call_istril(fixed_gpu_handle(&LOWER_PROVIDER)).unwrap()
946 ));
947 assert!(!expect_bool(
948 call_istriu(fixed_gpu_handle(&LOWER_PROVIDER)).unwrap()
949 ));
950
951 let _thread_provider =
952 runmat_accelerate_api::ThreadProviderGuard::set(Some(&UPPER_PROVIDER));
953 assert!(!expect_bool(
954 call_isdiag(fixed_gpu_handle(&UPPER_PROVIDER)).unwrap()
955 ));
956 assert!(!expect_bool(
957 call_istril(fixed_gpu_handle(&UPPER_PROVIDER)).unwrap()
958 ));
959 assert!(expect_bool(
960 call_istriu(fixed_gpu_handle(&UPPER_PROVIDER)).unwrap()
961 ));
962 }
963
964 #[test]
965 fn gpu_inputs_fall_back_to_download_when_bandwidth_is_unavailable() {
966 let _guard = test_support::accel_test_lock();
967 let _thread_provider =
968 runmat_accelerate_api::ThreadProviderGuard::set(Some(&FALLBACK_PROVIDER));
969 let handle = Value::GpuTensor(GpuTensorHandle {
970 shape: vec![2, 2],
971 device_id: FALLBACK_PROVIDER.device_id(),
972 buffer_id: 1,
973 });
974 assert!(!expect_bool(call_isdiag(handle.clone()).unwrap()));
975 assert!(expect_bool(call_istril(handle.clone()).unwrap()));
976 assert!(!expect_bool(call_istriu(handle).unwrap()));
977 }
978
979 #[test]
980 fn gpu_input_with_inprocess_provider_works() {
981 test_support::with_test_provider(|provider| {
982 let tensor = Tensor::new(vec![1.0, 0.0, 0.0, 1.0], vec![2, 2]).unwrap();
983 let handle = provider
984 .upload(&runmat_accelerate_api::HostTensorView {
985 data: &tensor.data,
986 shape: &tensor.shape,
987 })
988 .unwrap();
989 assert!(expect_bool(call_isdiag(Value::GpuTensor(handle)).unwrap()));
990 });
991 }
992
993 #[test]
994 #[cfg(feature = "wgpu")]
995 fn wgpu_inputs_use_bandwidth_predicates() {
996 if runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
997 runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
998 )
999 .is_err()
1000 {
1001 return;
1002 }
1003 let Some(provider) = runmat_accelerate_api::provider() else {
1004 return;
1005 };
1006 let lower = Tensor::new(
1007 vec![1.0, 4.0, 0.0, 0.0, 2.0, 5.0, 0.0, 0.0, 3.0],
1008 vec![3, 3],
1009 )
1010 .unwrap();
1011 let handle = provider
1012 .upload(&HostTensorView {
1013 data: &lower.data,
1014 shape: &lower.shape,
1015 })
1016 .expect("upload lower");
1017 assert!(!expect_bool(
1018 call_isdiag(Value::GpuTensor(handle.clone())).unwrap()
1019 ));
1020 assert!(expect_bool(
1021 call_istril(Value::GpuTensor(handle.clone())).unwrap()
1022 ));
1023 assert!(!expect_bool(call_istriu(Value::GpuTensor(handle)).unwrap()));
1024 }
1025
1026 #[test]
1027 fn unsupported_inputs_error() {
1028 let err = call_isdiag(Value::String("x".to_string())).unwrap_err();
1029 assert!(err.message.contains("isdiag: invalid input"));
1030 }
1031}