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

[Benjamin Maxwell]

https://github.com/MacDue created https://github.com/llvm/llvm-project/pull/77319

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

>From e8022d98cab11176ae8b20ae5ca246706923b7c1 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Mon, 8 Jan 2024 14:04:09 +0000
Subject: [PATCH] [mlir][ArmSME] Add `arm_sme.intr.cnts(b|h|w|d)` intrinsics

This adds MLIR versions of Arm streaming vector length intrinsics. These
allow reading the streaming vector length regardless of the streaming
mode.
---
 .../mlir/Dialect/ArmSME/IR/ArmSMEIntrinsicOps.td   | 13 +++++++++++++
 mlir/test/Target/LLVMIR/arm-sme-invalid.mlir       |  8 ++++++++
 mlir/test/Target/LLVMIR/arm-sme.mlir               | 14 ++++++++++++++
 3 files changed, 35 insertions(+)

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
+}