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, Tensor, Type, Value,
10};
11use runmat_macros::runtime_builtin;
12
13use super::common::{build_strides, dims_from_tokens, materialize_value, parse_dims};
14use crate::builtins::array::type_resolvers::is_scalar_type;
15use crate::builtins::common::arg_tokens::tokens_from_context;
16use crate::builtins::common::spec::{
17 BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
18 ProviderHook, ReductionNaN, ResidencyPolicy, ScalarType, ShapeRequirements,
19};
20use crate::builtins::common::tensor;
21use crate::{build_runtime_error, RuntimeError};
22use runmat_builtins::shape_rules::element_count_if_known;
23
24#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::array::indexing::sub2ind")]
25pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
26 name: "sub2ind",
27 op_kind: GpuOpKind::Custom("indexing"),
28 supported_precisions: &[ScalarType::F32, ScalarType::F64],
29 broadcast: BroadcastSemantics::Matlab,
30 provider_hooks: &[ProviderHook::Custom("sub2ind")],
31 constant_strategy: ConstantStrategy::InlineLiteral,
32 residency: ResidencyPolicy::NewHandle,
33 nan_mode: ReductionNaN::Include,
34 two_pass_threshold: None,
35 workgroup_size: None,
36 accepts_nan_mode: false,
37 notes: "Providers can implement the custom `sub2ind` hook to execute on device; runtimes fall back to host computation otherwise.",
38};
39
40#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::array::indexing::sub2ind")]
41pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
42 name: "sub2ind",
43 shape: ShapeRequirements::Any,
44 constant_strategy: ConstantStrategy::InlineLiteral,
45 elementwise: None,
46 reduction: None,
47 emits_nan: false,
48 notes: "Index conversion executes eagerly on the host; fusion does not apply.",
49};
50
51fn sub2ind_type(args: &[Type], ctx: &ResolveContext) -> Type {
52 if args.len() < 2 {
53 return Type::Unknown;
54 }
55 if let Some(dims) = dims_from_tokens(&tokens_from_context(ctx)) {
56 if args.len() - 1 != dims.len() {
57 return Type::Unknown;
58 }
59 }
60 let subscripts = &args[1..];
61 if subscripts.iter().all(|ty| is_scalar_type(ty)) {
62 return Type::Num;
63 }
64 for ty in subscripts {
65 if let Type::Tensor { shape: Some(shape) } | Type::Logical { shape: Some(shape) } = ty {
66 if element_count_if_known(shape).unwrap_or(0) > 1 {
67 return Type::Tensor {
68 shape: Some(shape.clone()),
69 };
70 }
71 }
72 }
73 Type::tensor()
74}
75
76const BUILTIN_NAME: &str = "sub2ind";
77
78const SUB2IND_OUTPUT: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
79 name: "ind",
80 ty: BuiltinParamType::NumericArray,
81 arity: BuiltinParamArity::Required,
82 default: None,
83 description: "Column-major linear indices corresponding to provided subscripts.",
84}];
85
86const SUB2IND_INPUTS: [BuiltinParamDescriptor; 3] = [
87 BuiltinParamDescriptor {
88 name: "sz",
89 ty: BuiltinParamType::SizeArg,
90 arity: BuiltinParamArity::Required,
91 default: None,
92 description: "Size vector describing source array dimensions.",
93 },
94 BuiltinParamDescriptor {
95 name: "I1",
96 ty: BuiltinParamType::Any,
97 arity: BuiltinParamArity::Required,
98 default: None,
99 description: "First-dimension subscript values.",
100 },
101 BuiltinParamDescriptor {
102 name: "In",
103 ty: BuiltinParamType::Any,
104 arity: BuiltinParamArity::Variadic,
105 default: None,
106 description: "Remaining per-dimension subscript arrays/scalars.",
107 },
108];
109
110const SUB2IND_SIGNATURES: [BuiltinSignatureDescriptor; 1] = [BuiltinSignatureDescriptor {
111 label: "ind = sub2ind(sz, I1, In...)",
112 inputs: &SUB2IND_INPUTS,
113 outputs: &SUB2IND_OUTPUT,
114}];
115
116const SUB2IND_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
117 code: "RM.SUB2IND.INVALID_INPUT",
118 identifier: Some("RunMat:sub2ind:InvalidInput"),
119 when: "Size vector, subscript count, or subscript types are invalid.",
120 message: "sub2ind: invalid input arguments",
121};
122
123const SUB2IND_ERROR_INDEX_BOUNDS: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
124 code: "RM.SUB2IND.INDEX_BOUNDS",
125 identifier: Some("RunMat:sub2ind:IndexBounds"),
126 when: "At least one subscript lies outside bounds for its dimension.",
127 message: "sub2ind: subscript index exceeds dimension bounds",
128};
129
130const SUB2IND_ERROR_PROVIDER: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
131 code: "RM.SUB2IND.PROVIDER",
132 identifier: Some("RunMat:sub2ind:ProviderError"),
133 when: "GPU provider sub2ind hook fails.",
134 message: "sub2ind: provider execution failed",
135};
136
137const SUB2IND_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
138 code: "RM.SUB2IND.INTERNAL",
139 identifier: Some("RunMat:sub2ind:InternalError"),
140 when: "Internal tensor conversion/output construction fails.",
141 message: "sub2ind: internal error",
142};
143
144const SUB2IND_ERRORS: [BuiltinErrorDescriptor; 4] = [
145 SUB2IND_ERROR_INVALID_INPUT,
146 SUB2IND_ERROR_INDEX_BOUNDS,
147 SUB2IND_ERROR_PROVIDER,
148 SUB2IND_ERROR_INTERNAL,
149];
150
151pub const SUB2IND_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
152 signatures: &SUB2IND_SIGNATURES,
153 output_mode: BuiltinOutputMode::Fixed,
154 completion_policy: BuiltinCompletionPolicy::Public,
155 errors: &SUB2IND_ERRORS,
156};
157
158fn sub2ind_error_with_message(
159 message: impl Into<String>,
160 error: &'static BuiltinErrorDescriptor,
161) -> RuntimeError {
162 let mut builder = build_runtime_error(message).with_builtin(BUILTIN_NAME);
163 if let Some(identifier) = error.identifier {
164 builder = builder.with_identifier(identifier);
165 }
166 builder.build()
167}
168
169fn sub2ind_input_error(message: impl Into<String>) -> RuntimeError {
170 sub2ind_error_with_message(message, &SUB2IND_ERROR_INVALID_INPUT)
171}
172
173fn sub2ind_bounds_error(message: impl Into<String>) -> RuntimeError {
174 sub2ind_error_with_message(message, &SUB2IND_ERROR_INDEX_BOUNDS)
175}
176
177fn sub2ind_provider_error(message: impl Into<String>) -> RuntimeError {
178 sub2ind_error_with_message(message, &SUB2IND_ERROR_PROVIDER)
179}
180
181fn sub2ind_internal_error(message: impl Into<String>) -> RuntimeError {
182 sub2ind_error_with_message(message, &SUB2IND_ERROR_INTERNAL)
183}
184
185#[runtime_builtin(
186 name = "sub2ind",
187 category = "array/indexing",
188 summary = "Convert N-D subscripts to MATLAB-style column-major linear indices.",
189 keywords = "sub2ind,linear index,column major,gpu indexing",
190 accel = "custom",
191 type_resolver(sub2ind_type),
192 descriptor(crate::builtins::array::indexing::sub2ind::SUB2IND_DESCRIPTOR),
193 builtin_path = "crate::builtins::array::indexing::sub2ind"
194)]
195async fn sub2ind_builtin(dims_val: Value, rest: Vec<Value>) -> crate::BuiltinResult<Value> {
196 let (dims_value, dims_was_gpu) = materialize_value(dims_val, "sub2ind").await?;
197 let dims = parse_dims(&dims_value, "sub2ind").await?;
198 if dims.is_empty() {
199 return Err(sub2ind_error("Size vector must have at least one element."));
200 }
201
202 if rest.len() != dims.len() {
203 return Err(sub2ind_error(
204 "The number of subscripts supplied must equal the number of dimensions in the size vector.",
205 ));
206 }
207
208 if let Some(value) = try_gpu_sub2ind(&dims, &rest)? {
209 return Ok(value);
210 }
211
212 let mut saw_gpu = dims_was_gpu;
213 let mut subscripts: Vec<Tensor> = Vec::with_capacity(rest.len());
214 for value in rest {
215 let (materialised, was_gpu) = materialize_value(value, "sub2ind").await?;
216 saw_gpu |= was_gpu;
217 let tensor = tensor::value_into_tensor_for("sub2ind", materialised)
218 .map_err(|message| sub2ind_error(message))?;
219 subscripts.push(tensor);
220 }
221
222 let (result_data, result_shape) = compute_indices(&dims, &subscripts)?;
223 let want_gpu_output = saw_gpu && runmat_accelerate_api::provider().is_some();
224
225 if want_gpu_output {
226 #[cfg(all(test, feature = "wgpu"))]
227 {
228 if runmat_accelerate_api::provider().is_none() {
229 let _ = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
230 runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
231 );
232 }
233 }
234 let shape = result_shape.clone().unwrap_or_else(|| vec![1, 1]);
235 if let Some(provider) = runmat_accelerate_api::provider() {
236 let view = HostTensorView {
237 data: &result_data,
238 shape: &shape,
239 };
240 if let Ok(handle) = provider.upload(&view) {
241 return Ok(Value::GpuTensor(handle));
242 }
243 }
244 }
245
246 build_host_value(result_data, result_shape)
247}
248
249fn try_gpu_sub2ind(dims: &[usize], subs: &[Value]) -> crate::BuiltinResult<Option<Value>> {
250 #[cfg(target_arch = "wasm32")]
251 {
252 let _ = (dims, subs);
253 Ok(None)
254 }
255 #[cfg(not(target_arch = "wasm32"))]
256 {
257 let provider = match runmat_accelerate_api::provider() {
258 Some(p) => p,
259 None => return Ok(None),
260 };
261 if !subs
262 .iter()
263 .all(|value| matches!(value, Value::GpuTensor(_)))
264 {
265 return Ok(None);
266 }
267 if dims.is_empty() {
268 return Ok(None);
269 }
270
271 let mut handles: Vec<&GpuTensorHandle> = Vec::with_capacity(subs.len());
272 for value in subs {
273 if let Value::GpuTensor(handle) = value {
274 handles.push(handle);
275 }
276 }
277
278 if handles.len() != dims.len() {
279 return Err(sub2ind_error(
280 "The number of subscripts supplied must equal the number of dimensions in the size vector.",
281 ));
282 }
283
284 let mut scalar_mask: Vec<bool> = Vec::with_capacity(handles.len());
285 let mut target_shape: Option<Vec<usize>> = None;
286 let mut result_len: usize = 1;
287 let mut saw_non_scalar = false;
288
289 for handle in &handles {
290 let len = tensor::element_count(&handle.shape);
291 let is_scalar = len == 1;
292 scalar_mask.push(is_scalar);
293 if !is_scalar {
294 saw_non_scalar = true;
295 if let Some(existing) = &target_shape {
296 if existing != &handle.shape {
297 return Err(sub2ind_error("Subscript inputs must have the same size."));
298 }
299 } else {
300 target_shape = Some(handle.shape.clone());
301 result_len = len;
302 }
303 }
304 }
305
306 if !saw_non_scalar {
307 target_shape = Some(vec![1, 1]);
308 result_len = 1;
309 } else if let Some(shape) = &target_shape {
310 result_len = tensor::element_count(shape);
311 }
312
313 let strides = build_strides(dims, "sub2ind")?;
314 if dims.iter().any(|&d| d > u32::MAX as usize)
315 || strides.iter().any(|&s| s > u32::MAX as usize)
316 || result_len > u32::MAX as usize
317 {
318 return Ok(None);
319 }
320
321 let output_shape = target_shape.clone().unwrap_or_else(|| vec![1, 1]);
322 match provider.sub2ind(
323 dims,
324 &strides,
325 &handles,
326 &scalar_mask,
327 result_len,
328 &output_shape,
329 ) {
330 Ok(handle) => Ok(Some(Value::GpuTensor(handle))),
331 Err(err) => Err(sub2ind_provider_error(err.to_string())),
332 }
333 }
334}
335
336fn compute_indices(
337 dims: &[usize],
338 subscripts: &[Tensor],
339) -> crate::BuiltinResult<(Vec<f64>, Option<Vec<usize>>)> {
340 let mut target_shape: Option<Vec<usize>> = None;
341 let mut result_len: usize = 1;
342 let mut has_non_scalar = false;
343
344 for tensor in subscripts {
345 if tensor.data.len() != 1 {
346 has_non_scalar = true;
347 if let Some(shape) = &target_shape {
348 if &tensor.shape != shape {
349 return Err(sub2ind_error("Subscript inputs must have the same size."));
350 }
351 } else {
352 target_shape = Some(tensor.shape.clone());
353 result_len = tensor.data.len();
354 }
355 }
356 }
357
358 if !has_non_scalar {
359 target_shape = Some(vec![1, 1]);
361 result_len = 1;
362 }
363
364 if result_len == 0 {
365 return Ok((Vec::new(), target_shape));
366 }
367
368 let strides = build_strides(dims, "sub2ind")?;
369 let mut output = Vec::with_capacity(result_len);
370
371 for idx in 0..result_len {
372 let mut offset: usize = 0;
373 for (dim_index, (&dim, tensor)) in dims.iter().zip(subscripts.iter()).enumerate() {
374 let raw = subscript_value(tensor, idx);
375 let coerced = coerce_subscript(raw, dim_index + 1, dim)?;
376 let term = coerced
377 .checked_sub(1)
378 .and_then(|v| v.checked_mul(strides[dim_index]))
379 .ok_or_else(|| sub2ind_bounds_error("Index exceeds array dimensions."))?;
380 offset = offset
381 .checked_add(term)
382 .ok_or_else(|| sub2ind_bounds_error("Index exceeds array dimensions."))?;
383 }
384 output.push((offset + 1) as f64);
385 }
386
387 Ok((output, target_shape))
388}
389
390fn subscript_value(tensor: &Tensor, idx: usize) -> f64 {
391 if tensor.data.len() == 1 {
392 tensor.data[0]
393 } else {
394 tensor.data[idx]
395 }
396}
397
398fn coerce_subscript(value: f64, dim_number: usize, dim_size: usize) -> crate::BuiltinResult<usize> {
399 if !value.is_finite() {
400 return Err(sub2ind_error(
401 "Subscript indices must either be real positive integers or logicals.",
402 ));
403 }
404 let rounded = value.round();
405 if (rounded - value).abs() > f64::EPSILON {
406 return Err(sub2ind_error(
407 "Subscript indices must either be real positive integers or logicals.",
408 ));
409 }
410 if rounded < 1.0 {
411 return Err(sub2ind_error(
412 "Subscript indices must either be real positive integers or logicals.",
413 ));
414 }
415 if rounded > dim_size as f64 {
416 return Err(dimension_bounds_error(dim_number));
417 }
418 Ok(rounded as usize)
419}
420
421fn dimension_bounds_error(dim_number: usize) -> RuntimeError {
422 let message = match dim_number {
423 1 => format!("Index exceeds the number of rows in dimension {dim_number}."),
424 2 => format!("Index exceeds the number of columns in dimension {dim_number}."),
425 3 => format!("Index exceeds the number of pages in dimension {dim_number}."),
426 _ => "Index exceeds array dimensions.".to_string(),
427 };
428 sub2ind_bounds_error(message)
429}
430
431fn build_host_value(data: Vec<f64>, shape: Option<Vec<usize>>) -> crate::BuiltinResult<Value> {
432 let shape = shape.unwrap_or_else(|| vec![1, 1]);
433 if data.len() == 1 && tensor::element_count(&shape) == 1 {
434 Ok(Value::Num(data[0]))
435 } else {
436 let tensor = Tensor::new(data, shape).map_err(|e| {
437 sub2ind_internal_error(format!("Unable to construct sub2ind output: {e}"))
438 })?;
439 Ok(Value::Tensor(tensor))
440 }
441}
442
443fn sub2ind_error(message: impl Into<String>) -> RuntimeError {
444 sub2ind_input_error(message)
445}
446
447#[cfg(test)]
448pub(crate) mod tests {
449 use super::*;
450 use crate::builtins::common::test_support;
451 use futures::executor::block_on;
452 use runmat_builtins::{IntValue, Tensor, Type, Value};
453
454 fn sub2ind_builtin(dims_val: Value, rest: Vec<Value>) -> crate::BuiltinResult<Value> {
455 block_on(super::sub2ind_builtin(dims_val, rest))
456 }
457
458 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
459 #[test]
460 fn converts_scalar_indices() {
461 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
462 let result =
463 sub2ind_builtin(Value::Tensor(dims), vec![Value::Num(2.0), Value::Num(3.0)]).unwrap();
464 assert_eq!(result, Value::Num(8.0));
465 }
466
467 #[test]
468 fn sub2ind_type_scalar_outputs_num() {
469 assert_eq!(
470 sub2ind_type(
471 &[Type::Tensor { shape: None }, Type::Num, Type::Int],
472 &ResolveContext::new(Vec::new()),
473 ),
474 Type::Num
475 );
476 }
477
478 #[test]
479 fn sub2ind_type_vector_outputs_tensor() {
480 let subs = Type::Tensor {
481 shape: Some(vec![Some(3), Some(1)]),
482 };
483 assert_eq!(
484 sub2ind_type(
485 &[Type::Tensor { shape: None }, subs.clone(), Type::Num],
486 &ResolveContext::new(Vec::new()),
487 ),
488 Type::Tensor {
489 shape: Some(vec![Some(3), Some(1)])
490 }
491 );
492 }
493
494 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
495 #[test]
496 fn broadcasts_scalars_over_vectors() {
497 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
498 let rows = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
499 let result = sub2ind_builtin(
500 Value::Tensor(dims),
501 vec![Value::Tensor(rows), Value::Num(4.0)],
502 )
503 .unwrap();
504 match result {
505 Value::Tensor(t) => {
506 assert_eq!(t.shape, vec![3, 1]);
507 assert_eq!(t.data, vec![10.0, 11.0, 12.0]);
508 }
509 other => panic!("expected tensor result, got {other:?}"),
510 }
511 }
512
513 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
514 #[test]
515 fn handles_three_dimensions() {
516 let dims = Tensor::new(vec![2.0, 3.0, 4.0], vec![1, 3]).unwrap();
517 let row = Tensor::new(vec![1.0, 1.0], vec![1, 2]).unwrap();
518 let col = Tensor::new(vec![2.0, 3.0], vec![1, 2]).unwrap();
519 let page = Tensor::new(vec![1.0, 2.0], vec![1, 2]).unwrap();
520 let result = sub2ind_builtin(
521 Value::Tensor(dims),
522 vec![Value::Tensor(row), Value::Tensor(col), Value::Tensor(page)],
523 )
524 .unwrap();
525 match result {
526 Value::Tensor(t) => {
527 assert_eq!(t.shape, vec![1, 2]);
528 assert_eq!(t.data, vec![3.0, 11.0]);
529 }
530 other => panic!("expected tensor result, got {other:?}"),
531 }
532 }
533
534 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
535 #[test]
536 fn rejects_out_of_range_subscripts() {
537 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
538 let err = sub2ind_builtin(Value::Tensor(dims), vec![Value::Num(4.0), Value::Num(1.0)])
539 .unwrap_err();
540 assert!(
541 err.to_string().contains("Index exceeds"),
542 "expected index bounds error, got {err}"
543 );
544 assert_eq!(
545 err.identifier(),
546 super::SUB2IND_ERROR_INDEX_BOUNDS.identifier
547 );
548 }
549
550 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
551 #[test]
552 fn rejects_shape_mismatch() {
553 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
554 let rows = Tensor::new(vec![1.0, 2.0], vec![2, 1]).unwrap();
555 let cols = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
556 let err = sub2ind_builtin(
557 Value::Tensor(dims),
558 vec![Value::Tensor(rows), Value::Tensor(cols)],
559 )
560 .unwrap_err();
561 assert!(
562 err.to_string().contains("same size"),
563 "expected size mismatch error, got {err}"
564 );
565 assert_eq!(
566 err.identifier(),
567 super::SUB2IND_ERROR_INVALID_INPUT.identifier
568 );
569 }
570
571 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
572 #[test]
573 fn rejects_non_integer_subscripts() {
574 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
575 let err = sub2ind_builtin(Value::Tensor(dims), vec![Value::Num(1.5), Value::Num(1.0)])
576 .unwrap_err();
577 assert!(
578 err.to_string().contains("real positive integers"),
579 "expected integer coercion error, got {err}"
580 );
581 assert_eq!(
582 err.identifier(),
583 super::SUB2IND_ERROR_INVALID_INPUT.identifier
584 );
585 }
586
587 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
588 #[test]
589 fn accepts_integer_value_variants() {
590 let dims = Value::Tensor(Tensor::new(vec![3.0], vec![1, 1]).unwrap());
591 let result = sub2ind_builtin(dims, vec![Value::Int(IntValue::I32(2))]).expect("sub2ind");
592 assert_eq!(result, Value::Num(2.0));
593 }
594
595 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
596 #[test]
597 fn sub2ind_gpu_roundtrip() {
598 test_support::with_test_provider(|provider| {
599 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
600 let rows = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
601 let cols = Tensor::new(vec![4.0, 4.0, 4.0], vec![3, 1]).unwrap();
602
603 let dims_handle = provider
604 .upload(&HostTensorView {
605 data: &dims.data,
606 shape: &dims.shape,
607 })
608 .expect("upload dims");
609 let rows_handle = provider
610 .upload(&HostTensorView {
611 data: &rows.data,
612 shape: &rows.shape,
613 })
614 .expect("upload rows");
615 let cols_handle = provider
616 .upload(&HostTensorView {
617 data: &cols.data,
618 shape: &cols.shape,
619 })
620 .expect("upload cols");
621
622 let result = sub2ind_builtin(
623 Value::GpuTensor(dims_handle),
624 vec![Value::GpuTensor(rows_handle), Value::GpuTensor(cols_handle)],
625 )
626 .expect("sub2ind");
627
628 match result {
629 Value::GpuTensor(handle) => {
630 let gathered = test_support::gather(Value::GpuTensor(handle)).unwrap();
631 assert_eq!(gathered.shape, vec![3, 1]);
632 assert_eq!(gathered.data, vec![10.0, 11.0, 12.0]);
633 }
634 other => panic!("expected gpu tensor, got {other:?}"),
635 }
636 });
637 }
638
639 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
640 #[test]
641 #[cfg(feature = "wgpu")]
642 fn sub2ind_wgpu_matches_cpu() {
643 let _ = runmat_accelerate::backend::wgpu::provider::register_wgpu_provider(
644 runmat_accelerate::backend::wgpu::provider::WgpuProviderOptions::default(),
645 );
646 let Some(provider) = runmat_accelerate_api::provider() else {
647 panic!("wgpu provider not available");
648 };
649
650 let dims = Tensor::new(vec![3.0, 4.0], vec![1, 2]).unwrap();
651 let rows = Tensor::new(vec![1.0, 2.0, 3.0], vec![3, 1]).unwrap();
652 let cols = Tensor::new(vec![4.0, 4.0, 4.0], vec![3, 1]).unwrap();
653
654 let cpu = sub2ind_builtin(
655 Value::Tensor(dims.clone()),
656 vec![Value::Tensor(rows.clone()), Value::Tensor(cols.clone())],
657 )
658 .expect("cpu sub2ind");
659
660 let rows_handle = provider
661 .upload(&HostTensorView {
662 data: &rows.data,
663 shape: &rows.shape,
664 })
665 .expect("upload rows");
666 let cols_handle = provider
667 .upload(&HostTensorView {
668 data: &cols.data,
669 shape: &cols.shape,
670 })
671 .expect("upload cols");
672
673 let result = sub2ind_builtin(
674 Value::Tensor(dims),
675 vec![Value::GpuTensor(rows_handle), Value::GpuTensor(cols_handle)],
676 )
677 .expect("wgpu sub2ind");
678
679 let gathered = test_support::gather(result).expect("gather");
680 let expected = match cpu {
681 Value::Tensor(t) => t,
682 Value::Num(v) => Tensor::new(vec![v], vec![1, 1]).unwrap(),
683 other => panic!("unexpected cpu result {other:?}"),
684 };
685 assert_eq!(gathered.shape, expected.shape);
686 assert_eq!(gathered.data, expected.data);
687 }
688}