1use runmat_builtins::{
4 BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
5 BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
6 ResolveContext, Type,
7};
8use runmat_builtins::{
9 BuiltinIntegerBackendRule, BuiltinIntegerCapabilityDescriptor, BuiltinIntegerComputationDomain,
10 BuiltinIntegerInputAvailability, BuiltinIntegerInputCapability, BuiltinIntegerOutputClassRule,
11 BuiltinIntegerOverflowRule, BuiltinIntegerOverloadKind, BuiltinIntegerScalarDoubleRule,
12};
13use runmat_macros::runtime_builtin;
14use runmat_value::{IntValue, Tensor, Value};
15
16use super::common::{
17 build_strides, dims_from_tokens, fits_positive_platform_index, materialize_value, parse_dims,
18 total_elements,
19};
20use crate::builtins::array::type_resolvers::size_vector_len;
21use crate::builtins::common::arg_tokens::tokens_from_context;
22use crate::builtins::common::spec::{
23 BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
24 ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
25};
26use crate::builtins::common::{gpu_helpers, tensor};
27use crate::{build_runtime_error, make_cell, RuntimeError};
28
29#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::array::indexing::ind2sub")]
30pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
31 name: "ind2sub",
32 op_kind: GpuOpKind::Custom("indexing"),
33 supported_precisions: &[ScalarType::F64],
34 broadcast: BroadcastSemantics::Matlab,
35 provider_hooks: &[ProviderHook::Custom("ind2sub")],
36 constant_strategy: ConstantStrategy::InlineLiteral,
37 residency: ResidencyPolicy::NewHandle,
38 nan_mode: ReductionNaN::Include,
39 two_pass_threshold: None,
40 workgroup_size: None,
41 accepts_nan_mode: false,
42 notes: "A binary64 WGPU provider executes admitted `ind2sub` calls entirely on-device. Kernel or adapter limits and other providers use exact host fallback; results return to the source owner only when that owner can truthfully represent double output.",
43};
44
45#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::array::indexing::ind2sub")]
46pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
47 name: "ind2sub",
48 shape: ShapeRequirements::Any,
49 constant_strategy: ConstantStrategy::InlineLiteral,
50 elementwise: None,
51 reduction: None,
52 emits_nan: false,
53 notes: "Index conversion is eager and does not participate in fusion today.",
54};
55
56fn ind2sub_type(args: &[Type], ctx: &ResolveContext) -> Type {
57 let Some(dims) = args.first() else {
58 return Type::Unknown;
59 };
60 let length = dims_from_tokens(&tokens_from_context(ctx))
61 .map(|values| values.len())
62 .or_else(|| size_vector_len(dims));
63 Type::Cell {
64 element_type: Some(Box::new(Type::tensor())),
65 length,
66 }
67}
68
69const BUILTIN_NAME: &str = "ind2sub";
70
71const IND2SUB_SIZE_CLASSES: [BuiltinIntegerInputCapability; 2] = [
72 BuiltinIntegerInputCapability {
73 name: "size vector",
74 classes: &crate::builtins::common::integer_capability::ALL_INTEGER_CLASSES,
75 availability: BuiltinIntegerInputAvailability::Documented,
76 scalar_double: BuiltinIntegerScalarDoubleRule::Allowed,
77 notes: "The size vector accepts single, double, and every integer class, must contain at least two positive integral elements, and is read exactly from authoritative storage.",
78 },
79 BuiltinIntegerInputCapability {
80 name: "linear indices",
81 classes: &crate::builtins::common::integer_capability::ALL_INTEGER_CLASSES,
82 availability: BuiltinIntegerInputAvailability::Documented,
83 scalar_double: BuiltinIntegerScalarDoubleRule::Allowed,
84 notes: "Linear indices accept single, double, and every integer class; outputs are double and preserve the index array shape.",
85 },
86];
87
88pub const IND2SUB_INTEGER_CAPABILITIES: [BuiltinIntegerCapabilityDescriptor; 1] =
89 [BuiltinIntegerCapabilityDescriptor {
90 form: "[I1,...,In] = ind2sub(integer_size, integer_indices)",
91 inputs: &IND2SUB_SIZE_CLASSES,
92 computation_domain: BuiltinIntegerComputationDomain::Structural,
93 output_class: BuiltinIntegerOutputClassRule::Double,
94 overflow: BuiltinIntegerOverflowRule::Error,
95 backend: BuiltinIntegerBackendRule::HostAndGpu,
96 overload: BuiltinIntegerOverloadKind::Multiple,
97 notes: "Requested output count controls trailing-dimension collapse; indices beyond prod(sz) expand the final effective dimension. Resident inputs always produce double output and return to the exact source owner only when that owner can physically preserve binary64.",
98 }];
99
100const IND2SUB_OUTPUT_CELL: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
101 name: "subs",
102 ty: BuiltinParamType::Any,
103 arity: BuiltinParamArity::Required,
104 default: None,
105 description: "Cell array containing one subscript output per dimension.",
106}];
107
108const IND2SUB_INPUTS: [BuiltinParamDescriptor; 2] = [
109 BuiltinParamDescriptor {
110 name: "sz",
111 ty: BuiltinParamType::SizeArg,
112 arity: BuiltinParamArity::Required,
113 default: None,
114 description: "Size vector describing source array dimensions.",
115 },
116 BuiltinParamDescriptor {
117 name: "ind",
118 ty: BuiltinParamType::Any,
119 arity: BuiltinParamArity::Required,
120 default: None,
121 description: "Linear indices to convert into per-dimension subscripts.",
122 },
123];
124
125const IND2SUB_SIGNATURES: [BuiltinSignatureDescriptor; 1] = [BuiltinSignatureDescriptor {
126 label: "subs = ind2sub(sz, ind)",
127 inputs: &IND2SUB_INPUTS,
128 outputs: &IND2SUB_OUTPUT_CELL,
129}];
130
131const IND2SUB_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
132 code: "RM.IND2SUB.INVALID_INPUT",
133 identifier: Some("RunMat:ind2sub:InvalidInput"),
134 when: "Size vector or linear index inputs are malformed or unsupported.",
135 message: "ind2sub: invalid input arguments",
136};
137
138const IND2SUB_ERROR_PROVIDER: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
139 code: "RM.IND2SUB.PROVIDER",
140 identifier: Some("RunMat:ind2sub:ProviderError"),
141 when: "Provider-side ind2sub execution fails or returns malformed outputs.",
142 message: "ind2sub: provider execution failed",
143};
144
145const IND2SUB_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
146 code: "RM.IND2SUB.INTERNAL",
147 identifier: Some("RunMat:ind2sub:InternalError"),
148 when: "Internal tensor/materialization logic fails while building outputs.",
149 message: "ind2sub: internal error",
150};
151
152const IND2SUB_ERRORS: [BuiltinErrorDescriptor; 3] = [
153 IND2SUB_ERROR_INVALID_INPUT,
154 IND2SUB_ERROR_PROVIDER,
155 IND2SUB_ERROR_INTERNAL,
156];
157
158pub const IND2SUB_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
159 signatures: &IND2SUB_SIGNATURES,
160 output_mode: BuiltinOutputMode::ByRequestedOutputCount,
161 completion_policy: BuiltinCompletionPolicy::Public,
162 errors: &IND2SUB_ERRORS,
163};
164
165fn ind2sub_error_with_message(
166 message: impl Into<String>,
167 error: &'static BuiltinErrorDescriptor,
168) -> RuntimeError {
169 let mut builder = build_runtime_error(message).with_builtin(BUILTIN_NAME);
170 if let Some(identifier) = error.identifier {
171 builder = builder.with_identifier(identifier);
172 }
173 builder.build()
174}
175
176fn ind2sub_input_error(message: impl Into<String>) -> RuntimeError {
177 ind2sub_error_with_message(message, &IND2SUB_ERROR_INVALID_INPUT)
178}
179
180fn ind2sub_internal_error(message: impl Into<String>) -> RuntimeError {
181 ind2sub_error_with_message(message, &IND2SUB_ERROR_INTERNAL)
182}
183
184fn ind2sub_provider_error(message: impl Into<String>) -> RuntimeError {
185 ind2sub_error_with_message(message, &IND2SUB_ERROR_PROVIDER)
186}
187
188#[runtime_builtin(
189 name = "ind2sub",
190 category = "array/indexing",
191 summary = "Convert linear indices to subscripts.",
192 keywords = "ind2sub,linear index,subscripts,column major,gpu indexing",
193 accel = "custom",
194 type_resolver(ind2sub_type),
195 descriptor(crate::builtins::array::indexing::ind2sub::IND2SUB_DESCRIPTOR),
196 integer_capabilities(crate::builtins::array::indexing::ind2sub::IND2SUB_INTEGER_CAPABILITIES),
197 builtin_path = "crate::builtins::array::indexing::ind2sub"
198)]
199async fn ind2sub_builtin(dims_val: Value, indices_val: Value) -> crate::BuiltinResult<Value> {
200 let output_source = match (&indices_val, &dims_val) {
201 (Value::GpuTensor(handle), _) => Some(handle.clone()),
202 (_, Value::GpuTensor(handle)) => Some(handle.clone()),
203 _ => None,
204 };
205 let (dims_value, dims_was_gpu) = materialize_value(dims_val, "ind2sub").await?;
206 let dims = parse_dims(&dims_value, "ind2sub").await?;
207 if dims.len() < 2 {
208 return Err(ind2sub_error(
209 "Size vector must have at least two elements.",
210 ));
211 }
212
213 let requested_outputs = crate::output_count::current_output_count().unwrap_or(dims.len());
214 if requested_outputs == 0 {
215 return Ok(Value::OutputList(Vec::new()));
216 }
217 let effective_dims = effective_output_dims(&dims, requested_outputs)?;
218 let total = total_elements(&effective_dims, "ind2sub")?;
219 let strides = build_strides(&effective_dims, "ind2sub")?;
220
221 if let Some(result) = try_gpu_ind2sub(&effective_dims, &strides, total, &indices_val)? {
222 return Ok(result);
223 }
224
225 let (indices_value, indices_was_gpu) = materialize_value(indices_val, "ind2sub").await?;
226 let indices_tensor = tensor::value_into_tensor_for("ind2sub", indices_value)
227 .map_err(|message| ind2sub_error(message))?;
228
229 let subscripts = compute_subscripts(&effective_dims, &strides, &indices_tensor)?;
230
231 let host_tensors = subscripts;
232 let outputs = if dims_was_gpu || indices_was_gpu {
233 restore_host_outputs_transactionally(output_source.as_ref(), &host_tensors)?
234 } else {
235 host_tensors
236 .into_iter()
237 .map(tensor::tensor_into_value)
238 .collect()
239 };
240
241 finish_outputs(outputs)
242}
243
244fn restore_host_outputs_transactionally(
245 source: Option<&runmat_accelerate_api::GpuTensorHandle>,
246 host_tensors: &[Tensor],
247) -> crate::BuiltinResult<Vec<Value>> {
248 let host_outputs = || {
249 host_tensors
250 .iter()
251 .cloned()
252 .map(tensor::tensor_into_value)
253 .collect::<Vec<_>>()
254 };
255 let Some(source) = source else {
256 return Ok(host_outputs());
257 };
258 let mut restored = Vec::with_capacity(host_tensors.len());
259 for tensor in host_tensors {
260 match gpu_helpers::restore_class_preserving_value(
261 source,
262 Value::Tensor(tensor.clone()),
263 BUILTIN_NAME,
264 ) {
265 Ok(Value::GpuTensor(handle)) => restored.push(Value::GpuTensor(handle)),
266 Ok(_) => {
267 free_restored_outputs(&restored, source);
268 return Ok(host_outputs());
269 }
270 Err(error) => {
271 free_restored_outputs(&restored, source);
272 return Err(error);
273 }
274 }
275 }
276 Ok(restored)
277}
278
279fn free_restored_outputs(outputs: &[Value], source: &runmat_accelerate_api::GpuTensorHandle) {
280 for output in outputs {
281 if let Value::GpuTensor(handle) = output {
282 gpu_helpers::free_unprotected_exact_owner(handle, &[source]);
283 }
284 }
285}
286
287fn effective_output_dims(
288 dims: &[usize],
289 requested_outputs: usize,
290) -> crate::BuiltinResult<Vec<usize>> {
291 if requested_outputs == 0 {
292 return Ok(Vec::new());
293 }
294 if requested_outputs >= dims.len() {
295 let mut effective = dims.to_vec();
296 effective.resize(requested_outputs, 1);
297 return Ok(effective);
298 }
299 let mut effective = dims[..requested_outputs.saturating_sub(1)].to_vec();
300 let trailing = dims[requested_outputs.saturating_sub(1)..]
301 .iter()
302 .try_fold(1usize, |product, dim| product.checked_mul(*dim))
303 .ok_or_else(|| {
304 ind2sub_input_error("ind2sub: collapsed dimensions exceed platform limits")
305 })?;
306 effective.push(trailing);
307 Ok(effective)
308}
309
310fn finish_outputs(outputs: Vec<Value>) -> crate::BuiltinResult<Value> {
311 if crate::output_count::current_output_count().is_some() {
312 return Ok(Value::OutputList(outputs));
313 }
314 let len = outputs.len();
315 make_cell(outputs, 1, len).map_err(|message| ind2sub_error(message))
316}
317
318fn try_gpu_ind2sub(
319 dims: &[usize],
320 strides: &[usize],
321 total: usize,
322 indices: &Value,
323) -> crate::BuiltinResult<Option<Value>> {
324 #[cfg(target_arch = "wasm32")]
325 {
326 let _ = (dims, strides, total, indices);
327 Ok(None)
328 }
329 #[cfg(not(target_arch = "wasm32"))]
330 {
331 let provider = match runmat_accelerate_api::provider() {
332 Some(p) => p,
333 None => return Ok(None),
334 };
335 if !provider.supports_ind2sub() {
336 return Ok(None);
337 }
338 let handle = match indices {
339 Value::GpuTensor(handle) => handle,
340 _ => return Ok(None),
341 };
342 let provider = match gpu_helpers::exact_provider_for_handle(handle) {
343 Some(provider) => provider,
344 None => {
345 return Err(ind2sub_provider_error(
346 "ind2sub: no acceleration provider owns the index handle.",
347 ))
348 }
349 };
350 if !provider.supports_ind2sub() {
351 return Ok(None);
352 }
353 if runmat_accelerate_api::handle_integer_type(handle).is_some() {
354 return Ok(None);
355 }
356 let expected_precision = provider.precision();
357 if expected_precision != runmat_accelerate_api::ProviderPrecision::F64 {
358 return Ok(None);
359 }
360 if runmat_accelerate_api::handle_storage(handle)
361 != runmat_accelerate_api::GpuTensorStorage::Real
362 || runmat_accelerate_api::handle_is_logical(handle)
363 || runmat_accelerate_api::handle_precision(handle) != Some(expected_precision)
364 || !gpu_helpers::gpu_class_metadata_matches(
365 handle,
366 Some(expected_precision),
367 None,
368 false,
369 )
370 {
371 return Err(ind2sub_provider_error(
372 "ind2sub: index handle metadata contradicts its provider payload.",
373 ));
374 }
375 if dims.len() != strides.len() {
376 return Err(ind2sub_error("Size vector must have at least one element."));
377 }
378 if dims.iter().any(|&d| d > u32::MAX as usize)
379 || strides.iter().any(|&s| s > u32::MAX as usize)
380 || total > u32::MAX as usize
381 {
382 return Ok(None);
383 }
384 let len = if handle.shape.is_empty() {
385 1usize
386 } else {
387 handle.shape.iter().copied().product()
388 };
389 if total == 0 && len > 0 {
390 return Err(ind2sub_error(
391 "Index exceeds number of array elements. Index must not exceed 0.",
392 ));
393 }
394 if len > u32::MAX as usize {
395 return Ok(None);
396 }
397 let output_shape = if handle.shape.is_empty() {
398 vec![len, 1]
399 } else {
400 handle.shape.clone()
401 };
402 let source_metadata = gpu_helpers::snapshot_handle_metadata(handle);
403 let result = provider.ind2sub(dims, strides, handle, total, len, &output_shape);
404 gpu_helpers::restore_handle_metadata(handle, &source_metadata);
405 match result {
406 Ok(mut handles) => {
407 if handles.len() != dims.len() {
408 for output in &handles {
409 gpu_helpers::free_unprotected_exact_owner(output, &[handle]);
410 }
411 return Err(ind2sub_provider_error(
412 "ind2sub: provider returned an unexpected number of outputs.",
413 ));
414 }
415 let valid = handles.iter().enumerate().all(|(index, output)| {
416 output.shape == output_shape
417 && output.device_id == provider.device_id()
418 && gpu_helpers::exact_provider_for_handle(output)
419 .is_some_and(|owner| std::ptr::eq(owner, provider))
420 && runmat_accelerate_api::handle_storage(output)
421 == runmat_accelerate_api::GpuTensorStorage::Real
422 && runmat_accelerate_api::handle_precision(output)
423 == Some(expected_precision)
424 && runmat_accelerate_api::handle_integer_type(output).is_none()
425 && !runmat_accelerate_api::handle_is_logical(output)
426 && gpu_helpers::gpu_class_metadata_matches(
427 output,
428 Some(expected_precision),
429 None,
430 false,
431 )
432 && !gpu_helpers::same_gpu_handle(output, handle)
433 && handles[..index]
434 .iter()
435 .all(|prior| !gpu_helpers::same_gpu_handle(output, prior))
436 });
437 if !valid {
438 for output in &handles {
439 gpu_helpers::free_unprotected_exact_owner(output, &[handle]);
440 }
441 return Err(ind2sub_provider_error(
442 "ind2sub: provider returned an invalid output handle.",
443 ));
444 }
445 let provenance = runmat_accelerate_api::handle_provenance(handle)
446 .unwrap_or(runmat_accelerate_api::GpuHandleProvenance::Automatic);
447 for output in &mut handles {
448 runmat_accelerate_api::set_handle_provenance(output, provenance);
449 runmat_accelerate_api::mark_residency(output);
450 }
451 let values: Vec<Value> = handles.into_iter().map(Value::GpuTensor).collect();
452 finish_outputs(values).map(Some)
453 }
454 Err(err) => {
455 let message = err.to_string();
456 if message.contains("GPU kernel limits")
457 || message.contains("storage buffers")
458 || message.contains("bind group entries")
459 {
460 Ok(None)
461 } else {
462 Err(ind2sub_provider_error(message))
463 }
464 }
465 }
466 }
467}
468
469fn compute_subscripts(
470 dims: &[usize],
471 strides: &[usize],
472 indices: &Tensor,
473) -> crate::BuiltinResult<Vec<Tensor>> {
474 if strides.len() != dims.len() {
475 return Err(ind2sub_error("Size vector must have at least one element."));
476 }
477
478 let len = tensor::tensor_element_len(indices);
479 let mut outputs: Vec<Vec<f64>> = dims.iter().map(|_| Vec::with_capacity(len)).collect();
480
481 for value_index in 0..len {
482 let idx = coerce_linear_index_value(linear_index_value(indices, value_index))?;
483 let zero_based = idx - 1;
484 for (dim_index, (&dim, &stride)) in dims.iter().zip(strides.iter()).enumerate() {
485 let coord = if dim_index + 1 == dims.len() {
486 (zero_based / stride) + 1
487 } else {
488 ((zero_based / stride) % dim) + 1
489 };
490 outputs[dim_index].push(coord as f64);
491 }
492 }
493
494 let output_shape = if indices.shape.is_empty() {
495 vec![len, 1]
496 } else {
497 indices.shape.clone()
498 };
499
500 let mut tensors = Vec::with_capacity(dims.len());
501 for data in outputs {
502 let tensor = Tensor::new(data, output_shape.clone())
503 .map_err(|e| ind2sub_internal_error(format!("ind2sub: {e}")))?;
504 tensors.push(tensor);
505 }
506 Ok(tensors)
507}
508
509enum LinearIndexValue {
510 Float(f64),
511 Integer(IntValue),
512}
513
514fn linear_index_value(indices: &Tensor, value_index: usize) -> LinearIndexValue {
515 if let Some(storage) = indices.integer_storage() {
516 return LinearIndexValue::Integer(
517 storage
518 .value_at(value_index)
519 .expect("linear index is within integer storage bounds"),
520 );
521 }
522 LinearIndexValue::Float(tensor::tensor_value_f64(indices, value_index))
523}
524
525fn coerce_linear_index_value(value: LinearIndexValue) -> crate::BuiltinResult<usize> {
526 match value {
527 LinearIndexValue::Float(value) => coerce_linear_index(value),
528 LinearIndexValue::Integer(value) => coerce_integer_linear_index(&value),
529 }
530}
531
532fn coerce_integer_linear_index(value: &IntValue) -> crate::BuiltinResult<usize> {
533 let Some(coerced) = value.try_to_usize() else {
534 return Err(ind2sub_error("Linear indices must be positive integers."));
535 };
536 if coerced < 1 {
537 return Err(ind2sub_error("Linear indices must be positive integers."));
538 }
539 Ok(coerced)
540}
541
542fn coerce_linear_index(value: f64) -> crate::BuiltinResult<usize> {
543 if !value.is_finite() {
544 return Err(ind2sub_error("Linear indices must be positive integers."));
545 }
546 let rounded = value.round();
547 if (rounded - value).abs() > f64::EPSILON {
548 return Err(ind2sub_error("Linear indices must be positive integers."));
549 }
550 if rounded < 1.0 {
551 return Err(ind2sub_error("Linear indices must be positive integers."));
552 }
553 if !fits_positive_platform_index(rounded) {
554 return Err(ind2sub_error(
555 "Index exceeds maximum supported size for this platform.",
556 ));
557 }
558 let coerced = rounded as usize;
559 Ok(coerced)
560}
561
562fn ind2sub_error(message: impl Into<String>) -> RuntimeError {
563 ind2sub_input_error(message)
564}
565
566#[cfg(test)]
567pub(crate) mod tests {
568 use crate::builtins::common::test_support;
569 use futures::executor::block_on;
570 use runmat_accelerate_api::HostTensorView;
571 use runmat_builtins::{ResolveContext, Type};
572 use runmat_value::{IntValue, IntegerStorage, Tensor, Value};
573
574 fn ind2sub_builtin(dims_val: Value, indices_val: Value) -> crate::BuiltinResult<Value> {
575 block_on(super::ind2sub_builtin(dims_val, indices_val))
576 }
577
578 fn cell_to_vec(cell: &runmat_value::CellArray) -> Vec<Value> {
579 cell.data.clone()
580 }
581
582 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
583 #[test]
584 fn recovers_tensor_indices() {
585 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
586 let result = ind2sub_builtin(Value::Tensor(dims), Value::Num(8.0)).unwrap();
587 match result {
588 Value::Cell(cell) => {
589 let values = cell_to_vec(&cell);
590 assert_eq!(values.len(), 2);
591 assert_eq!(values[0], Value::Num(2.0));
592 assert_eq!(values[1], Value::Num(3.0));
593 }
594 other => panic!("expected cell output, got {other:?}"),
595 }
596 }
597
598 #[test]
599 fn ind2sub_type_infers_cell_length() {
600 let dims = Type::Tensor {
601 shape: Some(vec![Some(1), Some(3)]),
602 };
603 assert_eq!(
604 super::ind2sub_type(&[dims, Type::Num], &ResolveContext::new(Vec::new())),
605 Type::Cell {
606 element_type: Some(Box::new(Type::tensor())),
607 length: Some(3)
608 }
609 );
610 }
611
612 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
613 #[test]
614 fn handles_vector_indices() {
615 let dims = Tensor::new(vec![3.0, 5.0], vec![1, 2]).unwrap();
616 let idx = Tensor::new(vec![7.0, 8.0, 9.0], vec![1, 3]).unwrap();
617 let result =
618 ind2sub_builtin(Value::Tensor(dims), Value::Tensor(idx)).expect("ind2sub result");
619 match result {
620 Value::Cell(cell) => {
621 let values = cell_to_vec(&cell);
622 assert_eq!(values.len(), 2);
623 match &values[0] {
624 Value::Tensor(t) => {
625 assert_eq!(t.shape, vec![1, 3]);
626 assert_eq!(t.materialize_f64(), vec![1.0, 2.0, 3.0]);
627 }
628 other => panic!("expected tensor rows, got {other:?}"),
629 }
630 match &values[1] {
631 Value::Tensor(t) => {
632 assert_eq!(t.shape, vec![1, 3]);
633 assert_eq!(t.materialize_f64(), vec![3.0, 3.0, 3.0]);
634 }
635 other => panic!("expected tensor cols, got {other:?}"),
636 }
637 }
638 other => panic!("expected cell output, got {other:?}"),
639 }
640 }
641
642 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
643 #[test]
644 fn rejects_non_integer_linear_index_identifier() {
645 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
646 let err = ind2sub_builtin(Value::Tensor(dims), Value::Num(1.25))
647 .expect_err("expected non-integer index error");
648 assert_eq!(
649 err.identifier(),
650 super::IND2SUB_ERROR_INVALID_INPUT.identifier
651 );
652 }
653
654 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
655 #[test]
656 fn indices_beyond_nominal_size_expand_the_final_dimension() {
657 let dims = Tensor::new(vec![2.0, 2.0], vec![1, 2]).unwrap();
658 let result = ind2sub_builtin(Value::Tensor(dims), Value::Num(9.0)).expect("expanded index");
659 let Value::Cell(result) = result else {
660 panic!("expected subscript outputs");
661 };
662 assert_eq!(result.data, vec![Value::Num(1.0), Value::Num(5.0)]);
663 }
664
665 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
666 #[test]
667 fn recovers_three_dimensional_indices() {
668 let dims = Tensor::new(vec![2.0, 3.0, 4.0], vec![1, 3]).unwrap();
669 let idx = Tensor::new(vec![3.0, 11.0], vec![1, 2]).unwrap();
670 let result =
671 ind2sub_builtin(Value::Tensor(dims), Value::Tensor(idx)).expect("ind2sub result");
672 if let Value::Cell(cell) = result {
673 let values = cell_to_vec(&cell);
674 assert_eq!(values.len(), 3);
675 assert_eq!(
676 values[0],
677 Value::Tensor(Tensor::new(vec![1.0, 1.0], vec![1, 2]).unwrap())
678 );
679 assert_eq!(
680 values[1],
681 Value::Tensor(Tensor::new(vec![2.0, 3.0], vec![1, 2]).unwrap())
682 );
683 assert_eq!(
684 values[2],
685 Value::Tensor(Tensor::new(vec![1.0, 2.0], vec![1, 2]).unwrap())
686 );
687 } else {
688 panic!("expected cell output");
689 }
690 }
691
692 #[test]
693 fn requested_outputs_collapse_trailing_dimensions_and_extend_with_singletons() {
694 let dims = Value::Tensor(Tensor::new(vec![2.0, 3.0, 4.0], vec![1, 3]).unwrap());
695 let indices = Value::Tensor(Tensor::new(vec![3.0, 11.0], vec![1, 2]).unwrap());
696 {
697 let _outputs = crate::output_count::push_output_count(Some(2));
698 let result = ind2sub_builtin(dims.clone(), indices.clone()).expect("two outputs");
699 let Value::OutputList(outputs) = result else {
700 panic!("expected output list");
701 };
702 assert_eq!(outputs.len(), 2);
703 assert_eq!(
704 outputs[0],
705 Value::Tensor(Tensor::new(vec![1.0, 1.0], vec![1, 2]).unwrap())
706 );
707 assert_eq!(
708 outputs[1],
709 Value::Tensor(Tensor::new(vec![2.0, 6.0], vec![1, 2]).unwrap())
710 );
711 }
712 {
713 let _outputs = crate::output_count::push_output_count(Some(4));
714 let result = ind2sub_builtin(dims, indices).expect("four outputs");
715 let Value::OutputList(outputs) = result else {
716 panic!("expected output list");
717 };
718 assert_eq!(outputs.len(), 4);
719 assert_eq!(
720 outputs[3],
721 Value::Tensor(Tensor::new(vec![1.0, 1.0], vec![1, 2]).unwrap())
722 );
723 }
724 }
725
726 #[test]
727 fn scalar_size_vector_is_rejected_by_current_contract() {
728 let error =
729 ind2sub_builtin(Value::Num(4.0), Value::Num(1.0)).expect_err("scalar size must reject");
730 assert_eq!(error.identifier(), Some("RunMat:ind2sub:InvalidInput"));
731 }
732
733 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
734 #[test]
735 fn ind2sub_linear_indices_read_typed_integer_storage_exactly() {
736 let dims = Tensor::new(vec![2.0, 3.0, 4.0], vec![1, 3]).unwrap();
737 let idx =
738 Tensor::new_integer(IntegerStorage::U16(vec![3, 11]), vec![1, 2]).expect("indices");
739
740 let result =
741 ind2sub_builtin(Value::Tensor(dims), Value::Tensor(idx)).expect("ind2sub result");
742 let Value::Cell(cell) = result else {
743 panic!("expected cell output");
744 };
745 let values = cell_to_vec(&cell);
746 assert_eq!(values.len(), 3);
747 assert_eq!(
748 values[0],
749 Value::Tensor(Tensor::new(vec![1.0, 1.0], vec![1, 2]).unwrap())
750 );
751 assert_eq!(
752 values[1],
753 Value::Tensor(Tensor::new(vec![2.0, 3.0], vec![1, 2]).unwrap())
754 );
755 assert_eq!(
756 values[2],
757 Value::Tensor(Tensor::new(vec![1.0, 2.0], vec![1, 2]).unwrap())
758 );
759 }
760
761 #[test]
762 fn ind2sub_accepts_every_integer_class_and_returns_double() {
763 for prototype in [
764 IntegerStorage::I8(Vec::new()),
765 IntegerStorage::I16(Vec::new()),
766 IntegerStorage::I32(Vec::new()),
767 IntegerStorage::I64(Vec::new()),
768 IntegerStorage::U8(Vec::new()),
769 IntegerStorage::U16(Vec::new()),
770 IntegerStorage::U32(Vec::new()),
771 IntegerStorage::U64(Vec::new()),
772 ] {
773 let typed = |values: &[i8], shape| {
774 let values = values
775 .iter()
776 .map(|value| prototype.cast_exact_assignment(&IntValue::I8(*value)))
777 .collect();
778 Tensor::new_integer(
779 prototype
780 .from_same_class_values(values)
781 .expect("same-class values"),
782 shape,
783 )
784 .expect("typed tensor")
785 };
786 let result = ind2sub_builtin(
787 Value::Tensor(typed(&[2, 3], vec![1, 2])),
788 Value::Tensor(typed(&[5, 6], vec![1, 2])),
789 )
790 .expect("ind2sub");
791 let Value::Cell(result) = result else {
792 panic!("expected cell output");
793 };
794 let outputs = cell_to_vec(&result);
795 assert_eq!(outputs.len(), 2);
796 for (output, expected) in outputs.iter().zip([vec![1.0, 2.0], vec![3.0, 3.0]]) {
797 let Value::Tensor(output) = output else {
798 panic!("expected double tensor output");
799 };
800 assert_eq!(output.shape, vec![1, 2]);
801 assert_eq!(output.materialize_f64(), expected);
802 assert!(output.integer_storage().is_none());
803 }
804 }
805 }
806
807 #[test]
808 fn ind2sub_uses_wide_typed_index_exactly_before_double_output() {
809 let max = 9_007_199_254_740_992_u64;
810 let wide_index = max + 1;
811 assert_eq!(max as f64, wide_index as f64);
812
813 let dims =
814 Tensor::new_integer(IntegerStorage::U64(vec![2, max]), vec![1, 2]).expect("typed dims");
815 let index = Tensor::new_integer(IntegerStorage::U64(vec![wide_index]), vec![1, 1])
816 .expect("typed index");
817
818 let result = ind2sub_builtin(Value::Tensor(dims), Value::Tensor(index))
819 .expect("wide index remains exact through subscript computation");
820 let Value::Cell(result) = result else {
821 panic!("expected cell output");
822 };
823 let outputs = cell_to_vec(&result);
824 assert_eq!(outputs[0], Value::Num(1.0));
825 assert_eq!(outputs[1], Value::Num(4_503_599_627_370_497.0));
826 }
827
828 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
829 #[test]
830 fn expands_last_dimension_for_out_of_range_index() {
831 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
832 let result = ind2sub_builtin(Value::Tensor(dims), Value::Num(13.0)).expect("expanded");
833 let Value::Cell(result) = result else {
834 panic!("expected subscript outputs");
835 };
836 assert_eq!(result.data, vec![Value::Num(1.0), Value::Num(5.0)]);
837 }
838
839 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
840 #[test]
841 fn errors_on_zero_index() {
842 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
843 let err =
844 ind2sub_builtin(Value::Tensor(dims), Value::Num(0.0)).expect_err("expected failure");
845 assert!(
846 err.contains("Linear indices must be positive integers"),
847 "unexpected error: {err}"
848 );
849 }
850
851 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
852 #[test]
853 fn errors_on_fractional_index() {
854 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
855 let err =
856 ind2sub_builtin(Value::Tensor(dims), Value::Num(2.5)).expect_err("expected failure");
857 assert!(
858 err.contains("Linear indices must be positive integers"),
859 "unexpected error: {err}"
860 );
861 }
862
863 #[test]
864 fn rejects_oversized_float_linear_indices_before_casting() {
865 let dims = Tensor::new_integer(
866 IntegerStorage::U64(vec![usize::MAX.saturating_sub(1) as u64]),
867 vec![1, 1],
868 )
869 .unwrap();
870
871 let err = ind2sub_builtin(Value::Tensor(dims.clone()), Value::Num(1.0e300))
872 .expect_err("huge float index must reject");
873 assert_eq!(
874 err.identifier(),
875 super::IND2SUB_ERROR_INVALID_INPUT.identifier
876 );
877
878 let err = ind2sub_builtin(Value::Tensor(dims), Value::Num(usize::MAX as f64))
879 .expect_err("platform boundary float index must reject");
880 assert_eq!(
881 err.identifier(),
882 super::IND2SUB_ERROR_INVALID_INPUT.identifier
883 );
884 }
885
886 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
887 #[test]
888 fn errors_on_invalid_size_elements() {
889 let dims = Tensor::new(vec![3.5, 4.0], vec![1, 2]).unwrap();
890 let err = ind2sub_builtin(Value::Tensor(dims), Value::Num(5.0)).expect_err("expected fail");
891 assert!(
892 err.to_string()
893 .contains("Size arguments must be positive integers"),
894 "unexpected error: {err}"
895 );
896 }
897
898 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
899 #[test]
900 fn ind2sub_gpu_roundtrip() {
901 test_support::with_test_provider(|provider| {
902 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
903 let idx_tensor = Tensor::new(vec![10.0, 11.0], vec![2, 1]).unwrap();
904 let view = HostTensorView {
905 data: &idx_tensor.materialize_f64(),
906 shape: &idx_tensor.shape,
907 };
908 let handle = provider.upload(&view).expect("upload indices");
909 let result = ind2sub_builtin(Value::Tensor(dims), Value::GpuTensor(handle)).unwrap();
910 match result {
911 Value::Cell(cell) => {
912 let values = cell_to_vec(&cell);
913 assert_eq!(values.len(), 2);
914 match &values[0] {
915 Value::GpuTensor(_) => {}
916 other => panic!("expected gpu tensor output, got {other:?}"),
917 }
918 match &values[1] {
919 Value::GpuTensor(_) => {}
920 other => panic!("expected gpu tensor output, got {other:?}"),
921 }
922 let rows = test_support::gather(values[0].clone()).expect("gather rows");
923 assert_eq!(rows.shape, vec![2, 1]);
924 assert_eq!(rows.materialize_f64(), vec![1.0, 2.0]);
925 let cols = test_support::gather(values[1].clone()).expect("gather cols");
926 assert_eq!(cols.shape, vec![2, 1]);
927 assert_eq!(cols.materialize_f64(), vec![4.0, 4.0]);
928 }
929 other => panic!("expected cell output, got {other:?}"),
930 }
931 });
932 }
933
934 #[test]
935 fn ind2sub_integer_gpu_input_falls_back_exactly_to_resident_double() {
936 test_support::with_test_provider(|provider| {
937 let dims =
938 Tensor::new_integer(IntegerStorage::U64(vec![2, 3]), vec![1, 2]).expect("dims");
939 let indices = provider
940 .upload_integer(&runmat_accelerate_api::HostIntegerTensorView {
941 data: runmat_accelerate_api::HostIntegerDataView::U64(&[5, 6]),
942 shape: &[1, 2],
943 })
944 .expect("indices");
945 let result =
946 ind2sub_builtin(Value::Tensor(dims), Value::GpuTensor(indices)).expect("ind2sub");
947 let Value::Cell(result) = result else {
948 panic!("expected cell output");
949 };
950 let outputs = cell_to_vec(&result);
951 assert_eq!(outputs.len(), 2);
952 for (output, expected) in outputs.into_iter().zip([vec![1.0, 2.0], vec![3.0, 3.0]]) {
953 let Value::GpuTensor(handle) = &output else {
954 panic!("expected resident double output");
955 };
956 assert_eq!(runmat_accelerate_api::handle_integer_type(handle), None);
957 let output = test_support::gather(output).expect("gather output");
958 assert_eq!(output.shape, vec![1, 2]);
959 assert_eq!(output.materialize_f64(), expected);
960 assert!(output.integer_storage().is_none());
961 }
962 });
963 }
964
965 #[test]
966 fn ind2sub_scalar_resident_fallback_restores_double_to_the_exact_owner() {
967 test_support::with_test_provider(|provider| {
968 let indices = provider
969 .upload_integer(&runmat_accelerate_api::HostIntegerTensorView {
970 data: runmat_accelerate_api::HostIntegerDataView::U64(&[5]),
971 shape: &[1, 1],
972 })
973 .expect("scalar index");
974 let indices =
975 indices.with_provenance(runmat_accelerate_api::GpuHandleProvenance::Explicit);
976 let result = ind2sub_builtin(
977 Value::Tensor(Tensor::new(vec![2.0, 3.0], vec![1, 2]).unwrap()),
978 Value::GpuTensor(indices.clone()),
979 )
980 .expect("scalar fallback");
981 let Value::Cell(cell) = result else {
982 panic!("expected cell output");
983 };
984 for output in cell_to_vec(&cell) {
985 let Value::GpuTensor(handle) = output else {
986 panic!("scalar fallback must restore resident output");
987 };
988 assert!(super::gpu_helpers::exact_provider_for_handle(&handle)
989 .is_some_and(|owner| std::ptr::eq(owner, provider)));
990 assert_eq!(
991 runmat_accelerate_api::handle_provenance(&handle),
992 Some(runmat_accelerate_api::GpuHandleProvenance::Explicit)
993 );
994 let _ = provider.free(&handle);
995 }
996 let _ = provider.free(&indices);
997 });
998 }
999
1000 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
1001 #[test]
1002 #[cfg(feature = "wgpu")]
1003 fn ind2sub_wgpu_matches_cpu() {
1004 let provider_init = std::panic::catch_unwind(|| {
1005 runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
1006 runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
1007 )
1008 });
1009 if let Ok(Ok(_)) = provider_init {
1010 } else {
1012 return;
1013 }
1014
1015 let dims_tensor = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
1016 let idx_tensor = Tensor::new(vec![7.0, 8.0, 9.0], vec![1, 3]).unwrap();
1017
1018 let cpu = ind2sub_builtin(
1019 Value::Tensor(dims_tensor.clone()),
1020 Value::Tensor(idx_tensor.clone()),
1021 )
1022 .expect("cpu ind2sub");
1023
1024 let provider = runmat_accelerate_api::provider().unwrap();
1025 let view = HostTensorView {
1026 data: &idx_tensor.materialize_f64(),
1027 shape: &idx_tensor.shape,
1028 };
1029 let handle = provider.upload(&view).expect("upload indices");
1030
1031 let gpu = ind2sub_builtin(Value::Tensor(dims_tensor), Value::GpuTensor(handle))
1032 .expect("gpu ind2sub");
1033
1034 let cpu_values = match cpu {
1035 Value::Cell(cell) => cell_to_vec(&cell),
1036 other => panic!("expected cell output, got {other:?}"),
1037 };
1038 let gpu_values = match gpu {
1039 Value::Cell(cell) => cell_to_vec(&cell),
1040 other => panic!("expected cell output, got {other:?}"),
1041 };
1042
1043 assert_eq!(cpu_values.len(), gpu_values.len());
1044
1045 for (cpu_val, gpu_val) in cpu_values.iter().zip(gpu_values.iter()) {
1046 let host_cpu = test_support::gather(cpu_val.clone()).expect("gather cpu");
1047 let host_gpu = test_support::gather(gpu_val.clone()).expect("gather gpu");
1048 assert_eq!(host_cpu.shape, host_gpu.shape);
1049 assert_eq!(host_cpu.materialize_f64(), host_gpu.materialize_f64());
1050 }
1051 }
1052}