add internal utilities for debugging XLA objects from Idris

Author: joelberkeley

backend/src/ffi.cpp

@@ -13,7 +13,13 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */
+#include <iostream>
+
 extern "C" {
+    void print_address(void* ptr) {
+        std::cout << ptr << std::endl;
+    }
+
     int sizeof_int() {
         return sizeof(int);
     }

backend/src/ffi.h

@@ -14,6 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 extern "C" {
+    void print_address(void* ptr);  // doesn't belong here?
+
     int sizeof_int();
 
     void set_array_int(int* arr, int idx, int value);

backend/src/tensorflow/compiler/xla/client/xla_builder.cpp

@@ -76,6 +76,13 @@ extern "C" {
         return reinterpret_cast<XlaComputation*>(non_stack);
     }
 
+    Shape* XlaBuilder_GetShape(XlaBuilder& s, XlaOp& op) {
+        auto& s_ = reinterpret_cast<xla::XlaBuilder&>(s);
+        auto& op_ = reinterpret_cast<xla::XlaOp&>(op);
+        xla::Shape shape = s_.GetShape(op_).ConsumeValueOrDie();
+        return reinterpret_cast<Shape*>(new xla::Shape(shape));
+    }
+
     const char* XlaBuilder_OpToString(XlaBuilder& s, XlaOp& op) {
         auto& s_ = reinterpret_cast<xla::XlaBuilder&>(s);
         auto& op_ = reinterpret_cast<xla::XlaOp&>(op);

backend/src/tensorflow/compiler/xla/client/xla_builder.h

@@ -48,6 +48,7 @@ extern "C" {
     const char* XlaBuilder_name(XlaBuilder& s);
     XlaBuilder* CreateSubBuilder(XlaBuilder& s, const char* computation_name);
     XlaComputation* XlaBuilder_Build(XlaBuilder& s, XlaOp& root);
+    Shape* XlaBuilder_GetShape(XlaBuilder& s, XlaOp& op);
     const char* XlaBuilder_OpToString(XlaBuilder& s, XlaOp& op);
 
     /*

backend/src/tensorflow/compiler/xla/shape.cpp

@@ -13,12 +13,30 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */
+#include <algorithm>
+#include <cstring>
+#include <string>
+
 #include "tensorflow/compiler/xla/shape.h"
 
 #include "shape.h"
 
 extern "C" {
+    const char* c_string_copy(std::string str) {
+        char *res = NULL;
+        auto len = str.length();
+        res = (char *) malloc(len + 1);
+        strncpy(res, str.c_str(), len);
+        res[len] = '\0';
+        return res;
+    }
+
     void Shape_delete(Shape* s) {
         delete reinterpret_cast<xla::Shape*>(s);
     }
+
+    const char* Shape_DebugString(Shape& s) {
+        auto& s_ = reinterpret_cast<xla::Shape&>(s);
+        return c_string_copy(s_.DebugString());
+    }
 }

spidr.ipkg

@@ -9,6 +9,7 @@ modules =
   BayesianOptimization.Acquisition,
   BayesianOptimization.Morphisms,
 
+  Compiler.Debug,
   Compiler.Eval,
   Compiler.Expr,
   Compiler.LiteralRW,

src/Compiler/Debug.idr

@@ -0,0 +1,23 @@
+{--
+Copyright 2022 Joel Berkeley
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+--}
+module Compiler.Debug
+
+import Compiler.Xla.TensorFlow.Compiler.Xla.Shape
+import Compiler.Xla.TensorFlow.Compiler.Xla.Client.XlaBuilder
+
+export
+shapeString : XlaBuilder -> XlaOp -> IO String
+shapeString builder op = map debugString (getShape builder op)

src/Compiler/Xla/Prim/TensorFlow/Compiler/Xla/Client/XlaBuilder.idr

@@ -36,6 +36,10 @@ export
 %foreign (libxla "XlaBuilder_Build")
 prim__build : GCAnyPtr -> GCAnyPtr -> AnyPtr
 
+export
+%foreign (libxla "XlaBuilder_GetShape")
+prim__getShape : GCAnyPtr -> GCAnyPtr -> PrimIO AnyPtr
+
 export
 %foreign (libxla "XlaBuilder_OpToString")
 prim__opToString : GCAnyPtr -> GCAnyPtr -> String

src/Compiler/Xla/Prim/TensorFlow/Compiler/Xla/Shape.idr

@@ -22,3 +22,7 @@ import Compiler.Xla.Prim.Util
 export
 %foreign (libxla "Shape_delete")
 prim__delete : AnyPtr -> PrimIO ()
+
+export
+%foreign (libxla "Shape_DebugString")
+prim__debugString : GCAnyPtr -> String

src/Compiler/Xla/Prim/Util.idr

@@ -21,6 +21,10 @@ public export
 libxla : String -> String
 libxla fname = "C:" ++ fname ++ ",libc_xla_extension"
 
+export
+%foreign (libxla "print_address")
+prim__printAddress : GCAnyPtr -> PrimIO ()
+
 export
 %foreign (libxla "sizeof_int")
 sizeofInt : Int

src/Compiler/Xla/TensorFlow/Compiler/Xla/Client/XlaBuilder.idr

@@ -65,6 +65,14 @@ build (MkXlaBuilder ptr) (MkXlaOp root)= do
   computationPtr <- onCollectAny computationPtr XlaComputation.delete
   pure (MkXlaComputation computationPtr)
 
+-- is this in IO?
+export
+getShape : HasIO io => XlaBuilder -> XlaOp -> io Xla.Shape
+getShape (MkXlaBuilder builder) (MkXlaOp op) = do
+  shape <- primIO $ prim__getShape builder op
+  shape <- onCollectAny shape Shape.delete
+  pure (MkShape shape)
+
 export
 opToString : XlaBuilder -> XlaOp -> String
 opToString (MkXlaBuilder builderPtr) (MkXlaOp opPtr) = prim__opToString builderPtr opPtr

src/Compiler/Xla/TensorFlow/Compiler/Xla/Shape.idr

@@ -25,3 +25,7 @@ namespace Xla
 export
 delete : AnyPtr -> IO ()
 delete = primIO . prim__delete
+
+export
+debugString : Shape -> String
+debugString (MkShape shape) = prim__debugString shape

src/Compiler/Xla/Util.idr

@@ -33,6 +33,10 @@ boolToCInt : Bool -> Int
 boolToCInt True = 1
 boolToCInt False = 0
 
+export
+printAddress : GCAnyPtr -> IO ()
+printAddress = primIO . prim__printAddress
+
 public export
 data IntArray : Type where
   MkIntArray : GCPtr Int -> IntArray

test.ipkg

@@ -16,6 +16,7 @@ executable = test
 main = Main
 
 modules =
+  Unit.Compiler.TestDebug,
   Unit.Model.TestKernel,
   Unit.Util.TestHashable,
   Unit.TestDistribution,

test/Main.idr

@@ -21,6 +21,7 @@ import Hedgehog
 import TestUtils
 import Utils.TestComparison
 
+import Unit.Compiler.TestDebug
 import Unit.Model.TestKernel
 import Unit.Util.TestHashable
 import Unit.TestDistribution
@@ -32,11 +33,12 @@ covering
 main : IO ()
 main = test [
       Utils.TestComparison.group
-    , TestUtils.group
-    , Unit.Util.TestHashable.group
-    , Unit.TestUtil.group
-    , Unit.TestLiteral.group
-    , Unit.TestTensor.group
-    , Unit.TestDistribution.group
-    , Unit.Model.TestKernel.group
+    -- , TestUtils.group
+    -- , Unit.Util.TestHashable.group
+    -- , Unit.TestUtil.group
+    -- , Unit.TestLiteral.group
+    , Unit.Compiler.TestDebug.group
+    -- , Unit.TestTensor.group
+    -- , Unit.TestDistribution.group
+    -- , Unit.Model.TestKernel.group
   ]

test/Unit/Compiler/TestDebug.idr

@@ -0,0 +1,54 @@
+{--
+Copyright 2022 Joel Berkeley
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+--}
+module Unit.Compiler.TestDebug
+
+import Compiler.Computation
+import Compiler.Debug
+import Compiler.LiteralRW
+import Compiler.Xla.TensorFlow.Compiler.Xla.Literal
+import Compiler.Xla.TensorFlow.Compiler.Xla.XlaData
+import Compiler.Xla.TensorFlow.Compiler.Xla.Client.XlaBuilder
+import Literal
+
+import Utils.Comparison
+import Utils.Cases
+
+export
+shapeString : Property
+shapeString = fixedProperty $ do
+  let str : String = unsafePerformIO $ do
+        lit <- write {dtype=S32} [[0, 1, 2], [3, 4, 5]]
+        builder <- mkXlaBuilder ""
+        op <- constantLiteral builder lit
+        shapeString builder op
+
+  str ===
+    "element_type: S32\n" ++
+    "dimensions: 2\n" ++
+    "dimensions: 3\n" ++
+    "layout {\n" ++
+    "  minor_to_major: 1\n" ++
+    "  minor_to_major: 0\n" ++
+    "  format: DENSE\n" ++
+    "}\n" ++
+    "is_dynamic_dimension: false\n" ++
+    "is_dynamic_dimension: false\n"
+
+export covering
+group : Group
+group = MkGroup "Debug" $ [
+      ("shape string", shapeString)
+  ]