Bit hacking

TensorLib/Common.lean

@@ -78,6 +78,30 @@ def bytesToInt (order : ByteOrder) (bytes : ByteArray) : Int := Id.run do
 #guard bytesToInt .bigEndian (ByteArray.mk #[0x80, 0]) == -32768
 #guard bytesToInt .littleEndian (ByteArray.mk #[0x80, 0]) == 0x80
 
+def bitVecToByteArray (order : ByteOrder) (n : Nat) (v : BitVec n) : ByteArray := Id.run do
+  let numBytes := natDivCeil n 8
+  let mut arr := ByteArray.mkEmpty numBytes
+  match order with
+  | .oneByte =>
+    let byte := (v &&& 0xFF).toNat.toUInt8
+    return arr.push byte
+  | .littleEndian =>
+    for i in [0 : numBytes] do
+      let byte := (v.ushiftRight (i * 8) &&& 0xFF).toNat.toUInt8
+      arr := arr.push byte
+    return arr
+  | .bigEndian =>
+    for i in [0 : numBytes] do
+      let byte := (v.ushiftRight ((numBytes - i - 1) * 8) &&& 0xFF).toNat.toUInt8
+      arr := arr.push byte
+    return arr
+
+#guard (bitVecToByteArray .bigEndian 16 0x0100).toList == [1, 0]
+#guard (bitVecToByteArray .littleEndian 16 0x0100).toList == [0, 1]
+#guard (bitVecToByteArray .bigEndian 20 0x01000).toList == [0, 16, 0]
+#guard (bitVecToByteArray .littleEndian 32 0x1).toList == [1, 0, 0, 0]
+#guard (bitVecToByteArray .bigEndian 32 0x1).toList == [0, 0, 0, 1]
+
 end ByteOrder
 
 /-!
@@ -522,6 +546,8 @@ abbrev BV64 := BitVec 64
 
 def BV64.ofNat (i : Nat) : BV64 := i.toUInt64.toBitVec
 
+def BV64.ofInt (i : Int) : BV64 := i.toInt64.toBitVec
+
 def BV64.toBytes (n : BV64) : BV8 × BV8 × BV8 × BV8 × BV8 × BV8 × BV8 × BV8 :=
   let n0 := (n >>> 0o00 &&& 0xFF).truncate 8
   let n1 := (n >>> 0o10 &&& 0xFF).truncate 8

TensorLib/Dtype.lean

@@ -73,6 +73,7 @@ def isIntLike (dtype : Dtype) : Bool := dtype.isInt || dtype.isUint
 
 def int8 : Dtype := Dtype.mk Dtype.Name.int8 ByteOrder.littleEndian
 def uint8 : Dtype := Dtype.mk Dtype.Name.uint8 ByteOrder.littleEndian
+def int64 : Dtype := Dtype.mk Dtype.Name.int64 ByteOrder.littleEndian
 def uint64 : Dtype := Dtype.mk Dtype.Name.uint64 ByteOrder.littleEndian
 def float64 : Dtype := Dtype.mk Dtype.Name.float64 ByteOrder.littleEndian
 

TensorLib/Tensor.lean

@@ -349,6 +349,9 @@ instance BV16Little : Element BV16 where
   toByteArray (x : BV16) : ByteArray := x.toByteArray .littleEndian
   fromByteArray arr startIndex := ByteArray.toBV16 arr startIndex .littleEndian
 
+#guard (arange BV16 10).size == 10
+#guard toList! BV16 (Tensor.Element.arange BV16 10) == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+
 instance BV32Little : Element BV32 where
   dtype := Dtype.mk .uint32 .littleEndian
   itemsize := 4
@@ -363,8 +366,22 @@ instance BV64Little : Element BV64 where
   toByteArray (x : BV64) : ByteArray := x.toByteArray .littleEndian
   fromByteArray arr startIndex := ByteArray.toBV64 arr startIndex .littleEndian
 
-#guard (arange BV16 10).size == 10
-#guard toList! BV16 (Tensor.Element.arange BV16 10) == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+instance Int64Little : Element Int64 where
+  dtype := Dtype.mk .int64 .littleEndian
+  itemsize := 8
+  ofNat := Int64.ofNat
+  toByteArray x : ByteArray := ByteOrder.bitVecToByteArray .littleEndian 64 x.toBitVec
+  fromByteArray arr startIndex := .ok (ByteOrder.bytesToInt .littleEndian (arr.extract startIndex (startIndex + 8))).toInt64
+
+#guard (Int64Little.toByteArray 7).toList == [7, 0, 0, 0, 0, 0, 0, 0]
+#guard (Int64Little.toByteArray (-2)).toList == [0xFE, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF]
+
+private def roundTrip64 (n : Int64) := Int64Little.fromByteArray (Int64Little.toByteArray n) 0
+
+#guard roundTrip64 (-2) == .ok (-2)
+#guard roundTrip64 0 == .ok 0
+#guard roundTrip64 2 == .ok 2
+#guard roundTrip64 0xFFFFFFFFFFFFFFFF == .ok 0xFFFFFFFFFFFFFFFF
 
 end Element