[llvm] [NVPTX] Improve folding to mad with immediate 1 (PR #93628)

[Alex MacLean via llvm-commits]

================
@@ -0,0 +1,101 @@
+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -O1 | FileCheck %s
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -O1 | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx -mcpu=sm_20 -O1 | %ptxas-verify %}
+; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_20 -O1 | %ptxas-verify %}
+
+define i32 @test1(i32 %n, i32 %m) {
+;
+; CHECK: ld.param.u32   %[[N:r[0-9]+]], [test1_param_0];
+; CHECK: ld.param.u32   %[[M:r[0-9]+]], [test1_param_1];
+; CHECK: mad.lo.s32     %[[MAD:r[0-9]+]], %[[M]], %[[N]], %[[M]];
+; CHECK: st.param.b32   [func_retval0+0], %[[MAD]];
+;
+  %add = add i32 %n, 1
+  %mul = mul i32 %add, %m
+  ret i32 %mul
+}
+
+define i32 @test1_rev(i32 %n, i32 %m) {
+;
+; CHECK: ld.param.u32   %[[N:r[0-9]+]], [test1_rev_param_0];
+; CHECK: ld.param.u32   %[[M:r[0-9]+]], [test1_rev_param_1];
+; CHECK: mad.lo.s32     %[[MAD:r[0-9]+]], %[[M]], %[[N]], %[[M]];
+; CHECK: st.param.b32   [func_retval0+0], %[[MAD]];
+;
+  %add = add i32 %n, 1
+  %mul = mul i32 %m, %add
+  ret i32 %mul
+}
+
+; Transpose (mul (select)) if it can then be folded to mad
----------------
AlexMaclean wrote:

By itself this transform doesn't help much, I agree. However, if `m` or `n` are `add(x,1)` then it enables the other transformation. In the code we're checking for this case and only running the transformation when it would enable further folding. A rare case to be sure, but better to support it than not.

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