Respond to some review comments

include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td

@@ -100,8 +100,8 @@ def TTG_SelectOp : TTG_Op<"select", [Pure, Elementwise,
   let description = [{}];
 
   let arguments = (ins TT_BoolLike:$condition,
-                       TT_Type:$true_value,
-                       TT_Type:$false_value);
+                       TT_Tensor:$true_value,
+                       TT_Tensor:$false_value);
 
   let results = (outs TT_Type:$result);
 }

lib/Analysis/Membar.cpp

@@ -76,7 +76,7 @@ void MembarAnalysis::visitTerminator(Operation *op,
   if (isa<triton::ReduceReturnOp>(op) || isa<func::ReturnOp>(op)) {
     return;
   }
-  op->emitOpError("Unknown terminator encountered in membar analysis");
+  llvm_unreachable("Unknown terminator encountered in membar analysis");
 }
 
 void MembarAnalysis::update(Operation *op, BlockInfo *blockInfo,

lib/Analysis/Utility.cpp

@@ -71,11 +71,11 @@ unsigned ReduceOpHelper::getScratchSizeInBytes() {
     elems = product<unsigned>(smemShape);
   }
 
-  unsigned bytes_per_elem = 0;
+  unsigned bytesPerElem = 0;
   for (const auto &ty : srcElementTypes) {
-    bytes_per_elem += ty.getIntOrFloatBitWidth() / 8;
+    bytesPerElem += ty.getIntOrFloatBitWidth() / 8;
   }
-  return bytes_per_elem * elems;
+  return bytesPerElem * elems;
 }
 
 bool isSharedEncoding(Value value) {

lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp

@@ -106,8 +106,8 @@ class SimplifyReduceCvt : public mlir::RewritePattern {
         continue;
       }
 
-      // TODO: This always takes layout from the first argument which
-      // is fine for argmin/argmax but may not be optimal generally
+      // TODO: This only moves conversions from the first argument which is
+      // fine for argmin/argmax but may not be optimal generally
       if (convert.getResult() != owner.getOperands()[0]) {
         continue;
       }
@@ -123,8 +123,8 @@ class SimplifyReduceCvt : public mlir::RewritePattern {
     auto newEncoding =
         newOperands[0].getType().cast<RankedTensorType>().getEncoding();
 
-    if (!newEncoding.isa<triton::gpu::BlockedEncodingAttr>()) {
-      // ReduceOpToLLVM requires block encoding
+    // this may generate unsupported conversions in the LLVM codegen
+    if (newEncoding.isa<triton::gpu::MmaEncodingAttr>()) {
       return failure();
     }
 

lib/Dialect/TritonGPU/Transforms/TritonGPUConversion.cpp

@@ -80,7 +80,7 @@ TritonGPUConversionTarget::TritonGPUConversionTarget(
   addIllegalOp<scf::ExecuteRegionOp, scf::ParallelOp, scf::ReduceOp,
                scf::ReduceReturnOp>();
   // We have custom versions of some arith operators
-  addIllegalOp<arith::CmpIOp, arith::CmpFOp, arith::SelectOp>();
+  addIllegalOp<arith::CmpIOp, arith::CmpFOp>();
 
   addDynamicallyLegalDialect<arith::ArithDialect, math::MathDialect,
                              func::FuncDialect, triton::TritonDialect,