Update on "[ET-VK] Enable int8 tiled compute shader to be used with buffer tensors"

## Context

As title. Allow the optimized int8 tiled compute shader to be usable for buffer-backed tensors as well.

## Changes

* Generate buffer variants for the int8 linear tiled shader
* Force the scales tensor to always be a buffer to reduce the number of shader variants that need to be generated.
* Generate an additional variant that computes only 1 output row
* Do not require output rows to be an exact multiple of 4 or 6 to use the tiled implementation

Differential Revision: [D73276277](https://our.internmc.facebook.com/intern/diff/D73276277/)

[ghstack-poisoned]

Author: SS-JIA

examples/models/llama/model.py

@@ -18,6 +18,7 @@
 from executorch.examples.models.llama.llama_transformer import Transformer
 
 from executorch.examples.models.llama.model_args import ModelArgs
+from torchao.utils import TorchAOBaseTensor
 
 try:
     from .fairseq2 import convert_to_llama_checkpoint
@@ -257,6 +258,9 @@ def __init__(self, **kwargs):
                 strict=False,
                 assign=True,
             )  # self.model_ = Transformer(gptconf)
+            for param in self.model_.parameters():
+                if isinstance(param, TorchAOBaseTensor):
+                    param.requires_grad = False
         else:
             print("Checkpoint not provided, defaulting weights to zeros.")
             self.model_.to_empty(device="cpu")

extension/llm/export/builder.py

@@ -41,6 +41,7 @@
 from torch.ao.quantization.quantizer.composable_quantizer import ComposableQuantizer
 from torch.export import export_for_training, ExportedProgram
 from torch.nn.attention import SDPBackend
+from torchao.utils import unwrap_tensor_subclass
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
 logging.basicConfig(level=logging.INFO, format=FORMAT)
@@ -199,6 +200,11 @@ def _get_edge_config(self) -> EdgeCompileConfig:
         return edge_config
 
     def _export(self, module: Optional[torch.nn.Module] = None) -> ExportedProgram:
+        if module is not None:
+            unwrap_tensor_subclass(module)
+        else:
+            unwrap_tensor_subclass(self.model)
+
         dynamic_shape = self._get_dynamic_shape()
         # 1. torch.nn.attention.sdpa_kernel([SDPBackend.MATH]) is for bypassing the dynamo error when tracing
         # 2. torch.no_grad() is for getting rid of the dropout (not sure why training ops will show up)

kernels/aten/functions.yaml

@@ -423,6 +423,8 @@
 
 - op: var.out
 
+- op: view_as_real_copy.out
+
 - op: view_copy.out
 
 - op: where.self_out

kernels/portable/cpu/op_view_as_real_copy.cpp

@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/portable/cpu/util/copy_ops_util.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+#include <executorch/runtime/platform/assert.h>
+
+namespace torch {
+namespace executor {
+namespace native {
+
+using Tensor = executorch::aten::Tensor;
+
+namespace {
+
+template <typename SELF_CTYPE, typename OUT_CTYPE>
+inline void _to_impl(const Tensor& self, Tensor& out) {
+  auto self_data = self.mutable_data_ptr<SELF_CTYPE>();
+  auto out_data = out.mutable_data_ptr<OUT_CTYPE>();
+
+  for (size_t i = 0, e = self.numel(); i < e; i++) {
+    auto val_in = self_data[i];
+    out_data[2 * i] = static_cast<OUT_CTYPE>(val_in.real_);
+    out_data[2 * i + 1] = static_cast<OUT_CTYPE>(val_in.imag_);
+  }
+}
+
+} // namespace
+
+// view_as_real_copy(Tensor self) -> Tensor
+Tensor& view_as_real_copy_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& self,
+    Tensor& out) {
+  (void)ctx;
+
+  // Get the output shape
+  Tensor::SizesType expected_output_size[kTensorDimensionLimit];
+  get_view_as_real_copy_out_target_size(self, expected_output_size);
+
+  // Resize for dynamic shape
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      resize_tensor(
+          out, {expected_output_size, static_cast<size_t>(out.dim())}) ==
+          Error::Ok,
+      InvalidArgument,
+      out,
+      "Failed to resize output tensor.");
+
+  // The input tensor must be complex type
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      executorch::runtime::isComplexType(self.scalar_type()),
+      InvalidArgument,
+      out,
+      "Input tensor must be complex type");
+
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(self, out), InvalidArgument, out);
+
+  constexpr auto op_name = "view_as_real_copy.out";
+
+  ET_SWITCH_COMPLEXH_TYPES(self.scalar_type(), ctx, op_name, CTYPE_IN, [&] {
+    ET_SWITCH_FLOATH_TYPES(out.scalar_type(), ctx, op_name, CTYPE_OUT, [&] {
+      _to_impl<CTYPE_IN, CTYPE_OUT>(self, out);
+    });
+  });
+
+  return out;
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/portable/cpu/util/copy_ops_util.cpp

@@ -1018,5 +1018,14 @@ void get_unfold_copy_out_target_size(
   *out_ndim = self.dim() + 1;
 }
 
+void get_view_as_real_copy_out_target_size(
+    const Tensor& self,
+    executorch::aten::SizesType* out_sizes) {
+  for (auto i : c10::irange(self.dim())) {
+    out_sizes[i] = self.size(i);
+  }
+  out_sizes[self.dim()] = 2;
+}
+
 } // namespace executor
 } // namespace torch

kernels/portable/cpu/util/copy_ops_util.h

@@ -247,5 +247,9 @@ void get_unfold_copy_out_target_size(
     executorch::aten::SizesType* out_sizes,
     size_t* out_ndim);
 
+void get_view_as_real_copy_out_target_size(
+    const Tensor& self,
+    executorch::aten::SizesType* out_sizes);
+
 } // namespace executor
 } // namespace torch

kernels/portable/functions.yaml

@@ -957,6 +957,11 @@
     - arg_meta: null
       kernel_name: torch::executor::var_out
 
+- op: view_as_real_copy.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::view_as_real_copy_out
+
 - op: view_copy.out
   kernels:
     - arg_meta: null

kernels/test/CMakeLists.txt

@@ -242,6 +242,7 @@ set(all_test_sources
     "op_upsample_bilinear2d_test.cpp"
     "op_upsample_nearest2d_test.cpp"
     "op_var_test.cpp"
+    "op_view_as_real_copy_test.cpp"
     "op_view_copy_test.cpp"
     "op_where_test.cpp"
     "op_zeros_test.cpp"

kernels/test/op_view_as_real_copy_test.cpp

@@ -0,0 +1,86 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
+#include <executorch/kernels/test/TestUtil.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
+
+#include <gtest/gtest.h>
+
+using namespace ::testing;
+using executorch::aten::ScalarType;
+using executorch::aten::Tensor;
+using torch::executor::testing::TensorFactory;
+
+class OpViewAsRealTest : public OperatorTest {
+ protected:
+  Tensor& view_as_real_copy_out(const Tensor& self, Tensor& out) {
+    return torch::executor::aten::view_as_real_copy_outf(context_, self, out);
+  }
+
+  template <typename CTYPE, ScalarType DTYPE>
+  void run_complex_smoke_test() {
+    TensorFactory<DTYPE> tf;
+    constexpr auto REAL_DTYPE = executorch::runtime::toRealValueType(DTYPE);
+    TensorFactory<REAL_DTYPE> tf_out;
+
+    Tensor in = tf.make(
+        {2, 2},
+        {CTYPE(3, 4), CTYPE(-1.7, 7.4), CTYPE(5, -12), CTYPE(8.3, 0.1)});
+    Tensor out = tf_out.zeros({2, 2, 2});
+    Tensor expected =
+        tf_out.make({2, 2, 2}, {3, 4, -1.7, 7.4, 5, -12, 8.3, 0.1});
+    Tensor ret = view_as_real_copy_out(in, out);
+
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
+
+  // Tests on tensors with 0 size
+  template <typename CTYPE, ScalarType DTYPE>
+  void test_empty_input() {
+    TensorFactory<DTYPE> tf;
+    constexpr auto REAL_DTYPE = executorch::runtime::toRealValueType(DTYPE);
+    TensorFactory<REAL_DTYPE> tf_out;
+
+    Tensor in = tf.make(/*sizes=*/{3, 0, 4}, /*data=*/{});
+    Tensor out = tf_out.zeros({3, 0, 4, 2});
+    Tensor expected = tf_out.make(/*sizes=*/{3, 0, 4, 2}, /*data=*/{});
+    Tensor ret = view_as_real_copy_out(in, out);
+
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
+
+  // Tests on 0-dim input tensors
+  template <typename CTYPE, ScalarType DTYPE>
+  void zero_dim_input() {
+    TensorFactory<DTYPE> tf;
+    constexpr auto REAL_DTYPE = executorch::runtime::toRealValueType(DTYPE);
+    TensorFactory<REAL_DTYPE> tf_out;
+
+    Tensor in = tf.make(/*sizes=*/{}, {CTYPE(0, 0)});
+    Tensor out = tf_out.zeros({2});
+    Tensor expected = tf_out.zeros(/*sizes=*/{2});
+    Tensor ret = view_as_real_copy_out(in, out);
+
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_EQ(out, expected);
+  }
+};
+
+TEST_F(OpViewAsRealTest, ComplexSmokeTest) {
+#define RUN_SMOKE_TEST(ctype, dtype)                  \
+  run_complex_smoke_test<ctype, ScalarType::dtype>(); \
+  test_empty_input<ctype, ScalarType::dtype>();       \
+  zero_dim_input<ctype, ScalarType::dtype>();
+  ET_FORALL_COMPLEXH_TYPES(RUN_SMOKE_TEST);
+#undef RUN_SMOKE_TEST
+}

kernels/test/targets.bzl

@@ -331,6 +331,7 @@ def define_common_targets():
     _common_op_test("op_upsample_bilinear2d_test", ["aten", "portable"])
     _common_op_test("op_upsample_nearest2d_test", ["aten", "portable"])
     _common_op_test("op_var_test", ["aten", "portable"])
+    _common_op_test("op_view_as_real_copy_test", ["aten", "portable"])
     _common_op_test("op_view_copy_test", ["aten", "portable"])
     _common_op_test("op_where_test", ["aten", "portable"])
     _common_op_test("op_zeros_test", ["aten", "portable"])