[NVPTX, LSV] Move the LSV optimization pass to later when the graph is cleaner

[nouiz]

Summary

This allow it to recognize more loads as being consecutive when the load's address are complex at the start.

[tra]

llvm/test/CodeGen/NVPTX/vector-loads-complex.ll

12 ↗

(On Diff #243996)

Couple of general notes:

I'd attempt to reduce IR to the bare minimum necessary to demonstrate/test intended behavior. This IR is not bad, but I think it may be reduced further to the point where we end up with just one vectorized 2-element load.

It's a good practice to 'anchor the checks with fixed start/end points

CHECK-LABEL: <function_name>

... # other checks on whatever happens inside the function.

CHECK: ret

It helps to delineate the search boundaries for FileCheck. Not a big deal now, as you only have one function, but I would not be surprised if more tests would be added over time.

14–18 ↗

(On Diff #243996)

Negative matching is tricky. As written, FileCheck will only ensure that there's no ld.global between the fourth instance of ld.global.v2.u8 and the end of the file.

I'd start by attempting to place the checks close to the load statements they coalesce into. It's not always possible, but it probably works in this case.

As long as the number of loads is fixed/predictable, positively matching all of them makes negative matching unnecessary, so I'd just drop it.

[nouiz]

I simplified the test. As it was way simpler, I added it to an existing test file.

[tra]

Minor test nits. LGTM otherwise.
I'm curious, what does generated ptx for the function look before/after the patch.

llvm/test/CodeGen/NVPTX/vector-loads.ll

73

CHECK-LABEL

75

I'm not sure we need it. The check above verifies that the load did get vectorized, which was the purpose of this patch and the test.

Any condition when CHECK-NOT may fail (e.g. we're loading more than one, or it's an unrelated load/store) is unlikely to be relevant. It will also only look in-between ld.v2.u8 and ret, and such issues will not be detected before ld.v2.u8. If you want to ensure that ld.v2.u8 is the only relevant load, then you may need to have a separate test case which would fail if it's found more than one ld of any kind. Or you can add another CHECK-NOT: ld between the start of the function and CHECK: ld.v2.u8. I think that may work, but I'm not completely sure.

I'd just skip CHECK-NOT altogether.

92

CHECK: ld.v2.u8

98

CHECK: ret

[nouiz]

Before:

.visible .func foo_complex(

.param .b32 foo_complex_param_0

) // @foo_complex
{

.reg .b16       %rs<4>;
.reg .b32       %r<11>;

// %bb.0:

ld.param.u32    %r1, [foo_complex_param_0];
mov.u32         %r2, %tid.x;
mov.u32         %r3, %ctaid.x;
shl.b32         %r4, %r3, 9;
and.b32         %r5, %r4, -131072;
and.b32         %r6, %r4, 130560;
shl.b32         %r7, %r2, 1;
or.b32          %r8, %r6, %r7;
add.s32         %r9, %r1, %r5;
add.s32         %r10, %r9, %r8;

• ld.u8 %rs1, [%r10+128]; ld.u8 %rs2, [%r10+129];** max.u16 %rs3, %rs1, %rs2; st.u8 [%r10+129], %rs3; ret; // -- End function

}

After:

.visible .func foo_complex(

.param .b32 foo_complex_param_0

) // @foo_complex
{

.reg .b16       %rs<4>;
.reg .b32       %r<11>;

// %bb.0:

ld.param.u32    %r1, [foo_complex_param_0];
mov.u32         %r2, %tid.x;
mov.u32         %r3, %ctaid.x;
shl.b32         %r4, %r3, 9;
and.b32         %r5, %r4, -131072;
and.b32         %r6, %r4, 130560;
shl.b32         %r7, %r2, 1;
or.b32          %r8, %r6, %r7;
add.s32         %r9, %r1, %r5;
add.s32         %r10, %r9, %r8;

• ld.v2.u8 {%rs1, %rs2}, [%r10+128];** max.u16 %rs3, %rs1, %rs2; st.u8 [%r10+129], %rs3; ret; // -- End function

}

[nouiz]

llvm/test/CodeGen/NVPTX/vector-loads.ll

73

Here I did as the other tests in this file. I'll fix all of them.

[nouiz]

I did all the small changes.

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

void NVPTXPassConfig::addIRPasses() {
   addPass(createNVPTXAssignValidGlobalNamesPass());
   addPass(createGenericToNVVMPass());
 
   // NVPTXLowerArgs is required for correctness and should be run right
   // before the address space inference passes.
   addPass(createNVPTXLowerArgsPass(&getNVPTXTargetMachine()));
   if (getOptLevel() != CodeGenOpt::None) {
     addAddressSpaceInferencePasses();
-    if (!DisableLoadStoreVectorizer)
-      addPass(createLoadStoreVectorizerPass());
     addStraightLineScalarOptimizationPasses();
   }
 
   // === LSR and other generic IR passes ===
   TargetPassConfig::addIRPasses();
   // EarlyCSE is not always strong enough to clean up what LSR produces. For
   // example, GVN can combine
   //
   //   %0 = add %a, %b
   //   %1 = add %b, %a
   //
   // and
   //
   //   %0 = shl nsw %a, 2
   //   %1 = shl %a, 2
   //
   // but EarlyCSE can do neither of them.
-  if (getOptLevel() != CodeGenOpt::None)
+  if (getOptLevel() != CodeGenOpt::None) {
     addEarlyCSEOrGVNPass();
+    if (!DisableLoadStoreVectorizer)
+      addPass(createLoadStoreVectorizerPass());
+  }
 }
 
 bool NVPTXPassConfig::addInstSelector() {
   const NVPTXSubtarget &ST = *getTM<NVPTXTargetMachine>().getSubtargetImpl();
 
   addPass(createLowerAggrCopies());
   addPass(createAllocaHoisting());
   addPass(createNVPTXISelDag(getNVPTXTargetMachine(), getOptLevel()));

llvm/test/CodeGen/NVPTX/vector-loads.ll

 ; RUN: llc < %s -march=nvptx -mcpu=sm_20 | FileCheck %s
 
 ; Even though general vector types are not supported in PTX, we can still
 ; optimize loads/stores with pseudo-vector instructions of the form:
 ;
 ; ld.v2.f32 {%f0, %f1}, [%r0]
 ;
 ; which will load two floats at once into scalar registers.
 
+; CHECK-LABEL foo
 define void @foo(<2 x float>* %a) {
-; CHECK: .func foo
 ; CHECK: ld.v2.f32 {%f{{[0-9]+}}, %f{{[0-9]+}}}
   %t1 = load <2 x float>, <2 x float>* %a
   %t2 = fmul <2 x float> %t1, %t1
   store <2 x float> %t2, <2 x float>* %a
   ret void
 }
 
+; CHECK-LABEL foo2
 define void @foo2(<4 x float>* %a) {
-; CHECK: .func foo2
 ; CHECK: ld.v4.f32 {%f{{[0-9]+}}, %f{{[0-9]+}}, %f{{[0-9]+}}, %f{{[0-9]+}}}
   %t1 = load <4 x float>, <4 x float>* %a
   %t2 = fmul <4 x float> %t1, %t1
   store <4 x float> %t2, <4 x float>* %a
   ret void
 }
 
+; CHECK-LABEL foo3
 define void @foo3(<8 x float>* %a) {
-; CHECK: .func foo3
 ; CHECK: ld.v4.f32 {%f{{[0-9]+}}, %f{{[0-9]+}}, %f{{[0-9]+}}, %f{{[0-9]+}}}
 ; CHECK-NEXT: ld.v4.f32 {%f{{[0-9]+}}, %f{{[0-9]+}}, %f{{[0-9]+}}, %f{{[0-9]+}}}
   %t1 = load <8 x float>, <8 x float>* %a
   %t2 = fmul <8 x float> %t1, %t1
   store <8 x float> %t2, <8 x float>* %a
   ret void
 }
 
 
 
+; CHECK-LABEL foo4
 define void @foo4(<2 x i32>* %a) {
-; CHECK: .func foo4
 ; CHECK: ld.v2.u32 {%r{{[0-9]+}}, %r{{[0-9]+}}}
   %t1 = load <2 x i32>, <2 x i32>* %a
   %t2 = mul <2 x i32> %t1, %t1
   store <2 x i32> %t2, <2 x i32>* %a
   ret void
 }
 
+; CHECK-LABEL foo5
 define void @foo5(<4 x i32>* %a) {
-; CHECK: .func foo5
 ; CHECK: ld.v4.u32 {%r{{[0-9]+}}, %r{{[0-9]+}}, %r{{[0-9]+}}, %r{{[0-9]+}}}
   %t1 = load <4 x i32>, <4 x i32>* %a
   %t2 = mul <4 x i32> %t1, %t1
   store <4 x i32> %t2, <4 x i32>* %a
   ret void
 }
 
+; CHECK-LABEL foo6
 define void @foo6(<8 x i32>* %a) {
-; CHECK: .func foo6
 ; CHECK: ld.v4.u32 {%r{{[0-9]+}}, %r{{[0-9]+}}, %r{{[0-9]+}}, %r{{[0-9]+}}}
 ; CHECK-NEXT: ld.v4.u32 {%r{{[0-9]+}}, %r{{[0-9]+}}, %r{{[0-9]+}}, %r{{[0-9]+}}}
   %t1 = load <8 x i32>, <8 x i32>* %a
   %t2 = mul <8 x i32> %t1, %t1
   store <8 x i32> %t2, <8 x i32>* %a
   ret void
 }
+
+; The following test wasn't passing previously as the address
+; computation was still too complex when LSV was called.
+declare i32 @llvm.nvvm.read.ptx.sreg.ctaid.x() #0
+declare i32 @llvm.nvvm.read.ptx.sreg.tid.x() #0
+; CHECK-LABEL foo_complex
+define void @foo_complex(i8* nocapture readonly align 16 dereferenceable(134217728) %alloc0) {
+  %targ0.1.typed = bitcast i8* %alloc0 to [1024 x [131072 x i8]]*
+  %t0 = tail call i32 @llvm.nvvm.read.ptx.sreg.tid.x(), !range !1
+  %t1 = tail call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
+  %t2 = lshr i32 %t1, 8
+  %t3 = shl nuw nsw i32 %t1, 9
+  %ttile_origin.2 = and i32 %t3, 130560
+  %tstart_offset_x_mul = shl nuw nsw i32 %t0, 1
+  %t4 = or i32 %ttile_origin.2, %tstart_offset_x_mul
+  %t6 = or i32 %t4, 1
+  %t8 = or i32 %t4, 128
+  %t9 = zext i32 %t8 to i64
+  %t10 = or i32 %t4, 129
+  %t11 = zext i32 %t10 to i64
+  %t20 = zext i32 %t2 to i64
+  %t27 = getelementptr inbounds [1024 x [131072 x i8]], [1024 x [131072 x i8]]* %targ0.1.typed, i64 0, i64 %t20, i64 %t9
+; CHECK: ld.v2.u8
+  %t28 = load i8, i8* %t27, align 2
+  %t31 = getelementptr inbounds [1024 x [131072 x i8]], [1024 x [131072 x i8]]* %targ0.1.typed, i64 0, i64 %t20, i64 %t11
+  %t32 = load i8, i8* %t31, align 1
+  %t33 = icmp ult i8 %t28, %t32
+  %t34 = select i1 %t33, i8 %t32, i8 %t28
+  store i8 %t34, i8* %t31
+; CHECK: ret
+  ret void
+}
+
+
+!1 = !{i32 0, i32 64}