[mlir][ArmSME] Add arm_sme.intr.cnts(b|h|w|d) intrinsics (#77319)

This adds MLIR versions of the Arm streaming vector length intrinsics.
These allow reading the streaming vector length regardless of the
streaming mode.

Author: MacDue

mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEIntrinsicOps.td

@@ -187,4 +187,17 @@ def LLVM_aarch64_sme_write_vert : LLVM_aarch64_sme_write<"vert">;
 def LLVM_aarch64_sme_read_horiz : LLVM_aarch64_sme_read<"horiz">;
 def LLVM_aarch64_sme_read_vert : LLVM_aarch64_sme_read<"vert">;
 
+class ArmSME_IntrCountOp<string mnemonic>
+    : ArmSME_IntrOp<mnemonic,
+                    /*immArgPositions=*/[],
+                    /*immArgAttrNames=*/[],
+                    /*overloadedOperands=*/[],
+                    /*traits*/[PredOpTrait<"`res` is i64", TypeIsPred<"res", I64>>],
+                    /*numResults=*/1, /*overloadedResults=*/[]>;
+
+def LLVM_aarch64_sme_cntsb : ArmSME_IntrCountOp<"cntsb">;
+def LLVM_aarch64_sme_cntsh : ArmSME_IntrCountOp<"cntsh">;
+def LLVM_aarch64_sme_cntsw : ArmSME_IntrCountOp<"cntsw">;
+def LLVM_aarch64_sme_cntsd : ArmSME_IntrCountOp<"cntsd">;
+
 #endif // ARMSME_INTRINSIC_OPS

mlir/test/Target/LLVMIR/arm-sme-invalid.mlir

@@ -31,3 +31,11 @@ llvm.func @arm_sme_tile_slice_to_vector_invalid_element_types(
       (vector<[4]xf32>, vector<[4]xi1>, i32) -> vector<[4]xi32>
   llvm.return %res : vector<[4]xi32>
 }
+
+// -----
+
+llvm.func @arm_sme_streaming_vl_invalid_return_type() -> i32 {
+  // expected-error @+1 {{failed to verify that `res` is i64}}
+  %res = "arm_sme.intr.cntsb"() : () -> i32
+  llvm.return %res : i32
+}

mlir/test/Target/LLVMIR/arm-sme.mlir

@@ -403,3 +403,17 @@ llvm.func @arm_sme_tile_slice_to_vector_vert(%tileslice : i32,
     : (vector<[2]xf64>, vector<[2]xi1>, i32) -> vector<[2]xf64>
   llvm.return
 }
+
+// -----
+
+llvm.func @arm_sme_streaming_vl() {
+  // CHECK: call i64 @llvm.aarch64.sme.cntsb()
+  %svl_b = "arm_sme.intr.cntsb"() : () -> i64
+  // CHECK: call i64 @llvm.aarch64.sme.cntsh()
+  %svl_h = "arm_sme.intr.cntsh"() : () -> i64
+  // CHECK: call i64 @llvm.aarch64.sme.cntsw()
+  %svl_w = "arm_sme.intr.cntsw"() : () -> i64
+  // CHECK: call i64 @llvm.aarch64.sme.cntsd()
+  %svl_d = "arm_sme.intr.cntsd"() : () -> i64
+  llvm.return
+}