[Mlir-commits] [clang] [llvm] [mlir] [AArch64][SME] Improve codegen for aarch64.sme.cnts* when not in streaming mode (PR #154761)

[Kerry McLaughlin]

================
@@ -845,15 +847,25 @@ struct StreamingVLOpConversion
     auto loc = streamingVlOp.getLoc();
     auto i64Type = rewriter.getI64Type();
     auto *intrOp = [&]() -> Operation * {
+      auto cntsd = arm_sme::aarch64_sme_cntsd::create(rewriter, loc, i64Type);
       switch (streamingVlOp.getTypeSize()) {
-      case arm_sme::TypeSize::Byte:
-        return arm_sme::aarch64_sme_cntsb::create(rewriter, loc, i64Type);
-      case arm_sme::TypeSize::Half:
-        return arm_sme::aarch64_sme_cntsh::create(rewriter, loc, i64Type);
-      case arm_sme::TypeSize::Word:
-        return arm_sme::aarch64_sme_cntsw::create(rewriter, loc, i64Type);
+      case arm_sme::TypeSize::Byte: {
+        auto mul = arith::ConstantIndexOp::create(rewriter, loc, 8);
+        auto mul64 = arith::IndexCastOp::create(rewriter, loc, i64Type, mul);
+        return arith::MulIOp::create(rewriter, loc, cntsd, mul64);
+      }
+      case arm_sme::TypeSize::Half: {
+        auto mul = arith::ConstantIndexOp::create(rewriter, loc, 4);
+        auto mul64 = arith::IndexCastOp::create(rewriter, loc, i64Type, mul);
+        return arith::MulIOp::create(rewriter, loc, cntsd, mul64);
+      }
+      case arm_sme::TypeSize::Word: {
+        auto mul = arith::ConstantIndexOp::create(rewriter, loc, 2);
+        auto mul64 = arith::IndexCastOp::create(rewriter, loc, i64Type, mul);
+        return arith::MulIOp::create(rewriter, loc, cntsd, mul64);
+      }
       case arm_sme::TypeSize::Double:
-        return arm_sme::aarch64_sme_cntsd::create(rewriter, loc, i64Type);
+        return cntsd;
       }
       llvm_unreachable("unknown type size in StreamingVLOpConversion");
     }();
----------------
kmclaughlin-arm wrote:

I've removed the lambda and added getSizeInBytes in Utils.cpp, as there seemed to be similar helpers in there for SME types.

https://github.com/llvm/llvm-project/pull/154761