add utilities for accessing XLA shape string within Idris

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());
+    }
 }

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/TensorFlow/Compiler/Xla/Client/XlaBuilder.idr

@@ -65,6 +65,13 @@ build (MkXlaBuilder ptr) (MkXlaOp root)= do
   computationPtr <- onCollectAny computationPtr XlaComputation.delete
   pure (MkXlaComputation computationPtr)
 
+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

test.ipkg

@@ -16,6 +16,8 @@ executable = test
 main = Main
 
 modules =
+  Unit.Compiler.Xla.TensorFlow.Compiler.Xla.Client.TestXlaBuilder,
+
   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.Xla.TensorFlow.Compiler.Xla.Client.TestXlaBuilder
 import Unit.Model.TestKernel
 import Unit.Util.TestHashable
 import Unit.TestDistribution
@@ -36,6 +37,7 @@ main = test [
     , Unit.Util.TestHashable.group
     , Unit.TestUtil.group
     , Unit.TestLiteral.group
+    , Unit.Compiler.Xla.TensorFlow.Compiler.Xla.Client.TestXlaBuilder.group
     , Unit.TestTensor.group
     , Unit.TestDistribution.group
     , Unit.Model.TestKernel.group

test/Unit/Compiler/Xla/TensorFlow/Compiler/Xla/Client/TestXlaBuilder.idr

@@ -0,0 +1,55 @@
+{--
+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.Xla.TensorFlow.Compiler.Xla.Client.TestXlaBuilder
+
+import Compiler.Eval
+import Compiler.LiteralRW
+import Compiler.Xla.TensorFlow.Compiler.Xla.Literal
+import Compiler.Xla.TensorFlow.Compiler.Xla.Shape
+import Compiler.Xla.TensorFlow.Compiler.Xla.XlaData
+import Compiler.Xla.TensorFlow.Compiler.Xla.Client.XlaBuilder
+import Literal
+
+import Utils.Comparison
+import Utils.Cases
+
+export
+xlaOpShapeDebugString : Property
+xlaOpShapeDebugString = fixedProperty $ do
+  let str : String = unsafePerformIO $ do
+        lit <- write {dtype=S32} [[0, 1, 2], [3, 4, 5]]
+        builder <- mkXlaBuilder ""
+        op <- constantLiteral builder lit
+        shape <- getShape builder op
+        pure (debugString shape)
+
+  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" $ [
+      ("XlaOp shape debug string", xlaOpShapeDebugString)
+  ]