Update on "[ET-VK][ez] Implement helper functions to get fastest moving dim"

Add C++ and GLSL helpers to query the fastest moving dimension (the
dimension with stride 1 in buffer layout). This is useful for optimizing
memory access patterns in shaders, as iterating along the fastest moving
dimension maximizes cache locality.

The C++ `fastest_whcn_dim()` method accounts for block-transposed layouts by
returning `outer_packed_dim` instead of `packed_dim` when applicable. A
corresponding GLSL macro extracts this info from the hashed layout.

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

[ghstack-poisoned]

Author: ssjia

.ci/scripts/export_model_artifact.sh

@@ -157,10 +157,11 @@ pip list
 if [ "$MODEL_NAME" = "parakeet" ]; then
   pip install -r examples/models/parakeet/install_requirements.txt
 
-  python examples/models/parakeet/export_parakeet_tdt.py \
+  python -m executorch.examples.models.parakeet.export_parakeet_tdt \
       --backend "$DEVICE" \
       --output-dir "${OUTPUT_DIR}" \
-      --dtype bf16
+      --dtype bf16 \
+      ${EXTRA_ARGS}
 
   test -f "${OUTPUT_DIR}/model.pte"
   # CUDA saves named data to separate .ptd file, Metal embeds in .pte

.github/workflows/cuda.yml

@@ -150,15 +150,6 @@ jobs:
               repo: "google"
               name: "gemma-3-4b-it"
             quant: "quantized-int4-weight-only"
-          # Parakeet only supports non-quantized
-          - model:
-              repo: "nvidia"
-              name: "parakeet-tdt"
-            quant: "quantized-int4-tile-packed"
-          - model:
-              repo: "nvidia"
-              name: "parakeet-tdt"
-            quant: "quantized-int4-weight-only"
     with:
       timeout: 90
       secrets-env: EXECUTORCH_HF_TOKEN
@@ -219,15 +210,6 @@ jobs:
               repo: "google"
               name: "gemma-3-4b-it"
             quant: "quantized-int4-weight-only"
-          # Parakeet only supports non-quantized
-          - model:
-              repo: "nvidia"
-              name: "parakeet-tdt"
-            quant: "quantized-int4-tile-packed"
-          - model:
-              repo: "nvidia"
-              name: "parakeet-tdt"
-            quant: "quantized-int4-weight-only"
     with:
       timeout: 90
       runner: linux.g5.4xlarge.nvidia.gpu

.github/workflows/trunk.yml

@@ -315,7 +315,6 @@ jobs:
           - test_arm_baremetal: test_run_ethos_u85
           - test_arm_baremetal: test_smaller_stories_llama
           - test_arm_baremetal: test_memory_allocation
-          - test_arm_baremetal: test_model_smollm2-135M
       fail-fast: false
     with:
       runner: linux.2xlarge.memory

.lintrunner.toml

@@ -503,3 +503,17 @@ command = [
     '--',
     '@{{PATHSFILE}}',
 ]
+
+[[linter]]
+code = 'DOCFORMATTER'
+include_patterns = []
+exclude_patterns = ['**']
+command = [
+  'python','-m','lintrunner_adapters','run','docformatter_linter','--config=pyproject.toml','--','@{{PATHSFILE}}'
+]
+init_command = [
+  'python','-m','lintrunner_adapters','run','pip_init',
+  '--dry-run={{DRYRUN}}',
+  '--requirement=requirements-lintrunner.txt',
+]
+is_formatter = true

backends/aoti/slim/c10/core/ScalarType.h

@@ -133,6 +133,24 @@ inline bool isBoolType(ScalarType t) {
   return t == ScalarType::Bool;
 }
 
+/// Checks if the scalar type is a valid/supported type.
+/// @param t The scalar type to check.
+/// @return true if the scalar type is valid, false otherwise.
+inline bool isValidScalarType(ScalarType t) {
+  switch (t) {
+    case ScalarType::Char:
+    case ScalarType::Short:
+    case ScalarType::Int:
+    case ScalarType::Long:
+    case ScalarType::Float:
+    case ScalarType::Bool:
+    case ScalarType::BFloat16:
+      return true;
+    default:
+      return false;
+  }
+}
+
 inline std::ostream& operator<<(std::ostream& stream, ScalarType scalar_type) {
   return stream << toString(scalar_type);
 }

backends/aoti/slim/core/slim_tensor.h

@@ -52,6 +52,7 @@ class SlimTensor {
         storage_offset_(storage_offset),
         dtype_(dtype) {
     set_sizes_and_strides(sizes, strides);
+    check_supportive();
   }
 
   /**
@@ -65,6 +66,7 @@ class SlimTensor {
         is_contiguous_(true) {
     sizes_and_strides_.set_sizes({0});
     sizes_and_strides_.set_strides({1});
+    check_supportive();
   }
 
   // Default copy/move operations
@@ -556,6 +558,13 @@ class SlimTensor {
         static_cast<int64_t>(numel_));
   }
 
+  void check_supportive() const {
+    ET_CHECK_MSG(
+        c10::isValidScalarType(dtype_),
+        "invalid dtype %d",
+        static_cast<int>(dtype_));
+  }
+
   Storage storage_;
   int64_t storage_offset_{0};
   c10::SizesAndStrides sizes_and_strides_;

backends/aoti/slim/core/storage.h

@@ -240,7 +240,11 @@ class MaybeOwningStorage {
         data_(data),
         capacity_(nbytes),
         deleter_(detail::noop),
-        is_owning_(false) {}
+        is_owning_(false) {
+    if (!device.is_cuda() && !device.is_cpu()) {
+      ET_CHECK_MSG(false, "Unsupported device type: %s", device.str().c_str());
+    }
+  }
 
   /// Default constructor is deleted - storage must have a device.
   MaybeOwningStorage() = delete;

backends/aoti/slim/core/test/test_slimtensor_basic.cpp

@@ -489,4 +489,38 @@ TEST(SlimTensorBasicTest, DataPtrWithOffset) {
   EXPECT_EQ(data, static_cast<char*>(base) + 5 * sizeof(float));
 }
 
+// =============================================================================
+// Dtype and Device Type Validation Tests
+// =============================================================================
+
+TEST(SlimTensorValidationTest, InvalidDtypeUndefined) {
+  std::vector<int64_t> sizes = {2, 3};
+  std::vector<int64_t> strides = {3, 1};
+  size_t nbytes = 6 * sizeof(float);
+  Storage storage = make_cpu_storage(nbytes);
+
+  EXPECT_DEATH(
+      SlimTensor(
+          std::move(storage),
+          makeArrayRef(sizes),
+          makeArrayRef(strides),
+          c10::ScalarType::Undefined),
+      "");
+}
+
+TEST(SlimTensorValidationTest, InvalidDtypeDouble) {
+  std::vector<int64_t> sizes = {2, 3};
+  std::vector<int64_t> strides = {3, 1};
+  size_t nbytes = 6 * sizeof(double);
+  Storage storage = make_cpu_storage(nbytes);
+
+  EXPECT_DEATH(
+      SlimTensor(
+          std::move(storage),
+          makeArrayRef(sizes),
+          makeArrayRef(strides),
+          static_cast<c10::ScalarType>(7)), // Double = 7
+      "");
+}
+
 } // namespace executorch::backends::aoti::slim

backends/aoti/slim/factory/test/test_empty.cpp

@@ -233,6 +233,65 @@ TEST(EmptyTest, CanWriteAndReadData) {
   }
 }
 
+// =============================================================================
+// Dtype and Device Type Validation Tests
+// =============================================================================
+
+TEST(EmptyStridedTest, InvalidDtypeUndefined) {
+  std::vector<int64_t> sizes = {2, 3};
+  std::vector<int64_t> strides = {3, 1};
+
+  EXPECT_DEATH(
+      empty_strided(
+          makeArrayRef(sizes),
+          makeArrayRef(strides),
+          c10::ScalarType::Undefined),
+      "");
+}
+
+TEST(EmptyStridedTest, InvalidDtypeDouble) {
+  std::vector<int64_t> sizes = {2, 3};
+  std::vector<int64_t> strides = {3, 1};
+
+  EXPECT_DEATH(
+      empty_strided(
+          makeArrayRef(sizes),
+          makeArrayRef(strides),
+          static_cast<c10::ScalarType>(7)), // Double = 7
+      "");
+}
+
+TEST(EmptyStridedTest, InvalidDeviceType) {
+  std::vector<int64_t> sizes = {2, 3};
+  std::vector<int64_t> strides = {3, 1};
+
+  c10::Device invalid_device(static_cast<c10::DeviceType>(100), 0);
+
+  EXPECT_DEATH(
+      empty_strided(
+          makeArrayRef(sizes),
+          makeArrayRef(strides),
+          c10::ScalarType::Float,
+          invalid_device),
+      "");
+}
+
+TEST(EmptyTest, InvalidDtypeUndefined) {
+  EXPECT_DEATH(empty({2, 3}, c10::ScalarType::Undefined), "");
+}
+
+TEST(EmptyTest, InvalidDtypeDouble) {
+  EXPECT_DEATH(
+      empty({2, 3}, static_cast<c10::ScalarType>(7)), // Double = 7
+      "");
+}
+
+TEST(EmptyTest, InvalidDeviceType) {
+  c10::Device invalid_device(static_cast<c10::DeviceType>(100), 0);
+
+  EXPECT_DEATH(empty({2, 3}, c10::ScalarType::Float, invalid_device), "");
+}
+
 #ifdef CUDA_AVAILABLE
 
 // =============================================================================

backends/aoti/slim/factory/test/test_from_blob.cpp

@@ -315,6 +315,37 @@ TEST(FromBlobTest, WithArrayRef) {
   EXPECT_TRUE(tensor.is_contiguous());
 }
 
+// =============================================================================
+// Dtype and Device Type Validation Tests
+// =============================================================================
+
+TEST(FromBlobTest, InvalidDtypeUndefined) {
+  constexpr size_t kNumFloats = 6;
+  float external_data[kNumFloats];
+
+  EXPECT_DEATH(
+      from_blob(external_data, {2, 3}, c10::ScalarType::Undefined), "");
+}
+
+TEST(FromBlobTest, InvalidDtypeDouble) {
+  constexpr size_t kNumFloats = 6;
+  float external_data[kNumFloats];
+
+  EXPECT_DEATH(
+      from_blob(external_data, {2, 3}, static_cast<c10::ScalarType>(7)), "");
+}
+
+TEST(FromBlobTest, InvalidDeviceType) {
+  constexpr size_t kNumFloats = 6;
+  float external_data[kNumFloats];
+
+  c10::Device invalid_device(static_cast<c10::DeviceType>(100), 0);
+
+  EXPECT_DEATH(
+      from_blob(external_data, {2, 3}, c10::ScalarType::Float, invalid_device),
+      "");
+}
+
 // =============================================================================
 // CUDA from_blob Tests
 // Tests are skipped at runtime if CUDA hardware is not available.