Merge ml/tensor-ops-d-bmm-int: batched matmul + integer/bit ops (ops D)

# Conflicts:
#	crates/rsscript/src/reg_vm/mod.rs
#	crates/rsscript/src/runtime_abi.rs
#	crates/rsscript/tests/vm_eval.rs
#	crates/rsscript/tests/vm_eval_parity/tensor.rs
#	crates/runtime/src/tensor.rs
#	packages/ml/interface/tensor.rssi

Author: olwang

crates/rsscript/src/reg_vm/mod.rs

@@ -2126,6 +2126,17 @@ enum RegIntrinsic {
     TensorLog2,
     TensorRsqrt,
     TensorPow,
+    // bmm+int/bit (ops D)
+    TensorBmm,
+    TensorIdiv,
+    TensorMod,
+    TensorShl,
+    TensorShr,
+    TensorAnd,
+    TensorOr,
+    TensorXor,
+    TensorBitcastF32ToI32,
+    TensorBitcastI32ToF32,
     TensorErrorMessage,
     CharCompare,
     CharFromCode,
@@ -5411,6 +5422,17 @@ fn qualified_intrinsic(namespace: &str, name: &str) -> Option<RegIntrinsic> {
         ("Tensor", "log2") => Some(RegIntrinsic::TensorLog2),
         ("Tensor", "rsqrt") => Some(RegIntrinsic::TensorRsqrt),
         ("Tensor", "pow") => Some(RegIntrinsic::TensorPow),
+        // bmm+int/bit (ops D)
+        ("Tensor", "bmm") => Some(RegIntrinsic::TensorBmm),
+        ("Tensor", "idiv") => Some(RegIntrinsic::TensorIdiv),
+        ("Tensor", "modulo") => Some(RegIntrinsic::TensorMod),
+        ("Tensor", "shl") => Some(RegIntrinsic::TensorShl),
+        ("Tensor", "shr") => Some(RegIntrinsic::TensorShr),
+        ("Tensor", "bit_and") => Some(RegIntrinsic::TensorAnd),
+        ("Tensor", "bit_or") => Some(RegIntrinsic::TensorOr),
+        ("Tensor", "bit_xor") => Some(RegIntrinsic::TensorXor),
+        ("Tensor", "bitcast_f32_to_i32") => Some(RegIntrinsic::TensorBitcastF32ToI32),
+        ("Tensor", "bitcast_i32_to_f32") => Some(RegIntrinsic::TensorBitcastI32ToF32),
         ("TensorError", "message") => Some(RegIntrinsic::TensorErrorMessage),
         ("Char", "compare") => Some(RegIntrinsic::CharCompare),
         ("Char", "from_code") => Some(RegIntrinsic::CharFromCode),
@@ -9594,6 +9616,44 @@ impl RegVm {
                     )),
                 }))
             }
+            // bmm+int/bit (ops D)
+            RegIntrinsic::TensorBmm
+            | RegIntrinsic::TensorIdiv
+            | RegIntrinsic::TensorMod
+            | RegIntrinsic::TensorShl
+            | RegIntrinsic::TensorShr
+            | RegIntrinsic::TensorAnd
+            | RegIntrinsic::TensorOr
+            | RegIntrinsic::TensorXor => {
+                let a = self.expect_tensor_ref(intrinsic_arg(&self.stack, base, args, 0)?)?;
+                let b = self.expect_tensor_ref(intrinsic_arg(&self.stack, base, args, 1)?)?;
+                let result = match intrinsic {
+                    RegIntrinsic::TensorBmm => rsscript_runtime::tensor_bmm(&a, &b),
+                    RegIntrinsic::TensorIdiv => rsscript_runtime::tensor_idiv(&a, &b),
+                    RegIntrinsic::TensorMod => rsscript_runtime::tensor_mod(&a, &b),
+                    RegIntrinsic::TensorShl => rsscript_runtime::tensor_shl(&a, &b),
+                    RegIntrinsic::TensorShr => rsscript_runtime::tensor_shr(&a, &b),
+                    RegIntrinsic::TensorAnd => rsscript_runtime::tensor_and(&a, &b),
+                    RegIntrinsic::TensorOr => rsscript_runtime::tensor_or(&a, &b),
+                    _ => rsscript_runtime::tensor_xor(&a, &b),
+                };
+                Ok(json_result(match result {
+                    Ok(tensor) => Ok(self.store_tensor(tensor)),
+                    Err(error) => Err(tensor_error_value(
+                        rsscript_runtime::tensor_error_message(&error),
+                    )),
+                }))
+            }
+            RegIntrinsic::TensorBitcastF32ToI32 | RegIntrinsic::TensorBitcastI32ToF32 => {
+                let t = self.expect_tensor_ref(intrinsic_arg(&self.stack, base, args, 0)?)?;
+                let result = match intrinsic {
+                    RegIntrinsic::TensorBitcastF32ToI32 => {
+                        rsscript_runtime::tensor_bitcast_f32_to_i32(&t)
+                    }
+                    _ => rsscript_runtime::tensor_bitcast_i32_to_f32(&t),
+                };
+                Ok(self.store_tensor(result))
+            }
             RegIntrinsic::CharCompare | RegIntrinsic::CharFromCode | RegIntrinsic::CharIsAlphanumeric | RegIntrinsic::CharIsAlpha | RegIntrinsic::CharIsDigit | RegIntrinsic::CharIsLower | RegIntrinsic::CharIsUpper | RegIntrinsic::CharIsWhitespace | RegIntrinsic::CharToCode | RegIntrinsic::CharToLower | RegIntrinsic::CharToString | RegIntrinsic::CharToUpper => self.exec_char_intrinsics(unit, intrinsic, args, base, next_base),
             RegIntrinsic::ClockNow => Ok(instant_value(clock_system_unix_ms())),
             RegIntrinsic::ClockSystemUnixMs => Ok(VmValue::Int(clock_system_unix_ms())),

crates/rsscript/src/runtime_abi.rs

@@ -494,6 +494,25 @@ const RUNTIME_INTRINSICS: &[RuntimeIntrinsic] = &[
     runtime_intrinsic("Tensor", "log2", "rsscript_runtime::tensor_log2"),
     runtime_intrinsic("Tensor", "rsqrt", "rsscript_runtime::tensor_rsqrt"),
     runtime_intrinsic("Tensor", "pow", "rsscript_runtime::tensor_pow"),
+    // bmm+int/bit (ops D)
+    runtime_intrinsic("Tensor", "bmm", "rsscript_runtime::tensor_bmm"),
+    runtime_intrinsic("Tensor", "idiv", "rsscript_runtime::tensor_idiv"),
+    runtime_intrinsic("Tensor", "modulo", "rsscript_runtime::tensor_mod"),
+    runtime_intrinsic("Tensor", "shl", "rsscript_runtime::tensor_shl"),
+    runtime_intrinsic("Tensor", "shr", "rsscript_runtime::tensor_shr"),
+    runtime_intrinsic("Tensor", "bit_and", "rsscript_runtime::tensor_and"),
+    runtime_intrinsic("Tensor", "bit_or", "rsscript_runtime::tensor_or"),
+    runtime_intrinsic("Tensor", "bit_xor", "rsscript_runtime::tensor_xor"),
+    runtime_intrinsic(
+        "Tensor",
+        "bitcast_f32_to_i32",
+        "rsscript_runtime::tensor_bitcast_f32_to_i32",
+    ),
+    runtime_intrinsic(
+        "Tensor",
+        "bitcast_i32_to_f32",
+        "rsscript_runtime::tensor_bitcast_i32_to_f32",
+    ),
     runtime_intrinsic(
         "TensorError",
         "message",

crates/rsscript/tests/vm_eval.rs

@@ -1075,6 +1075,7 @@ pub fn echo(message: &String) -> String {
 // parity: runtime:Tensor.cast_f32 runtime:Tensor.cast_i32 runtime:Tensor.cast_bool runtime:Tensor.dtype_code
 // parity: runtime:Tensor.pad runtime:Tensor.shrink runtime:Tensor.flip runtime:Tensor.gather
 // parity: runtime:Tensor.prod_axis runtime:Tensor.min_axis runtime:Tensor.sum_axes runtime:Tensor.prod_axes runtime:Tensor.max_axes runtime:Tensor.min_axes runtime:Tensor.mean_axes runtime:Tensor.reciprocal runtime:Tensor.exp2 runtime:Tensor.log2 runtime:Tensor.rsqrt runtime:Tensor.pow
+// parity: runtime:Tensor.bmm runtime:Tensor.idiv runtime:Tensor.modulo runtime:Tensor.shl runtime:Tensor.shr runtime:Tensor.bit_and runtime:Tensor.bit_or runtime:Tensor.bit_xor runtime:Tensor.bitcast_f32_to_i32 runtime:Tensor.bitcast_i32_to_f32
 // parity: runtime:Timer.sleep runtime:Timer.sleep_cancellable runtime:Timer.sleep_until
 // parity: runtime:Tcp.connect runtime:TcpError.message
 // parity: runtime:TcpStream.read runtime:TcpStream.shutdown

crates/rsscript/tests/vm_eval_parity/tensor.rs

@@ -1285,3 +1285,96 @@ fn main() -> Result<Unit, TensorError> {
         source,
     );
 }
+
+// bmm+int/bit (ops D)
+
+#[test]
+fn parity_tensor_bmm() {
+    let source = r#"
+features: native, local
+
+fn dump(t: read Tensor) -> Result<Unit, TensorError> {
+    let shape = Tensor.shape(tensor: read t)?
+    let mut si = 0
+    while si < List.len(list: read shape) {
+        Log.write(message: read String.from_int(value: List.get(list: read shape, index: si)))
+        si = si + 1
+    }
+    let values = Tensor.to_f32_slice(tensor: read t)?
+    let mut index = 0
+    while index < List.len(list: read values) {
+        Log.write(message: read Float.to_string(value: read List.get(list: read values, index: index)))
+        index = index + 1
+    }
+    return Ok(Unit)
+}
+
+fn main() -> Result<Unit, TensorError> {
+    let a = Tensor.from_f32_slice(data: read [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], shape: read [2, 2, 2])?
+    let b = Tensor.from_f32_slice(data: read [1.0, 0.0, 0.0, 1.0, 2.0, 0.0, 0.0, 2.0], shape: read [2, 2, 2])?
+    let c = Tensor.bmm(a: read a, b: read b)?
+    dump(t: read c)?
+    // Broadcast a single matrix over the batch dim.
+    let big = Tensor.from_f32_slice(data: read [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0], shape: read [2, 2, 3])?
+    let w = Tensor.from_f32_slice(data: read [1.0, 0.0, 0.0, 1.0, 1.0, 1.0], shape: read [1, 3, 2])?
+    let bc = Tensor.bmm(a: read big, b: read w)?
+    dump(t: read bc)?
+    return Ok(Unit)
+}
+"#;
+    common::assert_vm_eval_matches_backend(
+        "parity-tensor-bmm.rss",
+        "rsscript_parity_tensor_bmm",
+        source,
+    );
+}
+
+#[test]
+fn parity_tensor_int_bit_ops() {
+    let source = r#"
+features: native, local
+
+fn dump(t: read Tensor) -> Result<Unit, TensorError> {
+    Log.write(message: read String.from_int(value: Tensor.dtype_code(t: read t)))
+    let values = Tensor.to_f32_slice(tensor: read t)?
+    let mut index = 0
+    while index < List.len(list: read values) {
+        Log.write(message: read Float.to_string(value: read List.get(list: read values, index: index)))
+        index = index + 1
+    }
+    return Ok(Unit)
+}
+
+fn main() -> Result<Unit, TensorError> {
+    let af = Tensor.from_f32_slice(data: read [7.0, -7.0, 12.0, 6.0], shape: read [4])?
+    let bf = Tensor.from_f32_slice(data: read [2.0, 2.0, 3.0, 0.0], shape: read [4])?
+    let a = Tensor.cast_i32(t: read af)
+    let b = Tensor.cast_i32(t: read bf)
+    let q = Tensor.idiv(a: read a, b: read b)?
+    dump(t: read q)?
+    let r = Tensor.modulo(a: read a, b: read b)?
+    dump(t: read r)?
+    let shf = Tensor.from_f32_slice(data: read [1.0, 2.0, 1.0, 2.0], shape: read [4])?
+    let sh = Tensor.cast_i32(t: read shf)
+    let l = Tensor.shl(a: read a, b: read sh)?
+    dump(t: read l)?
+    let rs = Tensor.shr(a: read a, b: read sh)?
+    dump(t: read rs)?
+    let aa = Tensor.bit_and(a: read a, b: read b)?
+    dump(t: read aa)?
+    let oo = Tensor.bit_or(a: read a, b: read b)?
+    dump(t: read oo)?
+    let xx = Tensor.bit_xor(a: read a, b: read b)?
+    dump(t: read xx)?
+    let bits = Tensor.bitcast_f32_to_i32(t: read af)
+    let back = Tensor.bitcast_i32_to_f32(t: read bits)
+    dump(t: read back)?
+    return Ok(Unit)
+}
+"#;
+    common::assert_vm_eval_matches_backend(
+        "parity-tensor-int-bit-ops.rss",
+        "rsscript_parity_tensor_int_bit_ops",
+        source,
+    );
+}