1#[cfg(not(target_arch = "wasm32"))]
4use runmat_accelerate_api::GpuTensorHandle;
5use runmat_accelerate_api::HostTensorView;
6use runmat_builtins::{
7 BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
8 BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
9 ResolveContext, Type,
10};
11use runmat_builtins::{
12 BuiltinIntegerBackendRule, BuiltinIntegerCapabilityDescriptor, BuiltinIntegerComputationDomain,
13 BuiltinIntegerInputAvailability, BuiltinIntegerInputCapability, BuiltinIntegerOutputClassRule,
14 BuiltinIntegerOverflowRule, BuiltinIntegerOverloadKind, BuiltinIntegerScalarDoubleRule,
15};
16use runmat_macros::runtime_builtin;
17use runmat_value::{IntValue, Tensor, Value};
18
19use super::common::{
20 build_strides, dims_from_tokens, fits_positive_platform_index, materialize_value, parse_dims,
21};
22use crate::builtins::array::type_resolvers::is_scalar_type;
23use crate::builtins::common::arg_tokens::tokens_from_context;
24use crate::builtins::common::spec::{
25 BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
26 ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
27};
28use crate::builtins::common::tensor;
29use crate::{build_runtime_error, RuntimeError};
30use runmat_builtins::shape_rules::element_count_if_known;
31
32#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::array::indexing::sub2ind")]
33pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
34 name: "sub2ind",
35 op_kind: GpuOpKind::Custom("indexing"),
36 supported_precisions: &[ScalarType::F32, ScalarType::F64],
37 broadcast: BroadcastSemantics::Matlab,
38 provider_hooks: &[ProviderHook::Custom("sub2ind")],
39 constant_strategy: ConstantStrategy::InlineLiteral,
40 residency: ResidencyPolicy::NewHandle,
41 nan_mode: ReductionNaN::Include,
42 two_pass_threshold: None,
43 workgroup_size: None,
44 accepts_nan_mode: false,
45 notes: "Providers can implement the custom `sub2ind` hook to execute on device; runtimes fall back to host computation otherwise.",
46};
47
48#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::array::indexing::sub2ind")]
49pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
50 name: "sub2ind",
51 shape: ShapeRequirements::Any,
52 constant_strategy: ConstantStrategy::InlineLiteral,
53 elementwise: None,
54 reduction: None,
55 emits_nan: false,
56 notes: "Index conversion executes eagerly on the host; fusion does not apply.",
57};
58
59fn sub2ind_type(args: &[Type], ctx: &ResolveContext) -> Type {
60 if args.len() < 2 {
61 return Type::Unknown;
62 }
63 if let Some(dims) = dims_from_tokens(&tokens_from_context(ctx)) {
64 if args.len() - 1 != dims.len() {
65 return Type::Unknown;
66 }
67 }
68 let subscripts = &args[1..];
69 if subscripts.iter().all(|ty| is_scalar_type(ty)) {
70 return Type::Num;
71 }
72 for ty in subscripts {
73 if let Type::Tensor { shape: Some(shape) } | Type::Logical { shape: Some(shape) } = ty {
74 if element_count_if_known(shape).unwrap_or(0) > 1 {
75 return Type::Tensor {
76 shape: Some(shape.clone()),
77 };
78 }
79 }
80 }
81 Type::tensor()
82}
83
84const BUILTIN_NAME: &str = "sub2ind";
85
86const SUB2IND_INTEGER_INPUTS: [BuiltinIntegerInputCapability; 2] = [
87 BuiltinIntegerInputCapability {
88 name: "sz",
89 classes: &crate::builtins::common::integer_capability::ALL_INTEGER_CLASSES,
90 availability: BuiltinIntegerInputAvailability::Documented,
91 scalar_double: BuiltinIntegerScalarDoubleRule::Allowed,
92 notes: "The size vector accepts every native integer class and is parsed directly as positive platform dimensions.",
93 },
94 BuiltinIntegerInputCapability {
95 name: "I1...In",
96 classes: &crate::builtins::common::integer_capability::ALL_INTEGER_CLASSES,
97 availability: BuiltinIntegerInputAvailability::Documented,
98 scalar_double: BuiltinIntegerScalarDoubleRule::Allowed,
99 notes: "Every subscript role accepts all eight integer classes, including scalar expansion and full-width exact bounds checks.",
100 },
101];
102pub const SUB2IND_INTEGER_CAPABILITIES: [BuiltinIntegerCapabilityDescriptor; 1] =
103 [BuiltinIntegerCapabilityDescriptor {
104 form: "ind = sub2ind(integer_sz, integer_I1, integer_In...)",
105 inputs: &SUB2IND_INTEGER_INPUTS,
106 computation_domain: BuiltinIntegerComputationDomain::Structural,
107 output_class: BuiltinIntegerOutputClassRule::Double,
108 overflow: BuiltinIntegerOverflowRule::Error,
109 backend: BuiltinIntegerBackendRule::GpuRestricted,
110 overload: BuiltinIntegerOverloadKind::Multiple,
111 notes: "Dimensions and subscripts are read from authoritative integer storage, combined with checked column-major index arithmetic, and returned as documented double indices. Resident integer inputs use exact gather fallback and return resident double output when the owner supports upload.",
112 }];
113
114const SUB2IND_OUTPUT: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
115 name: "ind",
116 ty: BuiltinParamType::NumericArray,
117 arity: BuiltinParamArity::Required,
118 default: None,
119 description: "Column-major linear indices corresponding to provided subscripts.",
120}];
121
122const SUB2IND_INPUTS: [BuiltinParamDescriptor; 3] = [
123 BuiltinParamDescriptor {
124 name: "sz",
125 ty: BuiltinParamType::SizeArg,
126 arity: BuiltinParamArity::Required,
127 default: None,
128 description: "Size vector describing source array dimensions.",
129 },
130 BuiltinParamDescriptor {
131 name: "I1",
132 ty: BuiltinParamType::Any,
133 arity: BuiltinParamArity::Required,
134 default: None,
135 description: "First-dimension subscript values.",
136 },
137 BuiltinParamDescriptor {
138 name: "In",
139 ty: BuiltinParamType::Any,
140 arity: BuiltinParamArity::Variadic,
141 default: None,
142 description: "Remaining per-dimension subscript arrays/scalars.",
143 },
144];
145
146const SUB2IND_SIGNATURES: [BuiltinSignatureDescriptor; 1] = [BuiltinSignatureDescriptor {
147 label: "ind = sub2ind(sz, I1, In...)",
148 inputs: &SUB2IND_INPUTS,
149 outputs: &SUB2IND_OUTPUT,
150}];
151
152const SUB2IND_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
153 code: "RM.SUB2IND.INVALID_INPUT",
154 identifier: Some("RunMat:sub2ind:InvalidInput"),
155 when: "Size vector, subscript count, or subscript types are invalid.",
156 message: "sub2ind: invalid input arguments",
157};
158
159const SUB2IND_ERROR_INDEX_BOUNDS: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
160 code: "RM.SUB2IND.INDEX_BOUNDS",
161 identifier: Some("RunMat:sub2ind:IndexBounds"),
162 when: "At least one subscript lies outside bounds for its dimension.",
163 message: "sub2ind: subscript index exceeds dimension bounds",
164};
165
166const SUB2IND_ERROR_PROVIDER: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
167 code: "RM.SUB2IND.PROVIDER",
168 identifier: Some("RunMat:sub2ind:ProviderError"),
169 when: "GPU provider sub2ind hook fails.",
170 message: "sub2ind: provider execution failed",
171};
172
173const SUB2IND_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
174 code: "RM.SUB2IND.INTERNAL",
175 identifier: Some("RunMat:sub2ind:InternalError"),
176 when: "Internal tensor conversion/output construction fails.",
177 message: "sub2ind: internal error",
178};
179
180const SUB2IND_ERRORS: [BuiltinErrorDescriptor; 4] = [
181 SUB2IND_ERROR_INVALID_INPUT,
182 SUB2IND_ERROR_INDEX_BOUNDS,
183 SUB2IND_ERROR_PROVIDER,
184 SUB2IND_ERROR_INTERNAL,
185];
186
187pub const SUB2IND_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
188 signatures: &SUB2IND_SIGNATURES,
189 output_mode: BuiltinOutputMode::Fixed,
190 completion_policy: BuiltinCompletionPolicy::Public,
191 errors: &SUB2IND_ERRORS,
192};
193
194fn sub2ind_error_with_message(
195 message: impl Into<String>,
196 error: &'static BuiltinErrorDescriptor,
197) -> RuntimeError {
198 let mut builder = build_runtime_error(message).with_builtin(BUILTIN_NAME);
199 if let Some(identifier) = error.identifier {
200 builder = builder.with_identifier(identifier);
201 }
202 builder.build()
203}
204
205fn sub2ind_input_error(message: impl Into<String>) -> RuntimeError {
206 sub2ind_error_with_message(message, &SUB2IND_ERROR_INVALID_INPUT)
207}
208
209fn sub2ind_bounds_error(message: impl Into<String>) -> RuntimeError {
210 sub2ind_error_with_message(message, &SUB2IND_ERROR_INDEX_BOUNDS)
211}
212
213fn sub2ind_provider_error(message: impl Into<String>) -> RuntimeError {
214 sub2ind_error_with_message(message, &SUB2IND_ERROR_PROVIDER)
215}
216
217fn sub2ind_internal_error(message: impl Into<String>) -> RuntimeError {
218 sub2ind_error_with_message(message, &SUB2IND_ERROR_INTERNAL)
219}
220
221#[runtime_builtin(
222 name = "sub2ind",
223 category = "array/indexing",
224 summary = "Convert N-D subscripts to MATLAB-style column-major linear indices.",
225 keywords = "sub2ind,linear index,column major,gpu indexing",
226 accel = "custom",
227 type_resolver(sub2ind_type),
228 descriptor(crate::builtins::array::indexing::sub2ind::SUB2IND_DESCRIPTOR),
229 integer_capabilities(crate::builtins::array::indexing::sub2ind::SUB2IND_INTEGER_CAPABILITIES),
230 builtin_path = "crate::builtins::array::indexing::sub2ind"
231)]
232async fn sub2ind_builtin(dims_val: Value, rest: Vec<Value>) -> crate::BuiltinResult<Value> {
233 let (dims_value, dims_was_gpu) = materialize_value(dims_val, "sub2ind").await?;
234 let dims = parse_dims(&dims_value, "sub2ind").await?;
235 if dims.is_empty() {
236 return Err(sub2ind_error("Size vector must have at least one element."));
237 }
238
239 if rest.len() != dims.len() {
240 return Err(sub2ind_error(
241 "The number of subscripts supplied must equal the number of dimensions in the size vector.",
242 ));
243 }
244
245 if let Some(value) = try_gpu_sub2ind(&dims, &rest)? {
246 return Ok(value);
247 }
248
249 let mut saw_gpu = dims_was_gpu;
250 let mut subscripts: Vec<Tensor> = Vec::with_capacity(rest.len());
251 for value in rest {
252 let (materialised, was_gpu) = materialize_value(value, "sub2ind").await?;
253 saw_gpu |= was_gpu;
254 let tensor = tensor::value_into_tensor_for("sub2ind", materialised)
255 .map_err(|message| sub2ind_error(message))?;
256 subscripts.push(tensor);
257 }
258
259 let (result_data, result_shape) = compute_indices(&dims, &subscripts)?;
260 let want_gpu_output = saw_gpu && runmat_accelerate_api::provider().is_some();
261
262 if want_gpu_output {
263 #[cfg(all(test, feature = "wgpu"))]
264 {
265 if runmat_accelerate_api::provider().is_none() {
266 let _ = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
267 runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
268 );
269 }
270 }
271 let shape = result_shape.clone().unwrap_or_else(|| vec![1, 1]);
272 if let Some(provider) = runmat_accelerate_api::provider() {
273 let view = HostTensorView {
274 data: &result_data,
275 shape: &shape,
276 };
277 if let Ok(handle) = provider.upload(&view) {
278 return Ok(Value::GpuTensor(handle));
279 }
280 }
281 }
282
283 build_host_value(result_data, result_shape)
284}
285
286fn try_gpu_sub2ind(dims: &[usize], subs: &[Value]) -> crate::BuiltinResult<Option<Value>> {
287 #[cfg(target_arch = "wasm32")]
288 {
289 let _ = (dims, subs);
290 Ok(None)
291 }
292 #[cfg(not(target_arch = "wasm32"))]
293 {
294 let provider = match runmat_accelerate_api::provider() {
295 Some(p) => p,
296 None => return Ok(None),
297 };
298 if !subs
299 .iter()
300 .all(|value| matches!(value, Value::GpuTensor(_)))
301 {
302 return Ok(None);
303 }
304 if subs.iter().any(
305 |value| matches!(value, Value::GpuTensor(handle) if runmat_accelerate_api::handle_integer_type(handle).is_some()),
306 ) {
307 return Ok(None);
308 }
309 if dims.is_empty() {
310 return Ok(None);
311 }
312
313 let mut handles: Vec<&GpuTensorHandle> = Vec::with_capacity(subs.len());
314 for value in subs {
315 if let Value::GpuTensor(handle) = value {
316 handles.push(handle);
317 }
318 }
319
320 if handles.len() != dims.len() {
321 return Err(sub2ind_error(
322 "The number of subscripts supplied must equal the number of dimensions in the size vector.",
323 ));
324 }
325
326 let mut scalar_mask: Vec<bool> = Vec::with_capacity(handles.len());
327 let mut target_shape: Option<Vec<usize>> = None;
328 let mut result_len: usize = 1;
329 let mut saw_non_scalar = false;
330
331 for handle in &handles {
332 let len = tensor::element_count(&handle.shape);
333 let is_scalar = len == 1;
334 scalar_mask.push(is_scalar);
335 if !is_scalar {
336 saw_non_scalar = true;
337 if let Some(existing) = &target_shape {
338 if existing != &handle.shape {
339 return Err(sub2ind_error("Subscript inputs must have the same size."));
340 }
341 } else {
342 target_shape = Some(handle.shape.clone());
343 result_len = len;
344 }
345 }
346 }
347
348 if !saw_non_scalar {
349 target_shape = Some(vec![1, 1]);
350 result_len = 1;
351 } else if let Some(shape) = &target_shape {
352 result_len = tensor::element_count(shape);
353 }
354
355 let strides = build_strides(dims, "sub2ind")?;
356 if dims.iter().any(|&d| d > u32::MAX as usize)
357 || strides.iter().any(|&s| s > u32::MAX as usize)
358 || result_len > u32::MAX as usize
359 {
360 return Ok(None);
361 }
362
363 let output_shape = target_shape.clone().unwrap_or_else(|| vec![1, 1]);
364 match provider.sub2ind(
365 dims,
366 &strides,
367 &handles,
368 &scalar_mask,
369 result_len,
370 &output_shape,
371 ) {
372 Ok(handle) => Ok(Some(Value::GpuTensor(handle))),
373 Err(err) => Err(sub2ind_provider_error(err.to_string())),
374 }
375 }
376}
377
378fn compute_indices(
379 dims: &[usize],
380 subscripts: &[Tensor],
381) -> crate::BuiltinResult<(Vec<f64>, Option<Vec<usize>>)> {
382 let mut target_shape: Option<Vec<usize>> = None;
383 let mut result_len: usize = 1;
384 let mut has_non_scalar = false;
385
386 for tensor in subscripts {
387 let len = tensor_element_len(tensor);
388 if len != 1 {
389 has_non_scalar = true;
390 if let Some(shape) = &target_shape {
391 if &tensor.shape != shape {
392 return Err(sub2ind_error("Subscript inputs must have the same size."));
393 }
394 } else {
395 target_shape = Some(tensor.shape.clone());
396 result_len = len;
397 }
398 }
399 }
400
401 if !has_non_scalar {
402 target_shape = Some(vec![1, 1]);
404 result_len = 1;
405 }
406
407 if result_len == 0 {
408 return Ok((Vec::new(), target_shape));
409 }
410
411 let strides = build_strides(dims, "sub2ind")?;
412 let mut output = Vec::with_capacity(result_len);
413
414 for idx in 0..result_len {
415 let mut offset: usize = 0;
416 for (dim_index, (&dim, tensor)) in dims.iter().zip(subscripts.iter()).enumerate() {
417 let raw = subscript_value(tensor, idx);
418 let coerced = coerce_subscript_value(raw, dim_index + 1, dim)?;
419 let term = coerced
420 .checked_sub(1)
421 .and_then(|v| v.checked_mul(strides[dim_index]))
422 .ok_or_else(|| sub2ind_bounds_error("Index exceeds array dimensions."))?;
423 offset = offset
424 .checked_add(term)
425 .ok_or_else(|| sub2ind_bounds_error("Index exceeds array dimensions."))?;
426 }
427 output.push((offset + 1) as f64);
428 }
429
430 Ok((output, target_shape))
431}
432
433fn tensor_element_len(tensor: &Tensor) -> usize {
434 tensor.len()
435}
436
437enum SubscriptValue {
438 Float(f64),
439 Integer(IntValue),
440}
441
442fn subscript_value(tensor: &Tensor, idx: usize) -> SubscriptValue {
443 if let Some(storage) = tensor.integer_storage() {
444 let index = if storage.len() == 1 { 0 } else { idx };
445 return SubscriptValue::Integer(
446 storage
447 .value_at(index)
448 .expect("subscript index is within integer storage bounds"),
449 );
450 }
451 if tensor::is_scalar_tensor(tensor) {
452 SubscriptValue::Float(tensor::tensor_value_f64(tensor, 0))
453 } else {
454 SubscriptValue::Float(tensor::tensor_value_f64(tensor, idx))
455 }
456}
457
458fn coerce_subscript_value(
459 value: SubscriptValue,
460 dim_number: usize,
461 dim_size: usize,
462) -> crate::BuiltinResult<usize> {
463 match value {
464 SubscriptValue::Float(value) => coerce_subscript(value, dim_number, dim_size),
465 SubscriptValue::Integer(value) => coerce_integer_subscript(&value, dim_number, dim_size),
466 }
467}
468
469fn coerce_integer_subscript(
470 value: &IntValue,
471 dim_number: usize,
472 dim_size: usize,
473) -> crate::BuiltinResult<usize> {
474 let Some(index) = value.try_to_usize() else {
475 return Err(sub2ind_error(
476 "Subscript indices must either be real positive integers or logicals.",
477 ));
478 };
479 if index < 1 {
480 return Err(sub2ind_error(
481 "Subscript indices must either be real positive integers or logicals.",
482 ));
483 }
484 if index > dim_size {
485 return Err(dimension_bounds_error(dim_number));
486 }
487 Ok(index)
488}
489
490fn coerce_subscript(value: f64, dim_number: usize, dim_size: usize) -> crate::BuiltinResult<usize> {
491 if !value.is_finite() {
492 return Err(sub2ind_error(
493 "Subscript indices must either be real positive integers or logicals.",
494 ));
495 }
496 let rounded = value.round();
497 if (rounded - value).abs() > f64::EPSILON {
498 return Err(sub2ind_error(
499 "Subscript indices must either be real positive integers or logicals.",
500 ));
501 }
502 if rounded < 1.0 {
503 return Err(sub2ind_error(
504 "Subscript indices must either be real positive integers or logicals.",
505 ));
506 }
507 if !fits_positive_platform_index(rounded) {
508 return Err(sub2ind_error(
509 "Subscript indices exceed the maximum supported index range.",
510 ));
511 }
512 if rounded > dim_size as f64 {
513 return Err(dimension_bounds_error(dim_number));
514 }
515 Ok(rounded as usize)
516}
517
518fn dimension_bounds_error(dim_number: usize) -> RuntimeError {
519 let message = match dim_number {
520 1 => format!("Index exceeds the number of rows in dimension {dim_number}."),
521 2 => format!("Index exceeds the number of columns in dimension {dim_number}."),
522 3 => format!("Index exceeds the number of pages in dimension {dim_number}."),
523 _ => "Index exceeds array dimensions.".to_string(),
524 };
525 sub2ind_bounds_error(message)
526}
527
528fn build_host_value(data: Vec<f64>, shape: Option<Vec<usize>>) -> crate::BuiltinResult<Value> {
529 let shape = shape.unwrap_or_else(|| vec![1, 1]);
530 if data.len() == 1 && tensor::element_count(&shape) == 1 {
531 Ok(Value::Num(data[0]))
532 } else {
533 let tensor = Tensor::new(data, shape).map_err(|e| {
534 sub2ind_internal_error(format!("Unable to construct sub2ind output: {e}"))
535 })?;
536 Ok(Value::Tensor(tensor))
537 }
538}
539
540fn sub2ind_error(message: impl Into<String>) -> RuntimeError {
541 sub2ind_input_error(message)
542}
543
544#[cfg(test)]
545pub(crate) mod tests {
546 use super::*;
547 use crate::builtins::common::test_support;
548 use futures::executor::block_on;
549 use runmat_builtins::Type;
550 use runmat_value::{IntValue, IntegerStorage, Tensor, Value};
551
552 fn sub2ind_builtin(dims_val: Value, rest: Vec<Value>) -> crate::BuiltinResult<Value> {
553 block_on(super::sub2ind_builtin(dims_val, rest))
554 }
555
556 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
557 #[test]
558 fn converts_scalar_indices() {
559 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
560 let result =
561 sub2ind_builtin(Value::Tensor(dims), vec![Value::Num(2.0), Value::Num(3.0)]).unwrap();
562 assert_eq!(result, Value::Num(8.0));
563 }
564
565 #[test]
566 fn sub2ind_type_scalar_outputs_num() {
567 assert_eq!(
568 sub2ind_type(
569 &[Type::Tensor { shape: None }, Type::Num, Type::Int],
570 &ResolveContext::new(Vec::new()),
571 ),
572 Type::Num
573 );
574 }
575
576 #[test]
577 fn sub2ind_type_vector_outputs_tensor() {
578 let subs = Type::Tensor {
579 shape: Some(vec![Some(3), Some(1)]),
580 };
581 assert_eq!(
582 sub2ind_type(
583 &[Type::Tensor { shape: None }, subs.clone(), Type::Num],
584 &ResolveContext::new(Vec::new()),
585 ),
586 Type::Tensor {
587 shape: Some(vec![Some(3), Some(1)])
588 }
589 );
590 }
591
592 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
593 #[test]
594 fn broadcasts_scalars_over_vectors() {
595 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
596 let rows = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
597 let result = sub2ind_builtin(
598 Value::Tensor(dims),
599 vec![Value::Tensor(rows), Value::Num(4.0)],
600 )
601 .unwrap();
602 match result {
603 Value::Tensor(t) => {
604 assert_eq!(t.shape, vec![3, 1]);
605 assert_eq!(t.materialize_f64(), vec![10.0, 11.0, 12.0]);
606 }
607 other => panic!("expected tensor result, got {other:?}"),
608 }
609 }
610
611 #[test]
612 fn sub2ind_typed_integer_subscripts_read_exact_storage() {
613 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
614 let rows = Tensor::new_integer(IntegerStorage::U16(vec![1, 2, 3]), vec![3, 1]).unwrap();
615 let cols = Tensor::new_integer(IntegerStorage::I16(vec![4]), vec![1, 1]).unwrap();
616
617 let result = sub2ind_builtin(
618 Value::Tensor(dims),
619 vec![Value::Tensor(rows), Value::Tensor(cols)],
620 )
621 .unwrap();
622
623 match result {
624 Value::Tensor(tensor) => {
625 assert_eq!(tensor.shape, vec![3, 1]);
626 assert_eq!(tensor.materialize_f64(), vec![10.0, 11.0, 12.0]);
627 }
628 other => panic!("expected tensor result, got {other:?}"),
629 }
630 }
631
632 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
633 #[test]
634 fn handles_three_dimensions() {
635 let dims = Tensor::new(vec![2.0, 3.0, 4.0], vec![1, 3]).unwrap();
636 let row = Tensor::new(vec![1.0, 1.0], vec![1, 2]).unwrap();
637 let col = Tensor::new(vec![2.0, 3.0], vec![1, 2]).unwrap();
638 let page = Tensor::new(vec![1.0, 2.0], vec![1, 2]).unwrap();
639 let result = sub2ind_builtin(
640 Value::Tensor(dims),
641 vec![Value::Tensor(row), Value::Tensor(col), Value::Tensor(page)],
642 )
643 .unwrap();
644 match result {
645 Value::Tensor(t) => {
646 assert_eq!(t.shape, vec![1, 2]);
647 assert_eq!(t.materialize_f64(), vec![3.0, 11.0]);
648 }
649 other => panic!("expected tensor result, got {other:?}"),
650 }
651 }
652
653 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
654 #[test]
655 fn rejects_out_of_range_subscripts() {
656 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
657 let err = sub2ind_builtin(Value::Tensor(dims), vec![Value::Num(4.0), Value::Num(1.0)])
658 .unwrap_err();
659 assert!(
660 err.to_string().contains("Index exceeds"),
661 "expected index bounds error, got {err}"
662 );
663 assert_eq!(
664 err.identifier(),
665 super::SUB2IND_ERROR_INDEX_BOUNDS.identifier
666 );
667 }
668
669 #[test]
670 fn sub2ind_accepts_every_integer_class_and_returns_double() {
671 for prototype in [
672 IntegerStorage::I8(Vec::new()),
673 IntegerStorage::I16(Vec::new()),
674 IntegerStorage::I32(Vec::new()),
675 IntegerStorage::I64(Vec::new()),
676 IntegerStorage::U8(Vec::new()),
677 IntegerStorage::U16(Vec::new()),
678 IntegerStorage::U32(Vec::new()),
679 IntegerStorage::U64(Vec::new()),
680 ] {
681 let typed = |values: &[i8], shape| {
682 let values = values
683 .iter()
684 .map(|value| prototype.cast_exact_assignment(&IntValue::I8(*value)))
685 .collect();
686 Tensor::new_integer(
687 prototype
688 .from_same_class_values(values)
689 .expect("same-class values"),
690 shape,
691 )
692 .expect("typed tensor")
693 };
694 let result = sub2ind_builtin(
695 Value::Tensor(typed(&[2, 3], vec![1, 2])),
696 vec![
697 Value::Tensor(typed(&[1, 2], vec![1, 2])),
698 Value::Tensor(typed(&[3], vec![1, 1])),
699 ],
700 )
701 .expect("sub2ind");
702 let Value::Tensor(result) = result else {
703 panic!("expected double tensor output");
704 };
705 assert_eq!(result.shape, vec![1, 2]);
706 assert_eq!(result.materialize_f64(), vec![5.0, 6.0]);
707 assert!(result.integer_storage().is_none());
708 }
709 }
710
711 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
712 #[test]
713 fn rejects_shape_mismatch() {
714 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
715 let rows = Tensor::new(vec![1.0, 2.0], vec![2, 1]).unwrap();
716 let cols = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
717 let err = sub2ind_builtin(
718 Value::Tensor(dims),
719 vec![Value::Tensor(rows), Value::Tensor(cols)],
720 )
721 .unwrap_err();
722 assert!(
723 err.to_string().contains("same size"),
724 "expected size mismatch error, got {err}"
725 );
726 assert_eq!(
727 err.identifier(),
728 super::SUB2IND_ERROR_INVALID_INPUT.identifier
729 );
730 }
731
732 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
733 #[test]
734 fn rejects_non_integer_subscripts() {
735 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
736 let err = sub2ind_builtin(Value::Tensor(dims), vec![Value::Num(1.5), Value::Num(1.0)])
737 .unwrap_err();
738 assert!(
739 err.to_string().contains("real positive integers"),
740 "expected integer coercion error, got {err}"
741 );
742 assert_eq!(
743 err.identifier(),
744 super::SUB2IND_ERROR_INVALID_INPUT.identifier
745 );
746 }
747
748 #[test]
749 fn rejects_oversized_float_subscripts_before_casting() {
750 let dims = Value::Int(IntValue::U64(usize::MAX.saturating_sub(1) as u64));
751
752 let err = sub2ind_builtin(dims.clone(), vec![Value::Num(1.0e300)])
753 .expect_err("huge float subscript must reject");
754 assert_eq!(
755 err.identifier(),
756 super::SUB2IND_ERROR_INVALID_INPUT.identifier
757 );
758
759 let err = sub2ind_builtin(dims, vec![Value::Num(usize::MAX as f64)])
760 .expect_err("platform boundary float subscript must reject");
761 assert_eq!(
762 err.identifier(),
763 super::SUB2IND_ERROR_INVALID_INPUT.identifier
764 );
765 }
766
767 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
768 #[test]
769 fn accepts_integer_value_variants() {
770 let dims = Value::Tensor(Tensor::new(vec![3.0], vec![1, 1]).unwrap());
771 let result = sub2ind_builtin(dims, vec![Value::Int(IntValue::I32(2))]).expect("sub2ind");
772 assert_eq!(result, Value::Num(2.0));
773 }
774
775 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
776 #[test]
777 fn sub2ind_gpu_roundtrip() {
778 test_support::with_test_provider(|provider| {
779 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
780 let rows = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
781 let cols = Tensor::new(vec![4.0, 4.0, 4.0], vec![3, 1]).unwrap();
782
783 let dims_handle = provider
784 .upload(&HostTensorView {
785 data: &dims.materialize_f64(),
786 shape: &dims.shape,
787 })
788 .expect("upload dims");
789 let rows_handle = provider
790 .upload(&HostTensorView {
791 data: &rows.materialize_f64(),
792 shape: &rows.shape,
793 })
794 .expect("upload rows");
795 let cols_handle = provider
796 .upload(&HostTensorView {
797 data: &cols.materialize_f64(),
798 shape: &cols.shape,
799 })
800 .expect("upload cols");
801
802 let result = sub2ind_builtin(
803 Value::GpuTensor(dims_handle),
804 vec![Value::GpuTensor(rows_handle), Value::GpuTensor(cols_handle)],
805 )
806 .expect("sub2ind");
807
808 match result {
809 Value::GpuTensor(handle) => {
810 let gathered = test_support::gather(Value::GpuTensor(handle)).unwrap();
811 assert_eq!(gathered.shape, vec![3, 1]);
812 assert_eq!(gathered.materialize_f64(), vec![10.0, 11.0, 12.0]);
813 }
814 other => panic!("expected gpu tensor, got {other:?}"),
815 }
816 });
817 }
818
819 #[test]
820 fn sub2ind_integer_gpu_inputs_fall_back_exactly_to_resident_double() {
821 test_support::with_test_provider(|provider| {
822 let dims =
823 Tensor::new_integer(IntegerStorage::U64(vec![2, 3]), vec![1, 2]).expect("dims");
824 let rows = provider
825 .upload_integer(&runmat_accelerate_api::HostIntegerTensorView {
826 data: runmat_accelerate_api::HostIntegerDataView::U64(&[1, 2]),
827 shape: &[1, 2],
828 })
829 .expect("rows");
830 let cols = provider
831 .upload_integer(&runmat_accelerate_api::HostIntegerTensorView {
832 data: runmat_accelerate_api::HostIntegerDataView::U64(&[3]),
833 shape: &[1, 1],
834 })
835 .expect("cols");
836 let result = sub2ind_builtin(
837 Value::Tensor(dims),
838 vec![Value::GpuTensor(rows), Value::GpuTensor(cols)],
839 )
840 .expect("sub2ind");
841 let Value::GpuTensor(handle) = &result else {
842 panic!("expected resident double output");
843 };
844 assert_eq!(runmat_accelerate_api::handle_integer_type(handle), None);
845 let result = test_support::gather(result).expect("gather result");
846 assert_eq!(result.shape, vec![1, 2]);
847 assert_eq!(result.materialize_f64(), vec![5.0, 6.0]);
848 assert!(result.integer_storage().is_none());
849 });
850 }
851
852 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
853 #[test]
854 #[cfg(feature = "wgpu")]
855 fn sub2ind_wgpu_matches_cpu() {
856 let _ = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
857 runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
858 );
859 let Some(provider) = runmat_accelerate_api::provider() else {
860 panic!("wgpu provider not available");
861 };
862
863 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
864 let rows = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
865 let cols = Tensor::new(vec![4.0, 4.0, 4.0], vec![3, 1]).unwrap();
866
867 let cpu = sub2ind_builtin(
868 Value::Tensor(dims.clone()),
869 vec![Value::Tensor(rows.clone()), Value::Tensor(cols.clone())],
870 )
871 .expect("cpu sub2ind");
872
873 let rows_handle = provider
874 .upload(&HostTensorView {
875 data: &rows.materialize_f64(),
876 shape: &rows.shape,
877 })
878 .expect("upload rows");
879 let cols_handle = provider
880 .upload(&HostTensorView {
881 data: &cols.materialize_f64(),
882 shape: &cols.shape,
883 })
884 .expect("upload cols");
885
886 let result = sub2ind_builtin(
887 Value::Tensor(dims),
888 vec![Value::GpuTensor(rows_handle), Value::GpuTensor(cols_handle)],
889 )
890 .expect("wgpu sub2ind");
891
892 let gathered = test_support::gather(result).expect("gather");
893 let expected = match cpu {
894 Value::Tensor(t) => t,
895 Value::Num(v) => Tensor::new(vec![v], vec![1, 1]).unwrap(),
896 other => panic!("unexpected cpu result {other:?}"),
897 };
898 assert_eq!(gathered.shape, expected.shape);
899 assert_eq!(gathered.materialize_f64(), expected.materialize_f64());
900 }
901}