Remove all handling of EXLA.Shape tuple (#1475)

Author: jonatanklosko

exla/Makefile

@@ -66,7 +66,7 @@ HEADERS = $(EXLA_DIR)/mlir/ops.h $(EXLA_DIR)/mlir/builder.h $(EXLA_DIR)/mlir/cus
 OBJECTS = $(patsubst $(EXLA_DIR)/%.cc,$(EXLA_CACHE_OBJ_DIR)/%.o,$(SOURCES)) $(EXLA_CACHE_OBJ_DIR)/exla_cuda.o
 
 
-NVCC_RESULT := $(shell which nvcc 2> NULL)
+NVCC_RESULT := $(shell which nvcc 2> /dev/null)
 NVCC_TEST := $(notdir $(NVCC_RESULT))
 
 ifeq ($(NVCC_TEST),nvcc)

exla/c_src/exla/exla.cc

@@ -323,22 +323,6 @@ ERL_NIF_TERM make_shape(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
   return exla::nif::ok(env, exla::nif::make<xla::Shape>(env, shape));
 }
 
-ERL_NIF_TERM make_tuple_shape(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
-  if (argc != 1) {
-    return exla::nif::error(env, "Bad argument count.");
-  }
-
-  std::vector<xla::Shape> shapes;
-
-  if (!exla::nif::get_list<xla::Shape>(env, argv[0], shapes)) {
-    return exla::nif::error(env, "Unable to get shapes.");
-  }
-
-  xla::Shape shape = xla::ShapeUtil::MakeTupleShape(shapes);
-
-  return exla::nif::ok(env, exla::nif::make<xla::Shape>(env, shape));
-}
-
 ERL_NIF_TERM make_token_shape(ErlNifEnv* env, int argc, const ERL_NIF_TERM argv[]) {
   if (argc != 0) {
     return exla::nif::error(env, "Bad argument count.");
@@ -823,7 +807,6 @@ static ErlNifFunc exla_funcs[] = {
     // Shape
     {"make_shape", 2, make_shape},
     {"make_token_shape", 0, make_token_shape},
-    {"make_tuple_shape", 1, make_tuple_shape},
     {"get_shape_info", 1, get_shape_info},
     // Log Sink
     {"start_log_sink", 1, start_log_sink},

exla/c_src/exla/exla_client.cc

@@ -393,38 +393,6 @@ xla::Status ExlaClient::TransferToInfeed(ErlNifEnv* env,
                                          ERL_NIF_TERM data,
                                          const xla::Shape& shape,
                                          int device_id) {
-  // Tuples need to be decomposed a bit
-  if (shape.IsTuple()) {
-    // unsupported right now
-    if (xla::ShapeUtil::IsNestedTuple(shape)) {
-      return xla::InvalidArgument("nested tuples are not supported in infeed operation");
-    }
-
-    int num_elements = xla::ShapeUtil::TupleElementCount(shape);
-    std::vector<const char*> buf_ptrs;
-    buf_ptrs.reserve(num_elements);
-
-    ERL_NIF_TERM head, tail;
-    while (enif_get_list_cell(env, data, &head, &tail)) {
-      ErlNifBinary tmp_bin;
-      if (!nif::get_binary(env, head, &tmp_bin)) {
-        return xla::InvalidArgument("infeed operation expects a list of binaries");
-      }
-
-      const char* data_ptr = const_cast<char*>(reinterpret_cast<char*>(tmp_bin.data));
-      buf_ptrs.push_back(data_ptr);
-      data = tail;
-    }
-
-    xla::BorrowingLiteral literal(buf_ptrs, shape);
-
-    EXLA_ASSIGN_OR_RETURN(xla::PjRtDevice * device, client_->LookupDevice(device_id));
-
-    xla::Status status = device->TransferToInfeed(literal);
-
-    return status;
-  }
-
   // Fast path to avoid any traversal when not sending Tuples
   ERL_NIF_TERM head, tail;
   if (!enif_get_list_cell(env, data, &head, &tail)) {

exla/c_src/exla/exla_nif_util.cc

@@ -479,15 +479,8 @@ int get_primitive_type(ErlNifEnv* env, ERL_NIF_TERM term, xla::PrimitiveType* ty
 
 ERL_NIF_TERM make_shape_info(ErlNifEnv* env, xla::Shape shape) {
   if (shape.IsTuple()) {
-    int element_count = xla::ShapeUtil::TupleElementCount(shape);
-    std::vector<ERL_NIF_TERM> terms;
-    terms.reserve(element_count);
-    for (int i = 0; i < element_count; i++) {
-      xla::Shape shape_elem = xla::ShapeUtil::GetTupleElementShape(shape, i);
-      ERL_NIF_TERM shape_term = make<xla::Shape>(env, shape_elem);
-      terms.push_back(shape_term);
-    }
-    return enif_make_list_from_array(env, &terms[0], element_count);
+    std::cerr << "Unexpected tuple shape" << std::endl;
+    exit(1);
   } else if (shape.IsArray()) {
     xla::PrimitiveType type = shape.element_type();
     absl::Span<const int64> dims = shape.dimensions();

exla/lib/exla/defn/stream.ex

@@ -88,14 +88,7 @@ defmodule EXLA.Defn.Stream do
 
       data_and_shapes =
         if client.platform == :host do
-          # TODO: Remove first-clause once EXLA.OP is removed
-          shapes =
-            case send_shape do
-              %EXLA.Shape{dtype: {:tuple, shapes}} -> shapes
-              l when is_list(l) -> l
-            end
-
-          Enum.zip(buffers, shapes)
+          Enum.zip(buffers, send_shape)
         else
           [{buffers, send_shape}]
         end

exla/lib/exla/executable.ex

@@ -87,7 +87,6 @@ defmodule EXLA.Executable do
     shapes =
       Enum.flat_map(List.wrap(shapes), fn shape ->
         case shape do
-          %Shape{dtype: {:tuple, shapes}} -> shapes
           shapes when is_list(shapes) -> shapes
           %Shape{} -> [shape]
         end
@@ -102,18 +101,8 @@ defmodule EXLA.Executable do
     end)
   end
 
-  defp strip_shape(%Shape{dtype: {:tuple, shapes}}) do
-    subshapes = Enum.map(shapes, &strip_shape/1)
-    %{dtype: {:tuple, subshapes}, dims: {length(subshapes)}}
-  end
-
   defp strip_shape(%Shape{dtype: dtype, dims: dims}), do: %{dtype: dtype, dims: dims}
 
-  defp reconstruct_shapes(%{dtype: {:tuple, shapes}}) do
-    subshapes = Enum.map(shapes, &reconstruct_shapes/1)
-    EXLA.Shape.make_tuple_shape(subshapes)
-  end
-
   defp reconstruct_shapes(%{dtype: dtype, dims: dims}) do
     EXLA.Shape.make_shape(dtype, dims)
   end

exla/lib/exla/nif.ex

@@ -191,9 +191,6 @@ defmodule EXLA.NIF do
   def make_token_shape(),
     do: :erlang.nif_error(:undef)
 
-  def make_tuple_shape(_shapes),
-    do: :erlang.nif_error(:undef)
-
   def get_host_client(),
     do: :erlang.nif_error(:undef)
 

exla/lib/exla/shape.ex

@@ -11,14 +11,8 @@ defmodule EXLA.Shape do
 
   @doc false
   def get_shape_info(ref) when is_reference(ref) do
-    case EXLA.NIF.get_shape_info(ref) |> unwrap!() do
-      {dims_term, type_str} ->
-        %Shape{dims: dims_term, dtype: charlist_to_dtype(type_str), ref: ref}
-
-      children when is_list(children) ->
-        children = Enum.map(children, &get_shape_info/1)
-        %Shape{dims: {length(children)}, dtype: {:tuple, children}, ref: ref}
-    end
+    {dims_term, type_str} = EXLA.NIF.get_shape_info(ref) |> unwrap!()
+    %Shape{dims: dims_term, dtype: charlist_to_dtype(type_str), ref: ref}
   end
 
   @doc """
@@ -38,18 +32,6 @@ defmodule EXLA.Shape do
     %Shape{dims: {}, dtype: :token, ref: ref}
   end
 
-  @doc """
-  Creates a tuple shape with the given shapes.
-  """
-  def make_tuple_shape(shapes) when is_list(shapes) do
-    refs =
-      shapes
-      |> Enum.map(& &1.ref)
-
-    ref = EXLA.NIF.make_tuple_shape(refs) |> unwrap!()
-    %Shape{dims: {length(shapes)}, dtype: {:tuple, shapes}, ref: ref}
-  end
-
   defp validate_dims!(_dims, 0), do: :ok
 
   defp validate_dims!(dims, i)
@@ -63,10 +45,6 @@ defmodule EXLA.Shape do
   @doc """
   Returns the shape size in bytes.
   """
-  def byte_size(%EXLA.Shape{dtype: {:tuple, shapes}}) do
-    Enum.reduce(shapes, 0, &(byte_size(&1) + &2))
-  end
-
   def byte_size(%EXLA.Shape{dtype: {_, bit_size}, dims: dims}) do
     Tuple.product(dims) * div(bit_size, 8)
   end

exla/test/exla/executable_test.exs

@@ -23,7 +23,7 @@ defmodule EXLA.ExecutableTest do
       t2 = BinaryBuffer.from_binary(<<1::32-native>>, Shape.make_shape({:s, 32}, {}))
 
       assert [a = %DeviceBuffer{}] =
-               run_one([t1, t2], [], Shape.make_tuple_shape([t1.shape]), fn b, x, y ->
+               run_one([t1, t2], [], [t1.shape], fn b, x, y ->
                  [Value.add(b, x, y)]
                end)
 
@@ -53,7 +53,7 @@ defmodule EXLA.ExecutableTest do
                end)
 
       assert [%DeviceBuffer{}] =
-               run_one([t1, t2], [], Shape.make_tuple_shape([t1.shape]), fn b, x, y ->
+               run_one([t1, t2], [], [t1.shape], fn b, x, y ->
                  [Value.add(b, x, y)]
                end)
 
@@ -88,7 +88,7 @@ defmodule EXLA.ExecutableTest do
       t2 = BinaryBuffer.from_binary(<<2::32-native>>, Shape.make_shape({:s, 32}, {}))
 
       assert [a = %DeviceBuffer{}] =
-               run_one([t1, t2], [], {t1.shape}, fn b, x, y ->
+               run_one([t1, t2], [], [t1.shape], fn b, x, y ->
                  [Value.add(b, x, y)]
                end)
 
@@ -100,7 +100,7 @@ defmodule EXLA.ExecutableTest do
       t2 = BinaryBuffer.from_binary(<<2::32-native>>, Shape.make_shape({:s, 32}, {}))
 
       assert [a = %DeviceBuffer{}, b = %DeviceBuffer{}] =
-               run_one([t1, t2], [], Shape.make_tuple_shape([t1.shape, t2.shape]), fn _b, x, y ->
+               run_one([t1, t2], [], [t1.shape, t2.shape], fn _b, x, y ->
                  [x, y]
                end)
 
@@ -117,7 +117,7 @@ defmodule EXLA.ExecutableTest do
                run_one(
                  [t1, t2],
                  [device_id: 1],
-                 EXLA.Shape.make_tuple_shape([t1.shape, t2.shape, t1.shape]),
+                 [t1.shape, t2.shape, t1.shape],
                  fn b, x, y ->
                    [x, y, Value.add(b, x, y)]
                  end
@@ -158,7 +158,7 @@ defmodule EXLA.ExecutableFeedTest do
 
       assert res =
                Task.async(fn ->
-                 run_one([], [], Shape.make_tuple_shape([Shape.make_token_shape()]), fn b ->
+                 run_one([], [], [Shape.make_token_shape()], fn b ->
                    token = Value.create_token(b)
 
                    {new_token, [val]} = Value.infeed(token, t.shape)
@@ -183,7 +183,7 @@ defmodule EXLA.ExecutableFeedTest do
 
       assert res =
                Task.async(fn ->
-                 run_one([], [], {token_shape, t.shape}, fn b ->
+                 run_one([], [], [token_shape, t.shape], fn b ->
                    token = Value.create_token(b)
 
                    {token, [val]} = Value.infeed(token, t.shape)

exla/test/exla/shape_test.exs

@@ -16,38 +16,4 @@ defmodule EXLA.ShapeTest do
       assert Shape.byte_size(shape) == 2
     end
   end
-
-  describe "make_tuple_shape/1" do
-    test "creates tuple shape" do
-      s1 = Shape.make_shape({:s, 32}, {5, 5, 5})
-      s2 = Shape.make_shape({:bf, 16}, {})
-      s3 = Shape.make_shape({:f, 32}, {1, 1})
-
-      shape = Shape.make_tuple_shape([s1, s2, s3])
-      assert %Shape{dtype: {:tuple, [_, _, _]}, dims: {3}, ref: ref} = shape
-      assert Shape.byte_size(shape) == 506
-      assert %Shape{dtype: {:tuple, [_, _, _]}, dims: {3}, ref: ^ref} = Shape.get_shape_info(ref)
-    end
-
-    test "creates nested tuples" do
-      s1 = Shape.make_shape({:s, 32}, {5, 5, 5})
-      s2 = Shape.make_shape({:bf, 16}, {})
-      s3 = Shape.make_shape({:f, 32}, {1, 1})
-      s4 = Shape.make_shape({:s, 32}, {1})
-      t1 = Shape.make_tuple_shape([s1, s2, s3])
-
-      shape = Shape.make_tuple_shape([s4, t1])
-
-      assert %Shape{dtype: {:tuple, [_, %Shape{dtype: {:tuple, [_, _, _]}}]}, dims: {2}, ref: ref} =
-               shape
-
-      assert Shape.byte_size(shape) == 510
-
-      assert %Shape{
-               dtype: {:tuple, [_, %Shape{dtype: {:tuple, [_, _, _]}}]},
-               dims: {2},
-               ref: ^ref
-             } = Shape.get_shape_info(ref)
-    end
-  end
 end