[llvm] [AMDGPU] Only fold mul24+add into V_MAD_U64_U32 when operands fit in 24 bits (PR #198729)

[Jay Foad via llvm-commits]

================
@@ -1070,16 +1070,50 @@ multiclass IMAD32_Pats <VOP3_Pseudo inst> {
         >;
 }
 
+// AMDGPUmul_u24 only consumes the low 24 bits of its operands, and
+// simplifyMul24 may strip outer AND masks. V_MAD_U64_U32 does a full 32x32
+// multiply, so this fold is only valid when both operands are known to fit
+// in 24 unsigned bits at SDAG level.
+def AMDGPUmul_u24_24bit_ops : PatFrag<
+  (ops node:$src0, node:$src1),
+  (AMDGPUmul_u24 node:$src0, node:$src1),
+  [{
+    // AMDGPUmul_u24 is a PatFrags over an intrinsic and an SDNode; skip the
+    // intrinsic ID operand when needed.
+    unsigned BaseIdx = N->getOpcode() == ISD::INTRINSIC_WO_CHAIN ? 1 : 0;
+    return AMDGPUTargetLowering::numBitsUnsigned(N->getOperand(BaseIdx),
+                                                 *CurDAG) <= 24 &&
+           AMDGPUTargetLowering::numBitsUnsigned(N->getOperand(BaseIdx + 1),
+                                                 *CurDAG) <= 24;
+  }]>;
+
 // Handle cases where amdgpu-codegenprepare-mul24 made a mul24 instead of a normal mul.
 // We need to separate this because otherwise OtherPredicates would be overriden.
+// Preferred form: operands provably fit in 24 unsigned bits, no masking needed.
 class IMAD32_Mul24_Pats_Impl<VOP3_Pseudo inst, SDPatternOperator AddOp> : GCNPat <
-    (i64 (AddOp (i64 (AMDGPUmul_u24 i32:$src0, i32:$src1)), i64:$src2)),
+    (i64 (AddOp (i64 (AMDGPUmul_u24_24bit_ops i32:$src0, i32:$src1)), i64:$src2)),
     (inst $src0, $src1, $src2, 0 /* clamp */)
+    > {
+  let AddedComplexity = 1;
+}
+
+// Fallback: AMDGPUmul_u24 only uses the low 24 bits of its operands, but
+// V_MAD_U64_U32 does a full 32x32 multiply. Mask the high bits explicitly so
+// the fold remains valid even when SDAG cannot prove the operands already fit
+// in 24 bits. DAGCombine folds the AND away when the high bits are known zero,
+// recovering the single-instruction form.
----------------
jayfoad wrote:

I don't understand this part. No DAG combine will remove the V_AND_B32 instructions that you generate here, will it? E.g. if only one of the operands is known to fit in 24 bits already.

In general, it is not a good idea to write selection patterns that generate more than one instruction, for this reason.

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