[AArch64][SME] Sink tile offset operands into the loop for load/store instructions.

[sdesmalen]

Summary

This helps ISel decompose the generic offset for the tile into a base + offset.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12660–12661

This may need to protect against vector types now.

[sdesmalen]

Added isa<VectorType> check for fma.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

12660–12661

Good catch @dmgreen, thanks!

[dmgreen]

Thanks. LGTM

You may find that in certain situations the add gets transformed into an or if the operands have no bits are common, but I presume that would not break anything just be a little suboptimal. There is a haveNoCommonBitsSet method if it would be useful to support. I can imagine cases where base would known to be a multiple of a power2.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static bool isSplatShuffle(Value *V) {
   return false;
 }
 
 /// Check if sinking \p I's operands to I's basic block is profitable, because
 /// the operands can be folded into a target instruction, e.g.
 /// shufflevectors extracts and/or sext/zext can be folded into (u,s)subl(2).
 bool AArch64TargetLowering::shouldSinkOperands(
     Instruction *I, SmallVectorImpl<Use *> &Ops) const {
-  if (!I->getType()->isVectorTy())
-    return false;
-
   if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
     switch (II->getIntrinsicID()) {
     case Intrinsic::aarch64_neon_smull:
     case Intrinsic::aarch64_neon_umull:
       if (areExtractShuffleVectors(II->getOperand(0), II->getOperand(1))) {
         Ops.push_back(&II->getOperandUse(0));
         Ops.push_back(&II->getOperandUse(1));
         return true;
       }
       LLVM_FALLTHROUGH;
 
     case Intrinsic::fma:
-      if (cast<VectorType>(I->getType())->getElementType()->isHalfTy() &&
+      if (isa<VectorType>(I->getType()) &&
+          cast<VectorType>(I->getType())->getElementType()->isHalfTy() &&
           !Subtarget->hasFullFP16())
         return false;
       LLVM_FALLTHROUGH;
     case Intrinsic::aarch64_neon_sqdmull:
     case Intrinsic::aarch64_neon_sqdmulh:
     case Intrinsic::aarch64_neon_sqrdmulh:
       // Sink splats for index lane variants
       if (isSplatShuffle(II->getOperand(0)))
         Ops.push_back(&II->getOperandUse(0));
       if (isSplatShuffle(II->getOperand(1)))
         Ops.push_back(&II->getOperandUse(1));
       return !Ops.empty();
-
+    case Intrinsic::aarch64_sme_write_horiz:
+    case Intrinsic::aarch64_sme_write_vert:
+    case Intrinsic::aarch64_sme_writeq_horiz:
+    case Intrinsic::aarch64_sme_writeq_vert: {
+      auto *Idx = dyn_cast<Instruction>(II->getOperand(1));
+      if (!Idx || Idx->getOpcode() != Instruction::Add)
+        return false;
+      Ops.push_back(&II->getOperandUse(1));
+      return true;
+    }
+    case Intrinsic::aarch64_sme_read_horiz:
+    case Intrinsic::aarch64_sme_read_vert:
+    case Intrinsic::aarch64_sme_readq_horiz:
+    case Intrinsic::aarch64_sme_readq_vert:
+    case Intrinsic::aarch64_sme_ld1b_vert:
+    case Intrinsic::aarch64_sme_ld1h_vert:
+    case Intrinsic::aarch64_sme_ld1w_vert:
+    case Intrinsic::aarch64_sme_ld1d_vert:
+    case Intrinsic::aarch64_sme_ld1q_vert:
+    case Intrinsic::aarch64_sme_st1b_vert:
+    case Intrinsic::aarch64_sme_st1h_vert:
+    case Intrinsic::aarch64_sme_st1w_vert:
+    case Intrinsic::aarch64_sme_st1d_vert:
+    case Intrinsic::aarch64_sme_st1q_vert:
+    case Intrinsic::aarch64_sme_ld1b_horiz:
+    case Intrinsic::aarch64_sme_ld1h_horiz:
+    case Intrinsic::aarch64_sme_ld1w_horiz:
+    case Intrinsic::aarch64_sme_ld1d_horiz:
+    case Intrinsic::aarch64_sme_ld1q_horiz:
+    case Intrinsic::aarch64_sme_st1b_horiz:
+    case Intrinsic::aarch64_sme_st1h_horiz:
+    case Intrinsic::aarch64_sme_st1w_horiz:
+    case Intrinsic::aarch64_sme_st1d_horiz:
+    case Intrinsic::aarch64_sme_st1q_horiz: {
+      auto *Idx = dyn_cast<Instruction>(II->getOperand(3));
+      if (!Idx || Idx->getOpcode() != Instruction::Add)
+        return false;
+      Ops.push_back(&II->getOperandUse(3));
+      return true;
+    }
     case Intrinsic::aarch64_neon_pmull:
       if (!areExtractShuffleVectors(II->getOperand(0), II->getOperand(1)))
         return false;
       Ops.push_back(&II->getOperandUse(0));
       Ops.push_back(&II->getOperandUse(1));
       return true;
     case Intrinsic::aarch64_neon_pmull64:
       if (!areOperandsOfVmullHighP64(II->getArgOperand(0),
                                      II->getArgOperand(1)))
         return false;
       Ops.push_back(&II->getArgOperandUse(0));
       Ops.push_back(&II->getArgOperandUse(1));
       return true;
     default:
       return false;
     }
   }
 
+  if (!I->getType()->isVectorTy())
+    return false;
+
   switch (I->getOpcode()) {
   case Instruction::Sub:
   case Instruction::Add: {
     if (!areExtractExts(I->getOperand(0), I->getOperand(1)))
       return false;
 
     // If the exts' operands extract either the lower or upper elements, we
     // can sink them too.

llvm/test/CodeGen/AArch64/sme-intrinsics-loads.ll

 ; CHECK-NEXT:    ret
   %vscale = call i64 @llvm.vscale.i64()
   %mulvl = mul i64 %vscale, 256
   %base = getelementptr i8, ptr %ptr, i64 %mulvl
   call void @llvm.aarch64.sme.ldr(i32 0, ptr %base)
   ret void;
 }
 
+; Ensure that the tile offset is sunk, given that this is likely to be an 'add'
+; that's decomposed into a base + offset in ISel.
+define void @test_ld1_sink_tile0_offset_operand(<vscale x 16 x i1> %pg, ptr %src, i32 %base, i32 %N) {
+; CHECK-LABEL: test_ld1_sink_tile0_offset_operand:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov w12, w1
+; CHECK-NEXT:  .LBB14_1: // %for.body
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    ld1w {za0h.s[w12, 0]}, p0/z, [x0]
+; CHECK-NEXT:    subs w2, w2, #1
+; CHECK-NEXT:    ld1w {za0h.s[w12, 1]}, p0/z, [x0]
+; CHECK-NEXT:    ld1w {za0h.s[w12, 2]}, p0/z, [x0]
+; CHECK-NEXT:    b.ne .LBB14_1
+; CHECK-NEXT:  // %bb.2: // %exit
+; CHECK-NEXT:    ret
+entry:
+  %add1 = add i32 %base, 1
+  %add2 = add i32 %base, 2
+  br label %for.body
+
+for.body:
+  %i = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  call void @llvm.aarch64.sme.ld1w.horiz(<vscale x 16 x i1> %pg, ptr %src, i64 0, i32 %base)
+  call void @llvm.aarch64.sme.ld1w.horiz(<vscale x 16 x i1> %pg, ptr %src, i64 0, i32 %add1)
+  call void @llvm.aarch64.sme.ld1w.horiz(<vscale x 16 x i1> %pg, ptr %src, i64 0, i32 %add2)
+  %inc = add nuw nsw i32 %i, 1
+  %exitcond.not = icmp eq i32 %inc, %N
+  br i1 %exitcond.not, label %exit, label %for.body
+
+exit:
+  ret void
+}
+
+
 declare void @llvm.aarch64.sme.ld1b.horiz(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.ld1h.horiz(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.ld1w.horiz(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.ld1d.horiz(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.ld1q.horiz(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.ld1b.vert(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.ld1h.vert(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.ld1w.vert(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.ld1d.vert(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.ld1q.vert(<vscale x 16 x i1>, ptr, i64, i32)
 
 declare void @llvm.aarch64.sme.ldr(i32, ptr)
 declare i64 @llvm.vscale.i64()

llvm/test/CodeGen/AArch64/sme-intrinsics-mova-extract.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme -verify-machineinstrs < %s | FileCheck %s
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve,+sme -verify-machineinstrs < %s | FileCheck %s
 
 define <vscale x 16 x i8> @extract_row_b(<vscale x 16 x i8> %zd, <vscale x 16 x i1> %pg, i32 %tileslice) {
 ; CHECK-LABEL: extract_row_b:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov w12, w0
 ; CHECK-NEXT:    mov z1.d, z0.d
 ; CHECK-NEXT:    mov z2.d, z0.d
 ; CHECK-NEXT:    mov z3.d, z0.d

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov w12, wzr
 ; CHECK-NEXT:    mov z0.q, p0/m, za15v.q[w12, 0]
 ; CHECK-NEXT:    ret
   %res = call <vscale x 2 x double> @llvm.aarch64.sme.readq.vert.nxv2f64(<vscale x 2 x double> %zd, <vscale x 2 x i1> %pg, i64 15, i32 0)
   ret <vscale x 2 x double> %res
 }
 
+define <vscale x 4 x i32> @test_sink_offset_operand(<vscale x 4 x i1> %pg, i32 %base, i32 %N) {
+; CHECK-LABEL: test_sink_offset_operand:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov w12, w0
+; CHECK-NEXT:    mov z0.s, #0 // =0x0
+; CHECK-NEXT:  .LBB26_1: // %for.body
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    mov z1.d, z0.d
+; CHECK-NEXT:    mov z2.d, z0.d
+; CHECK-NEXT:    mov z1.s, p0/m, za0h.s[w12, 0]
+; CHECK-NEXT:    mov z3.d, z0.d
+; CHECK-NEXT:    mov z2.s, p0/m, za0h.s[w12, 1]
+; CHECK-NEXT:    subs w1, w1, #3
+; CHECK-NEXT:    mov z3.s, p0/m, za0h.s[w12, 2]
+; CHECK-NEXT:    b.ne .LBB26_1
+; CHECK-NEXT:  // %bb.2: // %exit
+; CHECK-NEXT:    add z0.s, z1.s, z2.s
+; CHECK-NEXT:    add z0.s, z0.s, z3.s
+; CHECK-NEXT:    ret
+entry:
+  %add1 = add i32 %base, 1
+  %add2 = add i32 %base, 2
+  br label %for.body
+
+for.body:
+  %i = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %z0 = call <vscale x 4 x i32> @llvm.aarch64.sme.read.horiz.nxv4i32(<vscale x 4 x i32> zeroinitializer, <vscale x 4 x i1> %pg, i64 0, i32 %base)
+  %z1 = call <vscale x 4 x i32> @llvm.aarch64.sme.read.horiz.nxv4i32(<vscale x 4 x i32> zeroinitializer, <vscale x 4 x i1> %pg, i64 0, i32 %add1)
+  %z2 = call <vscale x 4 x i32> @llvm.aarch64.sme.read.horiz.nxv4i32(<vscale x 4 x i32> zeroinitializer, <vscale x 4 x i1> %pg, i64 0, i32 %add2)
+  %inc = add nuw nsw i32 %i, 3
+  %exitcond.not = icmp eq i32 %inc, %N
+  br i1 %exitcond.not, label %exit, label %for.body
+
+exit:
+  %tmp1 = add <vscale x 4 x i32> %z0, %z1
+  %res = add <vscale x 4 x i32> %tmp1, %z2
+  ret <vscale x 4 x i32> %res
+}
+
 declare <vscale x 16 x i8> @llvm.aarch64.sme.read.horiz.nxv16i8(<vscale x 16 x i8>, <vscale x 16 x i1>, i64, i32)
 declare <vscale x 8 x i16> @llvm.aarch64.sme.read.horiz.nxv8i16(<vscale x 8 x i16>, <vscale x 8 x i1>, i64, i32)
 declare <vscale x 8 x half> @llvm.aarch64.sme.read.horiz.nxv8f16(<vscale x 8 x half>, <vscale x 8 x i1>, i64, i32)
 declare <vscale x 8 x bfloat> @llvm.aarch64.sme.read.horiz.nxv8bf16(<vscale x 8 x bfloat>, <vscale x 8 x i1>, i64, i32)
 declare <vscale x 4 x i32> @llvm.aarch64.sme.read.horiz.nxv4i32(<vscale x 4 x i32>, <vscale x 4 x i1>, i64, i32)
 declare <vscale x 4 x float> @llvm.aarch64.sme.read.horiz.nxv4f32(<vscale x 4 x float>, <vscale x 4 x i1>, i64, i32)
 declare <vscale x 2 x i64> @llvm.aarch64.sme.read.horiz.nxv2i64(<vscale x 2 x i64>, <vscale x 2 x i1>, i64, i32)
 declare <vscale x 2 x double> @llvm.aarch64.sme.read.horiz.nxv2f64(<vscale x 2 x double>, <vscale x 2 x i1>, i64, i32)

llvm/test/CodeGen/AArch64/sme-intrinsics-mova-insert.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme -verify-machineinstrs < %s | FileCheck %s
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve,+sme -verify-machineinstrs < %s | FileCheck %s
 
 define void @insert_row_b(i32 %tileslice, <vscale x 16 x i1> %pg,
 ; CHECK-LABEL: insert_row_b:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov w12, w0
 ; CHECK-NEXT:    mov za0h.b[w12, 0], p0/m, z0.b
 ; CHECK-NEXT:    mov za0h.b[w12, 2], p0/m, z1.b
 ; CHECK-NEXT:    mov za0h.b[w12, 4], p0/m, z2.b

 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov w12, wzr
 ; CHECK-NEXT:    mov za15v.q[w12, 0], p0/m, z0.q
 ; CHECK-NEXT:    ret
   call void @llvm.aarch64.sme.writeq.vert.nxv2f64(i64 15, i32 0, <vscale x 2 x i1> %pg, <vscale x 2 x double> %zn)
   ret void
 }
 
+define void @test_sink_offset_operand(<vscale x 4 x i1> %pg, i32 %base, i32 %N) {
+; CHECK-LABEL: test_sink_offset_operand:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov w12, w0
+; CHECK-NEXT:    mov z0.s, #0 // =0x0
+; CHECK-NEXT:  .LBB28_1: // %for.body
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    subs w1, w1, #3
+; CHECK-NEXT:    mov za0h.s[w12, 0], p0/m, z0.s
+; CHECK-NEXT:    mov za0h.s[w12, 1], p0/m, z0.s
+; CHECK-NEXT:    mov za0h.s[w12, 2], p0/m, z0.s
+; CHECK-NEXT:    b.ne .LBB28_1
+; CHECK-NEXT:  // %bb.2: // %exit
+; CHECK-NEXT:    ret
+entry:
+  %add1 = add i32 %base, 1
+  %add2 = add i32 %base, 2
+  br label %for.body
+
+for.body:
+  %i = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  call void @llvm.aarch64.sme.write.horiz.nxv4i32(i64 0, i32 %base, <vscale x 4 x i1> %pg, <vscale x 4 x i32> zeroinitializer)
+  call void @llvm.aarch64.sme.write.horiz.nxv4i32(i64 0, i32 %add1, <vscale x 4 x i1> %pg, <vscale x 4 x i32> zeroinitializer)
+  call void @llvm.aarch64.sme.write.horiz.nxv4i32(i64 0, i32 %add2, <vscale x 4 x i1> %pg, <vscale x 4 x i32> zeroinitializer)
+  %inc = add nuw nsw i32 %i, 3
+  %exitcond.not = icmp eq i32 %inc, %N
+  br i1 %exitcond.not, label %exit, label %for.body
+
+exit:
+  ret void
+}
 
 declare void @llvm.aarch64.sme.write.horiz.nxv16i8(i64, i32, <vscale x 16 x i1>, <vscale x 16 x i8>)
 declare void @llvm.aarch64.sme.write.horiz.nxv8i16(i64, i32, <vscale x 8 x i1>, <vscale x 8 x i16>)
 declare void @llvm.aarch64.sme.write.horiz.nxv8f16(i64, i32, <vscale x 8 x i1>, <vscale x 8 x half>)
 declare void @llvm.aarch64.sme.write.horiz.nxv8bf16(i64, i32, <vscale x 8 x i1>, <vscale x 8 x bfloat>)
 declare void @llvm.aarch64.sme.write.horiz.nxv4i32(i64, i32, <vscale x 4 x i1>, <vscale x 4 x i32>)
 declare void @llvm.aarch64.sme.write.horiz.nxv4f32(i64, i32, <vscale x 4 x i1>, <vscale x 4 x float>)
 declare void @llvm.aarch64.sme.write.horiz.nxv2i64(i64, i32, <vscale x 2 x i1>, <vscale x 2 x i64>)

llvm/test/CodeGen/AArch64/sme-intrinsics-stores.ll

 ; CHECK-NEXT:    ret
   %vscale = call i64 @llvm.vscale.i64()
   %mulvl = mul i64 %vscale, 256
   %base = getelementptr i8, ptr %ptr, i64 %mulvl
   call void @llvm.aarch64.sme.str(i32 0, ptr %base)
   ret void;
 }
 
+; Ensure that the tile offset is sunk, given that this is likely to be an 'add'
+; that's decomposed into a base + offset in ISel.
+define void @test_sink_tile0_offset_operand(<vscale x 16 x i1> %pg, ptr %src, i32 %base, i32 %N) {
+; CHECK-LABEL: test_sink_tile0_offset_operand:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    mov w12, w1
+; CHECK-NEXT:  .LBB14_1: // %for.body
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    st1w {za0h.s[w12, 0]}, p0, [x0]
+; CHECK-NEXT:    subs w2, w2, #1
+; CHECK-NEXT:    st1w {za0h.s[w12, 1]}, p0, [x0]
+; CHECK-NEXT:    st1w {za0h.s[w12, 2]}, p0, [x0]
+; CHECK-NEXT:    b.ne .LBB14_1
+; CHECK-NEXT:  // %bb.2: // %exit
+; CHECK-NEXT:    ret
+entry:
+  %add0 = add i32 %base, 1
+  %add1 = add i32 %base, 2
+  br label %for.body
+
+for.body:
+  %i = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  tail call void @llvm.aarch64.sme.st1w.horiz(<vscale x 16 x i1> %pg, ptr %src, i64 0, i32 %base)
+  tail call void @llvm.aarch64.sme.st1w.horiz(<vscale x 16 x i1> %pg, ptr %src, i64 0, i32 %add0)
+  tail call void @llvm.aarch64.sme.st1w.horiz(<vscale x 16 x i1> %pg, ptr %src, i64 0, i32 %add1)
+  %inc = add nuw nsw i32 %i, 1
+  %exitcond.not = icmp eq i32 %inc, %N
+  br i1 %exitcond.not, label %exit, label %for.body
+
+exit:
+  ret void
+}
+
 declare void @llvm.aarch64.sme.st1b.horiz(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.st1h.horiz(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.st1w.horiz(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.st1d.horiz(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.st1q.horiz(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.st1b.vert(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.st1h.vert(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.st1w.vert(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.st1d.vert(<vscale x 16 x i1>, ptr, i64, i32)
 declare void @llvm.aarch64.sme.st1q.vert(<vscale x 16 x i1>, ptr, i64, i32)
 
 declare void @llvm.aarch64.sme.str(i32, ptr)
 declare i64 @llvm.vscale.i64()