- Remove lambda from StreamingVLOpConversion

- Add getSizeInBytes helper

Author: kmclaughlin-arm

clang/lib/CodeGen/TargetBuiltins/ARM.cpp

@@ -4919,25 +4919,20 @@ Value *CodeGenFunction::EmitAArch64SMEBuiltinExpr(unsigned BuiltinID,
   // Handle builtins which require their multi-vector operands to be swapped
   swapCommutativeSMEOperands(BuiltinID, Ops);
 
-  auto isCntsBuiltin = [&](int64_t &Mul) {
+  auto isCntsBuiltin = [&]() {
     switch (BuiltinID) {
     default:
-      Mul = 0;
-      return false;
+      return 0;
     case SME::BI__builtin_sme_svcntsb:
-      Mul = 8;
-      return true;
+      return 8;
     case SME::BI__builtin_sme_svcntsh:
-      Mul = 4;
-      return true;
+      return 4;
     case SME::BI__builtin_sme_svcntsw:
-      Mul = 2;
-      return true;
+      return 2;
     }
   };
 
-  int64_t Mul = 0;
-  if (isCntsBuiltin(Mul)) {
+  if (auto Mul = isCntsBuiltin()) {
     llvm::Value *Cntd =
         Builder.CreateCall(CGM.getIntrinsic(Intrinsic::aarch64_sme_cntsd));
     return Builder.CreateMul(Cntd, llvm::ConstantInt::get(Int64Ty, Mul),

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

@@ -940,6 +940,9 @@ bool AArch64DAGToDAGISel::SelectRDVLImm(SDValue N, SDValue &Imm) {
   return false;
 }
 
+// Given cntsd = (rdsvl, #1) >> 3, attempt to return a suitable multiplier
+// for RDSVL to calculate the streaming vector length in bytes * N. i.e.
+//   rdsvl, #(ShlImm - 3)
 template <signed Low, signed High>
 bool AArch64DAGToDAGISel::SelectRDSVLShiftImm(SDValue N, SDValue &Imm) {
   if (!isa<ConstantSDNode>(N))

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

@@ -2102,8 +2102,8 @@ instCombineSVECntElts(InstCombiner &IC, IntrinsicInst &II, unsigned NumElts) {
 }
 
 static std::optional<Instruction *>
-instCombineSMECntsElts(InstCombiner &IC, IntrinsicInst &II,
-                       const AArch64Subtarget *ST) {
+instCombineSMECntsd(InstCombiner &IC, IntrinsicInst &II,
+                    const AArch64Subtarget *ST) {
   if (!ST->isStreaming())
     return std::nullopt;
 
@@ -2825,7 +2825,7 @@ AArch64TTIImpl::instCombineIntrinsic(InstCombiner &IC,
   case Intrinsic::aarch64_sve_cntb:
     return instCombineSVECntElts(IC, II, 16);
   case Intrinsic::aarch64_sme_cntsd:
-    return instCombineSMECntsElts(IC, II, ST);
+    return instCombineSMECntsd(IC, II, ST);
   case Intrinsic::aarch64_sve_ptest_any:
   case Intrinsic::aarch64_sve_ptest_first:
   case Intrinsic::aarch64_sve_ptest_last:

mlir/include/mlir/Dialect/ArmSME/Utils/Utils.h

@@ -32,6 +32,9 @@ namespace mlir::arm_sme {
 
 constexpr unsigned MinStreamingVectorLengthInBits = 128;
 
+/// Return the size represented by arm_sme::TypeSize in bytes.
+unsigned getSizeInBytes(TypeSize type);
+
 /// Return minimum number of elements for the given element `type` in
 /// a vector of SVL bits.
 unsigned getSMETileSliceMinNumElts(Type type);

mlir/lib/Conversion/ArmSMEToLLVM/ArmSMEToLLVM.cpp

@@ -846,31 +846,13 @@ struct StreamingVLOpConversion
                   ConversionPatternRewriter &rewriter) const override {
     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: {
-        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 cntsd;
-      }
-      llvm_unreachable("unknown type size in StreamingVLOpConversion");
-    }();
-    rewriter.replaceOpWithNewOp<arith::IndexCastOp>(
-        streamingVlOp, rewriter.getIndexType(), intrOp->getResult(0));
+    auto cntsd = arm_sme::aarch64_sme_cntsd::create(rewriter, loc, i64Type);
+    auto cntsdIdx = arith::IndexCastOp::create(rewriter, loc,
+                                               rewriter.getIndexType(), cntsd);
+    auto scale = arith::ConstantIndexOp::create(
+        rewriter, loc,
+        8 / arm_sme::getSizeInBytes(streamingVlOp.getTypeSize()));
+    rewriter.replaceOpWithNewOp<arith::MulIOp>(streamingVlOp, cntsdIdx, scale);
     return success();
   }
 };

mlir/lib/Dialect/ArmSME/IR/Utils.cpp

@@ -14,6 +14,21 @@
 
 namespace mlir::arm_sme {
 
+unsigned getSizeInBytes(TypeSize type) {
+  switch (type) {
+  case arm_sme::TypeSize::Byte:
+    return 1;
+  case arm_sme::TypeSize::Half:
+    return 2;
+  case arm_sme::TypeSize::Word:
+    return 4;
+  case arm_sme::TypeSize::Double:
+    return 8;
+  default:
+    llvm_unreachable("unknown type size");
+  }
+}
+
 unsigned getSMETileSliceMinNumElts(Type type) {
   assert(isValidSMETileElementType(type) && "invalid tile type!");
   return MinStreamingVectorLengthInBits / type.getIntOrFloatBitWidth();

mlir/test/Conversion/ArmSMEToLLVM/arm-sme-to-llvm.mlir

@@ -586,10 +586,10 @@ func.func @arm_sme_extract_tile_slice_ver_i128(%tile_slice_index : index) -> vec
 // -----
 
 // CHECK-LABEL: @arm_sme_streaming_vl_bytes
-// CHECK: %[[CONST:.*]] = arith.constant 8 : i64
+// CHECK: %[[CONST:.*]] = arith.constant 8 : index
 // CHECK: %[[CNTSD:.*]] = "arm_sme.intr.cntsd"() : () -> i64
-// CHECK: %[[MUL:.*]] = arith.muli %[[CNTSD]], %[[CONST]] : i64
-// CHECK: %[[INDEX_COUNT:.*]] = arith.index_cast %[[MUL]] : i64 to index
+// CHECK: %[[CNTSD_IDX:.*]] = arith.index_cast %[[CNTSD]] : i64 to index
+// CHECK: %[[MUL:.*]] = arith.muli %[[CNTSD_IDX]], %[[CONST]] : index
 func.func @arm_sme_streaming_vl_bytes() -> index {
   %svl_b = arm_sme.streaming_vl <byte>
   return %svl_b : index
@@ -598,10 +598,10 @@ func.func @arm_sme_streaming_vl_bytes() -> index {
 // -----
 
 // CHECK-LABEL: @arm_sme_streaming_vl_half_words
-// CHECK: %[[CONST:.*]] = arith.constant 4 : i64
+// CHECK: %[[CONST:.*]] = arith.constant 4 : index
 // CHECK: %[[CNTSD:.*]] = "arm_sme.intr.cntsd"() : () -> i64
-// CHECK: %[[MUL:.*]] = arith.muli %[[CNTSD]], %[[CONST]] : i64
-// CHECK: %[[INDEX_COUNT:.*]] = arith.index_cast %[[MUL]] : i64 to index
+// CHECK: %[[CNTSD_IDX:.*]] = arith.index_cast %[[CNTSD]] : i64 to index
+// CHECK: %[[MUL:.*]] = arith.muli %[[CNTSD_IDX]], %[[CONST]] : index
 func.func @arm_sme_streaming_vl_half_words() -> index {
   %svl_h = arm_sme.streaming_vl <half>
   return %svl_h : index
@@ -610,10 +610,10 @@ func.func @arm_sme_streaming_vl_half_words() -> index {
 // -----
 
 // CHECK-LABEL: @arm_sme_streaming_vl_words
-// CHECK: %[[CONST:.*]] = arith.constant 2 : i64
+// CHECK: %[[CONST:.*]] = arith.constant 2 : index
 // CHECK: %[[CNTSD:.*]] = "arm_sme.intr.cntsd"() : () -> i64
-// CHECK: %[[MUL:.*]] = arith.muli %[[CNTSD]], %[[CONST]] : i64
-// CHECK: %[[INDEX_COUNT:.*]] = arith.index_cast %[[MUL]] : i64 to index
+// CHECK: %[[CNTSD_IDX:.*]] = arith.index_cast %[[CNTSD]] : i64 to index
+// CHECK: %[[MUL:.*]] = arith.muli %[[CNTSD_IDX]], %[[CONST]] : index
 func.func @arm_sme_streaming_vl_words() -> index {
   %svl_w = arm_sme.streaming_vl <word>
   return %svl_w : index