Add warp reduce for integer min and max (#2933)

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

third_party/intel/lib/TritonIntelGPUToLLVM/SPIRVSubgroupOps.h

@@ -35,6 +35,18 @@ template <> struct SPIRVArithmeticGroupOp<arith::MulFOp> {
 template <> struct SPIRVArithmeticGroupOp<arith::MulIOp> {
   using type = spirv::GroupNonUniformIMulOp;
 };
+template <> struct SPIRVArithmeticGroupOp<arith::MaxSIOp> {
+  using type = spirv::GroupNonUniformSMaxOp;
+};
+template <> struct SPIRVArithmeticGroupOp<arith::MaxUIOp> {
+  using type = spirv::GroupNonUniformUMaxOp;
+};
+template <> struct SPIRVArithmeticGroupOp<arith::MinSIOp> {
+  using type = spirv::GroupNonUniformSMinOp;
+};
+template <> struct SPIRVArithmeticGroupOp<arith::MinUIOp> {
+  using type = spirv::GroupNonUniformUMinOp;
+};
 template <> struct SPIRVArithmeticGroupOp<arith::MaxNumFOp> {
   using type = spirv::GroupNonUniformFMaxOp;
 };

third_party/intel/lib/TritonIntelGPUToLLVM/TargetInfo.cpp

@@ -136,6 +136,7 @@ Value warpReduceHelper(RewriterBase &rewriter, Location loc, Value acc,
   Value warpReduce =
       TypeSwitch<mlir::Operation *, Value>(reduceOp)
           .Case<arith::AddFOp, arith::AddIOp, arith::MulFOp, arith::MulIOp,
+                arith::MaxSIOp, arith::MaxUIOp, arith::MinSIOp, arith::MinUIOp,
                 arith::MaxNumFOp, arith::MinNumFOp>([&](auto groupOp) {
             return createSPIRVGroupOp<
                 SPIRVArithmeticGroupOpTy<decltype(groupOp)>>(
@@ -182,9 +183,11 @@ bool TargetInfo::warpReduce(RewriterBase &rewriter, Location loc,
       reduceOp->getOperand(1) != block.getArgument(1))
     return false;
 
-  auto supportedOp = isa<arith::AddFOp, arith::AddIOp, arith::MulFOp,
-                         arith::MulIOp, arith::MaxNumFOp, arith::MinNumFOp,
-                         arith::AndIOp, arith::OrIOp, arith::XOrIOp>(reduceOp);
+  auto supportedOp =
+      isa<arith::AddFOp, arith::AddIOp, arith::MulFOp, arith::MulIOp,
+          arith::MaxSIOp, arith::MaxUIOp, arith::MinSIOp, arith::MinUIOp,
+          arith::MaxNumFOp, arith::MinNumFOp, arith::AndIOp, arith::OrIOp,
+          arith::XOrIOp>(reduceOp);
 
   if (!supportedOp)
     return false;