shlo_ref: Add convolution op and unit tests #66299

LokeshReddyOVS-MCW · 2024-04-23T15:07:15Z

Adds initial version of convolution op
Adds unit tests for convolution op

- Adds initial version of convolution op - Adds unit tests for convolution op

google-cla · 2024-04-23T15:07:19Z

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

rascani

Some initial comments, but I still have a lot more to go through.

rascani · 2024-04-23T22:01:14Z

tensorflow/lite/experimental/shlo/ops/convolution.cc

+  return std::unique(vec.begin(), vec.end()) == vec.end();
+}
+
+bool IsInRange(DimVector<int64_t>& vec, size_t N) {


vec can be const.

removed the vec, using set now(you can check it in PR#67245)
using const in appropriate places

rascani · 2024-04-23T22:05:04Z

tensorflow/lite/experimental/shlo/ops/convolution.cc

+using DimVector = absl::InlinedVector<T, 6>;
+
+bool IsUnique(DimVector<int64_t>& vec) {
+  std::sort(vec.begin(), vec.end());


This will modify the vec. You may want to consider making a copy, sorting, then testing for uniqueness. You could also use a set.

changed to set, you can check it in PR#67245

rascani · 2024-04-23T23:07:15Z

tensorflow/lite/experimental/shlo/ops/convolution.cc

+  }
+  // malloc is used to have the storage space available out of prepare function
+  // scope and it's pointer is stored in class data member to
+  // deallocate the memory in destructor.


I'd prefer to figure out a container for storing these rather than using malloc directly. I made a utility class TensorWithData that could pair a Tensor with the backing data in a std::vector<std::byte>, but it was originally aimed at tests. Would something like that work?

That said, do all of these need to be Tensors? The permutations & transposed seem to simply be 1d vectors.

Alternatively, you could use a std::vector<std::byte> directly that would handle the memory for you in place of malloc calls.

Ideally, the ops should not own any memory and only keep a pointer to the external data that is provided. I understand that if you want to call another op you may need to create that op's parameters.

changed to std::vectorstd::byte(you can check it in PR#67245)

rascani · 2024-04-23T23:08:36Z

tensorflow/lite/experimental/shlo/ops/convolution.cc

+  const int64_t* window_strides_pointer =
+      op.attributes.window_strides.GetDataAs<DataType::kSI64>();
+
+  // Constraints Check


We should perform constraint checks before any other prepare work is done.

rascani · 2024-04-23T23:15:28Z

tensorflow/lite/experimental/shlo/ops/convolution_test.cc

+      .precision_configs = precision_configs});
+
+  Vector<StorageT> expected_data;
+  if (std::is_same<StorageT, BF16>::value) {


These expected results seem to be differing to a larger degree than I would expect. Where were these generated from?

we got the expected from StableHlo repo, now float types are seperated into different tests(you can check it in PR#67245)

rascani · 2024-04-23T23:18:39Z

tensorflow/lite/experimental/shlo/ops/convolution_test.cc

+  ASSERT_OK(Evaluate(op, lhs, rhs, output_tensor));
+
+  constexpr double kEpsilon = 0.1;
+  EXPECT_THAT(output_data, Pointwise(FloatEq(), expected_data));


Why is this using FloatEq for integer storage?

changed to Eq() for integer storage(you can check it in PR#67245)

rascani · 2024-04-23T23:19:04Z

tensorflow/lite/experimental/shlo/ops/convolution_test.cc

+TYPED_TEST_SUITE(QuantizedIntConvolutionTest, QuantizedTestTypes,
+                 TestParamNames);
+
+TYPED_TEST(QuantizedIntConvolutionTest, PerTensorsRaiseAnError) {


This test doesn't seem to expect an error to be raised.

changed to appropriate tests names

rascani · 2024-04-23T23:19:35Z

tensorflow/lite/experimental/shlo/ops/convolution_test.cc

+
+  Vector<StorageT> expected_data =
+      Vector<StorageT>{5, 10, 15, 20, 25, 5, 10, 15, 20, -25};
+  ;


Extra semicolon

rascani · 2024-04-23T23:20:05Z

tensorflow/lite/experimental/shlo/ops/convolution_test.cc

+  EXPECT_THAT(output_data, Pointwise(FloatEq(), expected_data));
+}
+
+TYPED_TEST(QuantizedIntConvolutionTest, PerAxisRaiseAnError) {


This test doesn't seem to raise an error.

changed to appropriate test name

rascani · 2024-04-23T23:26:27Z

tensorflow/lite/experimental/shlo/ops/convolution_test.cc

+  Vector<StorageT> expected_data =
+      Vector<StorageT>{5, 10, 15, 20, 25, 5, 10, 15, 20, -25};
+  ;
+  if (std::is_same<StorageT, I4>::value) {


I'd suggest splitting out a separate test case for i4 where we can ensure the input and expected data are within the i4 range without having to clamp most of the values.

resolved by calling quantize function ,now works for all int types without clamping

qukhan · 2024-04-24T17:23:58Z

tensorflow/lite/experimental/shlo/ops/convolution.cc

+  using StorageT = StorageType<storage_type>;
+
+  // Transpose prepare
+  const int64_t* window_spacial_pointer =


~~spacial~~ -> spatial

qukhan · 2024-04-24T17:37:17Z

tensorflow/lite/experimental/shlo/ops/convolution.cc

+  for (int64_t dim : vec) {
+    if (dim >= N || dim < 0) {
+      return false;
+    }
+  }
+  return true;


Suggested change

for (int64_t dim : vec) {

if (dim >= N || dim < 0) {

return false;

}

}

return true;

return absl::c_all_of(dim, [N](int64_t v) { return v >= N || v < 0; });

changed to absl::c_all_of(you can check it in PR#67245)

qukhan · 2024-04-24T17:38:46Z

tensorflow/lite/experimental/shlo/ops/convolution.cc

+      .rhs_contracting_dimensions = rhs_contracting_dimensions,
+      .precision_configs = precision_configs});
+
+  auto state = Prepare(op.dot_general_op, lhs_dot_general, rhs_dot_general,


Unless the type is ridiculously complicated, prefer explicit types.

stopped auto usage

qukhan · 2024-04-24T17:42:03Z

tensorflow/lite/experimental/shlo/ops/convolution.cc

+        "stablehlo.convolution: Size of precision_config must be two.");
+  }
+  if (op.attributes.precision_configs[0] != PrecisionTypes::DEFAULT &&
+      op.attributes.precision_configs[1] != PrecisionTypes::DEFAULT) {


This should probably be an ||. You could also use the same implementation for all the configs.

qukhan · 2024-04-24T17:42:42Z

tensorflow/lite/experimental/shlo/ops/convolution.cc

+  for (size_t i = 0; i < n; ++i) {
+    if (check_buffer[i] == 0) {
+      is_greater_than_zero = false;
+      exit;


qukhan · 2024-04-24T17:57:40Z

tensorflow/lite/experimental/shlo/ops/convolution.cc

+// Padding op basic implimentation in context of Convolution
+template <DataType storage_type>
+void PaddingOp(ConvolutionOp& op, const Tensor& x, const Tensor& padding,
+               const Tensor& lhs_dilations) {
+  using StorageT = StorageType<storage_type>;
+  using int64_t = StorageType<DataType::kSI64>;
+  const StorageT* x_buffer = x.GetDataAs<storage_type>();
+  const int64_t* lhs_dilation_buffer =
+      lhs_dilations.GetDataAs<DataType::kSI64>();
+  const int64_t* padding_buffer = padding.GetDataAs<DataType::kSI64>();
+  StorageT* lhs_buffer = op.lhs_padded.GetDataAs<storage_type>();
+  size_t j = 0;
+  for (size_t i = 0; i < op.lhs_padded.NumElements(); ++i) {
+    int x_spacials[x.Rank() - 2];
+    size_t depth = 1;
+
+    for (int64_t m = x.Rank() - 3; m >= 0; --m) {
+      x_spacials[m] =
+          (i / depth) % static_cast<size_t>(op.lhs_padded.shape().Dim(m + 2));
+      depth *= static_cast<size_t>(op.lhs_padded.shape().Dim(m + 2));
+    }
+    bool check = true;
+    for (int64_t k = x.Rank() - 3; k >= 0; --k) {
+      check *= x_spacials[k] >= (padding_buffer[2 * k]) &&
+               x_spacials[k] < x.shape().Dim(k + 2) +
+                                   (lhs_dilation_buffer[k] - 1) *
+                                       (x.shape().Dim(k + 2) - 1) +
+                                   padding_buffer[2 * k];
+    }
+
+    if (check) {
+      for (int64_t k = x.Rank() - 3; k >= 0; --k) {
+        check *= static_cast<size_t>(lhs_dilation_buffer[k]) != 0;
+      }
+      if (check) {
+        for (int64_t k = x.Rank() - 3; k >= 0; --k) {
+          check *= static_cast<size_t>(x_spacials[k] - padding_buffer[2 * k]) %
+                       static_cast<size_t>(lhs_dilation_buffer[k]) ==
+                   0;
+        }
+        if (check) {
+          lhs_buffer[i] = x_buffer[j++];
+        }
+      }
+    }
+  }
+}


StableHLO Pad already has a somewhat optimized version in TFLite that you should try to adapt: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/kernels/stablehlo_pad.cc

adapted stablehlo_pad function(you can check it in PR#67246)

qukhan · 2024-04-24T17:58:28Z

tensorflow/lite/experimental/shlo/ops/convolution.h

+    free(lhs_permutation_data);
+    free(lhs_transposed_data);
+    free(rhs_permutation_data);
+    free(rhs_transposed_data);
+    free(output_permutation_data);
+    free(output_transposed_data);
+    free(lhs_padded_data);
+    free(lhs_dot_general_data);
+    free(rhs_dot_general_data);
+    free(output_dot_general_data);
+    free(lhs_contracting_dimensions_data);
+    free(rhs_contracting_dimensions_data);


Make use of RAII.

qukhan · 2024-04-24T18:02:36Z

tensorflow/lite/experimental/shlo/ops/convolution.h

+    Tensor window_strides;
+    Tensor padding;
+    Tensor lhs_dilation;
+    Tensor rhs_dilation;
+    Tensor window_reversal;
+    int64_t input_batch_dimension;
+    int64_t input_feature_dimension;
+    Tensor input_spacial_dimensions;
+    int64_t kernel_input_feature_dimension;
+    int64_t kernel_output_feature_dimension;
+    Tensor kernel_spacial_dimensions;
+    int64_t output_batch_dimension;
+    int64_t output_feature_dimension;
+    Tensor output_spacial_dimensions;
+    int64_t feature_group_count;
+    int64_t batch_group_count;


Most of these tensors are defined in the spec as 64 bit int 1D tensors. Use an absl::Span<int64_t>.

changed 1D tensors to span(you can check it in PR#67245)

qukhan · 2024-04-24T18:03:36Z

tensorflow/lite/experimental/shlo/ops/convolution.h

+    Tensor output_spacial_dimensions;
+    int64_t feature_group_count;
+    int64_t batch_group_count;
+    absl::InlinedVector<PrecisionTypes, 2> precision_configs;


This has a fixed size in the spec. No need for the flexibility of absl::InlinedVector. Use std::array<PrecisionTypes, 2>.

changed to std::array<PrecisionTypes, 2>

qukhan · 2024-04-24T18:04:37Z

tensorflow/lite/experimental/shlo/ops/convolution.h

+  void* lhs_permutation_data;
+  void* lhs_transposed_data;
+  void* rhs_permutation_data;
+  void* rhs_transposed_data;
+  void* output_permutation_data;
+  void* output_transposed_data;
+  void* lhs_padded_data;
+  void* lhs_dot_general_data;
+  void* rhs_dot_general_data;
+  void* output_dot_general_data;
+  void* lhs_contracting_dimensions_data;
+  void* rhs_contracting_dimensions_data;
+  void* lhs_dequantized_data;
+  void* rhs_dequantized_data;
+  void* output_dequantized_data;


If these are owned by the op, make use of RAII.

gbaned · 2024-04-26T12:30:30Z

Hi @LokeshReddyOVS-MCW This PR is in draft, any update on this? Please. Thank you!

shlo_ref add convolution op

55683d8

- Adds initial version of convolution op - Adds unit tests for convolution op

google-ml-butler bot added the size:XL CL Change Size:Extra Large label Apr 23, 2024

google-ml-butler bot assigned gbaned Apr 23, 2024

LokeshReddyOVS-MCW marked this pull request as draft April 23, 2024 16:07

rascani unassigned gbaned Apr 23, 2024

rascani reviewed Apr 23, 2024

View reviewed changes

qukhan reviewed Apr 24, 2024

View reviewed changes

gbaned added this to Assigned Reviewer in PR Queue via automation Apr 26, 2024

shlo_ref: Add convolution op and unit tests #66299

Are you sure you want to change the base?

shlo_ref: Add convolution op and unit tests #66299

Conversation

LokeshReddyOVS-MCW commented Apr 23, 2024

google-cla bot commented Apr 23, 2024

rascani left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gbaned commented Apr 26, 2024