[Mlir-commits] [mlir] ae5575d - [mlir][ArmSME] Add `arm_sme.intr.cnts(b|h|w|d)` intrinsics (#77319)

[llvmlistbot at llvm.org]

Author: Benjamin Maxwell
Date: 2024-01-09T09:05:31Z
New Revision: ae5575db1561c0606582346d5f0cbc799c1c02f3

URL: https://github.com/llvm/llvm-project/commit/ae5575db1561c0606582346d5f0cbc799c1c02f3
DIFF: https://github.com/llvm/llvm-project/commit/ae5575db1561c0606582346d5f0cbc799c1c02f3.diff

LOG: [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.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEIntrinsicOps.td
    mlir/test/Target/LLVMIR/arm-sme-invalid.mlir
    mlir/test/Target/LLVMIR/arm-sme.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEIntrinsicOps.td b/mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEIntrinsicOps.td
index 4d96e04c886fa3..d85ef963ae5dc4 100644
--- a/mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEIntrinsicOps.td
+++ b/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

diff  --git a/mlir/test/Target/LLVMIR/arm-sme-invalid.mlir b/mlir/test/Target/LLVMIR/arm-sme-invalid.mlir
index 7c9976bed91273..14821da838726f 100644
--- a/mlir/test/Target/LLVMIR/arm-sme-invalid.mlir
+++ b/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
+}

diff  --git a/mlir/test/Target/LLVMIR/arm-sme.mlir b/mlir/test/Target/LLVMIR/arm-sme.mlir
index edc1f749130440..7a42033dc04bc0 100644
--- a/mlir/test/Target/LLVMIR/arm-sme.mlir
+++ b/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
+}