Replace streaming_vector_bits with Feature::SME_SVL

Reason:
  While vector_bits is used across multiple target architectures,
  streaming_vector_bits is aarch64 specific. So we choose to
  use Target::Feature rather than a new member for arbitrary bits.

- Removed Target::streaming_vector_bits member variable
- Added Feature::SME_SVL{128,256,512,1024,2048}

Author: stevesuzuki-arm

python_bindings/src/halide/halide_/PyEnums.cpp

@@ -188,6 +188,11 @@ void define_enums(py::module &m) {
         .value("SVE", Target::Feature::SVE)
         .value("SVE2", Target::Feature::SVE2)
         .value("SME2", Target::Feature::SME2)
+        .value("SME_SVL128", Target::Feature::SME_SVL128)
+        .value("SME_SVL256", Target::Feature::SME_SVL256)
+        .value("SME_SVL512", Target::Feature::SME_SVL512)
+        .value("SME_SVL1024", Target::Feature::SME_SVL1024)
+        .value("SME_SVL2048", Target::Feature::SME_SVL2048)
         .value("ARMDotProd", Target::Feature::ARMDotProd)
         .value("ARMFp16", Target::Feature::ARMFp16)
         .value("LLVMLargeCodeModel", Target::Feature::LLVMLargeCodeModel)

python_bindings/src/halide/halide_/PyTarget.cpp

@@ -55,9 +55,11 @@ void define_target(py::module &m) {
             .def("supports_device_api", &Target::supports_device_api, py::arg("device"))
             .def("natural_vector_size", natural_vector_size1_method, py::arg("type"))
             .def("natural_vector_size", natural_vector_size2_method, py::arg("type"), py::arg("is_sme_streaming"))
+            .def("sme_streaming_vector_bits", &Target::sme_streaming_vector_bits)
             .def("has_large_buffers", &Target::has_large_buffers)
             .def("maximum_buffer_size", &Target::maximum_buffer_size)
             .def("supported", &Target::supported)
+            .def_static("sme_svl_feature_from_bits", &Target::sme_svl_feature_from_bits, py::arg("bits"))
             .def_static("validate_target_string", &Target::validate_target_string, py::arg("name"));
     ;
 

src/CodeGen_ARM.cpp

@@ -1070,8 +1070,8 @@ void CodeGen_ARM::init_module() {
         user_warning << "Halide does not support SVE for now. Use SVE2 if your target device supports it.\n";
     }
     if (target.has_feature(Target::SME2)) {
-        user_assert(target.streaming_vector_bits != 0) << "For SME2 support, Target::streaming_vector_bits must be set. For generator target strings, add \"streaming_vector_bits_<bits>\".\n";
-        user_assert((target.streaming_vector_bits % 128) == 0) << "For SME2 support, Target::streaming_vector_bits must be a multiple of 128.\n";
+        user_assert(target.sme_streaming_vector_bits() != 0)
+            << "For SME2 support, exactly one Target::SME_SVL* feature must be set. For generator target strings, add \"sme_svl<bits>\".\n";
     }
 
     const bool has_neon = !target.has_feature(Target::NoNEON);
@@ -1141,7 +1141,7 @@ void CodeGen_ARM::init_module() {
                 intrinsics_map = &intrinsics_sve2;
                 break;
             case SIMDFlavors::Streaming:
-                vscale = target.streaming_vector_bits / 128;
+                vscale = target.sme_streaming_vector_bits() / 128;
                 intrinsics_map = &intrinsics_streaming;
                 break;
             default:
@@ -1295,7 +1295,7 @@ void CodeGen_ARM::compile_func(const LoweredFunc &f,
         });
 
         if (is_streaming_task) {
-            feasible_vscale = check_feasible_vscale(target.streaming_vector_bits,  // SVL
+            feasible_vscale = check_feasible_vscale(target.sme_streaming_vector_bits(),  // SVL
                                                     lanes_used, "streaming_", simple_name);
         }
         in_streaming = (feasible_vscale > 0) && is_streaming_task;
@@ -1344,7 +1344,7 @@ void CodeGen_ARM::compile_func(const LoweredFunc &f,
             // We check regardless of streaming mode enabled or not
             // because streaming task is basically internal linkage.
             Expr runtime_vscale = Call::make(Int(32), Call::get_runtime_streaming_vscale, {}, Call::PureIntrinsic);
-            Expr compiletime_vscale = Expr(target.streaming_vector_bits / 128);
+            Expr compiletime_vscale = Expr(target.sme_streaming_vector_bits() / 128);
             std::vector<Expr> args{simple_name, std::string("streaming"), runtime_vscale, compiletime_vscale};
             Expr error = Call::make(Int(32), "halide_error_vscale_invalid", args, Call::Extern);
             func.body = Block::make(AssertStmt::make(runtime_vscale == compiletime_vscale, error), func.body);

src/Target.cpp

@@ -274,7 +274,6 @@ Target calculate_host_target() {
     bool use_64_bits = (sizeof(size_t) == 8);
     int bits = use_64_bits ? 64 : 32;
     int vector_bits = 0;
-    int streaming_vector_bits = 0;
     Target::Processor processor = Target::Processor::ProcessorGeneric;
     std::vector<Target::Feature> initial_features;
 
@@ -386,7 +385,11 @@ Target calculate_host_target() {
         vector_bits = get_sve_vector_length();
     }
     if (has_streaming_scalable_vector) {
-        streaming_vector_bits = get_sme_streaming_vector_length();
+        const int streaming_vector_bits = get_sme_streaming_vector_length();
+        Target::Feature sme_svl = Target::sme_svl_feature_from_bits(streaming_vector_bits);
+        user_assert(sme_svl != Target::FeatureEnd)
+            << "Detected unsupported SME streaming vector length " << streaming_vector_bits << " bits.\n";
+        initial_features.push_back(sme_svl);
     }
 #endif
 
@@ -461,7 +464,7 @@ Target calculate_host_target() {
         processor = get_amd_processor(family, model, have_sse3);
 
         if (processor == Target::Processor::ZnVer4) {
-            Target t{os, arch, bits, processor, initial_features, vector_bits, streaming_vector_bits};
+            Target t{os, arch, bits, processor, initial_features, vector_bits};
             t.set_feature(Target::SSE41);
             if (os_avx) {
                 t.set_features({Target::AVX, Target::F16C, Target::FMA, Target::AVX2});
@@ -472,7 +475,7 @@ Target calculate_host_target() {
             }
             return t;
         } else if (processor == Target::Processor::ZnVer5) {
-            Target t{os, arch, bits, processor, initial_features, vector_bits, streaming_vector_bits};
+            Target t{os, arch, bits, processor, initial_features, vector_bits};
             t.set_feature(Target::SSE41);
             if (os_avx) {
                 t.set_features({Target::AVX, Target::F16C, Target::FMA,
@@ -589,7 +592,7 @@ Target calculate_host_target() {
 #endif
 #endif
 
-    return {os, arch, bits, processor, initial_features, vector_bits, streaming_vector_bits};
+    return {os, arch, bits, processor, initial_features, vector_bits};
 }
 
 bool is_using_hexagon(const Target &t) {
@@ -842,6 +845,11 @@ const std::map<std::string, Target::Feature> feature_name_map = {
     {"sve", Target::SVE},
     {"sve2", Target::SVE2},
     {"sme2", Target::SME2},
+    {"sme_svl128", Target::SME_SVL128},
+    {"sme_svl256", Target::SME_SVL256},
+    {"sme_svl512", Target::SME_SVL512},
+    {"sme_svl1024", Target::SME_SVL1024},
+    {"sme_svl2048", Target::SME_SVL2048},
     {"arm_dot_prod", Target::ARMDotProd},
     {"arm_fp16", Target::ARMFp16},
     {"llvm_large_code_model", Target::LLVMLargeCodeModel},
@@ -1000,8 +1008,6 @@ bool merge_string(Target &t, const std::string &target) {
             features_specified = true;
         } else if (auto vb = parse_vector_bits(tok, "vector_bits_"); vb >= 0) {
             t.vector_bits = vb;
-        } else if (auto svb = parse_vector_bits(tok, "streaming_vector_bits_"); svb >= 0) {
-            t.streaming_vector_bits = svb;
         } else {
             return false;
         }
@@ -1123,6 +1129,11 @@ void Target::validate_features() const {
                                 POWER_ARCH_2_07,
                                 RVV,
                                 SME2,
+                                SME_SVL128,
+                                SME_SVL256,
+                                SME_SVL512,
+                                SME_SVL1024,
+                                SME_SVL2048,
                                 SVE,
                                 SVE2,
                                 VSX,
@@ -1184,12 +1195,31 @@ void Target::validate_features() const {
                                 POWER_ARCH_2_07,
                                 RVV,
                                 SSE41,
+                                SME_SVL128,
+                                SME_SVL256,
+                                SME_SVL512,
+                                SME_SVL1024,
+                                SME_SVL2048,
                                 SME2,
                                 SVE,
                                 SVE2,
                                 VSX,
                             });
     }
+
+    const int num_sme_svl_features =
+        (int)has_feature(SME_SVL128) +
+        (int)has_feature(SME_SVL256) +
+        (int)has_feature(SME_SVL512) +
+        (int)has_feature(SME_SVL1024) +
+        (int)has_feature(SME_SVL2048);
+
+    user_assert(num_sme_svl_features <= 1)
+        << "Target may have at most one SME_SVL feature.\n";
+    user_assert(!has_feature(SME2) || num_sme_svl_features == 1)
+        << "Target feature sme2 requires exactly one SME_SVL feature.\n";
+    user_assert(has_feature(SME2) || num_sme_svl_features == 0)
+        << "Target features SME_SVL128, SME_SVL256, SME_SVL512, SME_SVL1024, and SME_SVL2048 require target feature sme2.\n";
 }
 
 Target::Target(const std::string &target) {
@@ -1233,6 +1263,23 @@ Target::Feature Target::feature_from_name(const std::string &name) {
     return Target::FeatureEnd;
 }
 
+Target::Feature Target::sme_svl_feature_from_bits(int bits) {
+    switch (bits) {
+    case 128:
+        return Target::SME_SVL128;
+    case 256:
+        return Target::SME_SVL256;
+    case 512:
+        return Target::SME_SVL512;
+    case 1024:
+        return Target::SME_SVL1024;
+    case 2048:
+        return Target::SME_SVL2048;
+    default:
+        return Target::FeatureEnd;
+    }
+}
+
 std::string Target::to_string() const {
     string result;
     for (const auto &arch_entry : arch_name_map) {
@@ -1269,9 +1316,6 @@ std::string Target::to_string() const {
     if (vector_bits != 0) {
         result += "-vector_bits_" + std::to_string(vector_bits);
     }
-    if (streaming_vector_bits != 0) {
-        result += "-streaming_vector_bits_" + std::to_string(streaming_vector_bits);
-    }
 
     return result;
 }
@@ -1600,6 +1644,31 @@ int Target::natural_vector_size(const Halide::Type &t) const {
     return natural_vector_size(t, false);
 }
 
+int Target::sme_streaming_vector_bits() const {
+    int result = 0;
+    auto set_result = [&result](int bits) {
+        user_assert(result == 0)
+            << "Target may have at most one SME_SVL feature.\n";
+        result = bits;
+    };
+    if (has_feature(Target::SME_SVL128)) {
+        set_result(128);
+    }
+    if (has_feature(Target::SME_SVL256)) {
+        set_result(256);
+    }
+    if (has_feature(Target::SME_SVL512)) {
+        set_result(512);
+    }
+    if (has_feature(Target::SME_SVL1024)) {
+        set_result(1024);
+    }
+    if (has_feature(Target::SME_SVL2048)) {
+        set_result(2048);
+    }
+    return result;
+}
+
 int Target::natural_vector_size(const Halide::Type &t, bool is_sme_streaming) const {
     user_assert(!has_unknowns())
         << "natural_vector_size cannot be used on a Target with Unknown values.\n";
@@ -1609,9 +1678,9 @@ int Target::natural_vector_size(const Halide::Type &t, bool is_sme_streaming) co
 
     if (arch == Target::ARM) {
         if (is_sme_streaming &&
-            streaming_vector_bits != 0 &&
+            sme_streaming_vector_bits() != 0 &&
             has_feature(Halide::Target::SME2)) {
-            return streaming_vector_bits / (data_size * 8);
+            return sme_streaming_vector_bits() / (data_size * 8);
         } else if (vector_bits != 0 &&
                    (has_feature(Halide::Target::SVE2) ||
                     (t.is_float() && has_feature(Halide::Target::SVE)))) {
@@ -1965,14 +2034,13 @@ void target_test() {
     internal_assert(Target(with_vector_bits.to_string()).vector_bits == 512) << "Vector bits not round tripped properly.\n";
     internal_assert(with_vector_bits.natural_vector_size(Int(32)) == 16) << "Wrong natural_vector_size.\n";
 
-    // Tests for streaming_vector_bits
-    internal_assert(Target().streaming_vector_bits == 0) << "Default Target streaming_vector_bits not 0.\n";
-    internal_assert(Target("arm-64-linux-sme2-vector_bits_512-streaming_vector_bits_1024").streaming_vector_bits == 1024) << "Streaming vector bits not parsed correctly.\n";
-    Target with_streaming_vector_bits(Target::Linux, Target::ARM, 64, Target::ProcessorGeneric, {Target::SVE2, Target::SME2}, 512, 1024);
-    internal_assert(with_streaming_vector_bits.streaming_vector_bits == 1024) << "Streaming vector bits not populated in constructor.\n";
-    internal_assert(Target(with_streaming_vector_bits.to_string()).streaming_vector_bits == 1024) << "Streaming vector bits not round tripped properly.\n";
-    internal_assert(with_streaming_vector_bits.natural_vector_size(Int(32), true) == 32) << "Wrong natural_vector_size with SME streaming.\n";
-    internal_assert(with_streaming_vector_bits.natural_vector_size(Int(32), false) == 16) << "Wrong natural_vector_size without SME streaming.\n";
+    // Tests for SME streaming vector length
+    internal_assert(Target().sme_streaming_vector_bits() == 0) << "Default Target SME SVL not 0.\n";
+    internal_assert(Target::sme_svl_feature_from_bits(1024) == Target::SME_SVL1024) << "SME SVL feature lookup failed.\n";
+    Target with_sme_svl(Target::Linux, Target::ARM, 64, Target::ProcessorGeneric, {Target::SVE2, Target::SME2, Target::SME_SVL1024}, 512);
+    internal_assert(with_sme_svl.sme_streaming_vector_bits() == 1024) << "SME SVL not populated in constructor.\n";
+    internal_assert(with_sme_svl.natural_vector_size(Int(32), true) == 32) << "Wrong natural_vector_size with SME streaming.\n";
+    internal_assert(with_sme_svl.natural_vector_size(Int(32), false) == 16) << "Wrong natural_vector_size without SME streaming.\n";
 
     std::cout << "Target test passed\n";
 }

src/Target.h

@@ -51,10 +51,8 @@ struct Target {
 
     /** The bit-width of a vector register for targets where this is configurable and
      * targeting a fixed size is desired. The default of 0 indicates no assumption of
-     * fixed size is allowed.
-     * streaming_vector_bits is for streaming mode in aarch64 SME. */
+     * fixed size is allowed. */
     int vector_bits = 0;
-    int streaming_vector_bits = 0;
 
     /** The specific processor to be targeted, tuned for.
      * Corresponds to processor_name_map in Target.cpp.
@@ -157,6 +155,11 @@ struct Target {
         SVE = halide_target_feature_sve,
         SVE2 = halide_target_feature_sve2,
         SME2 = halide_target_feature_sme2,
+        SME_SVL128 = halide_target_feature_sme_svl128,
+        SME_SVL256 = halide_target_feature_sme_svl256,
+        SME_SVL512 = halide_target_feature_sme_svl512,
+        SME_SVL1024 = halide_target_feature_sme_svl1024,
+        SME_SVL2048 = halide_target_feature_sme_svl2048,
         ARMDotProd = halide_target_feature_arm_dot_prod,
         ARMFp16 = halide_target_feature_arm_fp16,
         LLVMLargeCodeModel = halide_llvm_large_code_model,
@@ -191,8 +194,8 @@ struct Target {
     };
     Target() = default;
     Target(OS o, Arch a, int b, Processor pt, const std::vector<Feature> &initial_features = std::vector<Feature>(),
-           int vb = 0, int svb = 0)
-        : os(o), arch(a), bits(b), vector_bits(vb), streaming_vector_bits(svb), processor_tune(pt) {
+           int vb = 0)
+        : os(o), arch(a), bits(b), vector_bits(vb), processor_tune(pt) {
         for (const auto &f : initial_features) {
             set_feature(f);
         }
@@ -285,8 +288,7 @@ struct Target {
                bits == other.bits &&
                processor_tune == other.processor_tune &&
                features == other.features &&
-               vector_bits == other.vector_bits &&
-               streaming_vector_bits == other.streaming_vector_bits;
+               vector_bits == other.vector_bits;
     }
 
     bool operator!=(const Target &other) const {
@@ -325,6 +327,10 @@ struct Target {
      * for that data type in streaming mode in aarch64 SME. */
     int natural_vector_size(const Halide::Type &t, bool is_sme_streaming) const;
 
+    /** Return the fixed SME streaming vector length in bits selected by this target,
+     * or 0 if no SME_SVL feature is set. */
+    int sme_streaming_vector_bits() const;
+
     /** Given a data type, return an estimate of the "natural" vector size
      * for that data type when compiling for this Target. */
     template<typename data_t>
@@ -381,6 +387,10 @@ struct Target {
      * If the string is not a known feature name, return FeatureEnd. */
     static Target::Feature feature_from_name(const std::string &name);
 
+    /** Return the SME_SVL feature corresponding to an SME streaming vector
+     * length in bits, or FeatureEnd if no exact SME_SVL feature exists. */
+    static Target::Feature sme_svl_feature_from_bits(int bits);
+
 private:
     /** A bitmask that stores the active features. */
     std::bitset<FeatureEnd> features;

src/runtime/HalideRuntime.h

@@ -1448,6 +1448,11 @@ typedef enum halide_target_feature_t {
     halide_target_feature_sve,                    ///< Enable ARM Scalable Vector Extensions
     halide_target_feature_sve2,                   ///< Enable ARM Scalable Vector Extensions v2
     halide_target_feature_sme2,                   ///< Enable ARM Scalable Matrix Extensions v2
+    halide_target_feature_sme_svl128,             ///< Assume ARM SME streaming vector length is 128 bits.
+    halide_target_feature_sme_svl256,             ///< Assume ARM SME streaming vector length is 256 bits.
+    halide_target_feature_sme_svl512,             ///< Assume ARM SME streaming vector length is 512 bits.
+    halide_target_feature_sme_svl1024,            ///< Assume ARM SME streaming vector length is 1024 bits.
+    halide_target_feature_sme_svl2048,            ///< Assume ARM SME streaming vector length is 2048 bits.
     halide_target_feature_egl,                    ///< Force use of EGL support.
     halide_target_feature_arm_dot_prod,           ///< Enable ARMv8.2-a dotprod extension (i.e. udot and sdot instructions)
     halide_target_feature_arm_fp16,               ///< Enable ARMv8.2-a half-precision floating point data processing

src/runtime/aarch64_cpu_features.cpp

@@ -141,6 +141,11 @@ extern "C" WEAK int halide_get_cpu_features(CpuFeatures *features) {
     halide_set_known_cpu_feature(features, halide_target_feature_sve);
     halide_set_known_cpu_feature(features, halide_target_feature_sve2);
     halide_set_known_cpu_feature(features, halide_target_feature_sme2);
+    halide_set_known_cpu_feature(features, halide_target_feature_sme_svl128);
+    halide_set_known_cpu_feature(features, halide_target_feature_sme_svl256);
+    halide_set_known_cpu_feature(features, halide_target_feature_sme_svl512);
+    halide_set_known_cpu_feature(features, halide_target_feature_sme_svl1024);
+    halide_set_known_cpu_feature(features, halide_target_feature_sme_svl2048);
 
     // All ARM architectures support "No Neon".
     halide_set_available_cpu_feature(features, halide_target_feature_no_neon);

test/correctness/fallback_vscale_sve.cpp

@@ -96,7 +96,12 @@ bool test_streaming_vscale(int vectorization_factor, int vector_bits, int stream
     }
     if (streaming_vector_bits != 0) {
         t = t.with_feature(Target::SME2);
-        t.streaming_vector_bits = streaming_vector_bits;
+        Target::Feature sme_svl = Target::sme_svl_feature_from_bits(streaming_vector_bits);
+        if (sme_svl == Target::FeatureEnd) {
+            printf("[%s] Unsupported streaming_vector_bits %d\n", name.c_str(), streaming_vector_bits);
+            return false;
+        }
+        t.set_feature(sme_svl);
     }
 
     // sve or neon