<table border="1" cellspacing="0" cellpadding="8">
    <tr>
        <th>Issue</th>
        <td>
            <a href=https://github.com/llvm/llvm-project/issues/210855>210855</a>
        </td>
    </tr>

    <tr>
        <th>Summary</th>
        <td>
            [mlir][math] vectorized bool-result ops trigger assertion in MathToNVVM
        </td>
    </tr>

    <tr>
      <th>Labels</th>
      <td>
      </td>
    </tr>

    <tr>
      <th>Assignees</th>
      <td>
      </td>
    </tr>

    <tr>
      <th>Reporter</th>
      <td>
          ashermancinelli
      </td>
    </tr>
</table>

<pre>
    I added this test based on a crash I encountered:
```
> cat mlir/test/Conversion/MathToNVVM/math-to-nvvm.mlir
// RUN: mlir-opt %s -convert-math-to-nvvm | FileCheck %s

// CHECK-LABEL:   func.func @fpclass_vector(
func.func @fpclass_vector(%arg: vector<2xf32>) -> (vector<2xi1>, vector<2xi1>, vector<2xi1>) {
  %inf = math.isinf %arg : vector<2xf32>
  // CHECK-COUNT-2: llvm.call @__nv_isinff(
  %finite = math.isfinite %arg : vector<2xf32>
  // CHECK-COUNT-2: llvm.call @__nv_finitef(
  %nan = math.isnan %arg : vector<2xf32>
  // CHECK-COUNT-2: llvm.call @__nv_isnanf(
  return %inf, %finite, %nan : vector<2xi1>, vector<2xi1>, vector<2xi1>
}

> build/bin/mlir-opt mlir/test/Conversion/MathToNVVM/math-to-nvvm.mlir -convert-math-to-nvvm
mlir-opt: /proj/llvm-project/mlir/lib/Conversion/MathToNVVM/../GPUCommon/OpToFuncCallLowering.h:80: 
llvm::LogicalResult mlir::OpToFuncCallLowering<SourceOp>::matchAndRewrite(
SourceOp, typename SourceOp::Adaptor, mlir::ConversionPatternRewriter&) const 
[with SourceOp = mlir::math::IsInfOp; typename SourceOp::Adaptor = mlir::math::IsInfOpAdaptor]: 
Assertion `(op->getResultTypes().front() == op->getOperand(0).getType() || isResultBool) 
&& "expected op with same operand and result types"' failed.
```

Looks like it's coming from here https://github.com/llvm/llvm-project/blob/main/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h#L75

I considered checking for bool vectors like the diff below but I have not looked at all the consumers of `OpToFuncCallLowering` to be sure this is right. Maybe we also need to check for a vectorizable trait?
```diff
diff --git a/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h b/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
index cb9b6da07183..07776620bd1f 100644
--- a/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
+++ b/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
@@ -13,6 +13,7 @@
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/TypeUtilities.h"

 namespace mlir {

@@ -72,7 +73,8 @@ public:
 std::is_base_of<OpTrait::OneResult<SourceOp>, SourceOp>::value,
 "expected single result op");

-    bool isResultBool = op->getResultTypes().front().isInteger(1);
+    bool isResultBool =
+ getElementTypeOrSelf(op->getResultTypes().front()).isInteger(1);
     if constexpr (!std::is_base_of<OpTrait::SameOperandsAndResultType<SourceOp>,
 SourceOp>::value) {
 assert(op->getNumOperands() > 0 &&
```

I'm hoping someone can chime in on whether this fix looks right or not, or if they have another suggestion. TIV!
</pre>
<img width="1" height="1" alt="" src="http://email.email.llvm.org/o/eJysV11zmzwW_jXyzRkYWdjGvvAFJvG7nnWbTpv0NiPgANoKiZGEk-yv35HAjvdt2ma3nXEyQpxvnfM8glsrGoW4JcsdWd7M-OBabbbctmg6rkqhUEoxK3T1sj0AryqswLXCgkProOAWK9AKOJSG2xYOgKrUg3JosCJJRmhGVnT60Ywkt1ByB50UhrC9t0HYPtfqhMYKrQjbf-Cuvdcfv379QNi-466NnI7U6dTFQYlmhO0J28Pnh48kyYKlSPcOCFtaiMpgykXXikDSHPZCYt5i-S0IejMXS_k_bvN_Rsdsd3v0FgHqQZWx_wdkQeu-lNzaxxOWThvC1oRmvxBgS24ab2raSnL2XCeMJLeEbSDyRSBsffVSzMO7HN63twGS7gjNwOciVA0kuQGfcSxseAwBwNsRjGpXied3Dx_vI-alpTx1ccml9Hk9PqrTYzBYj1kHd7VQwuG1x_POn3M6Wrz2qri6dhke_2SSiquLO4NuMGoqrT-AS9bTwxhM9s7DemOPZiS9mVowuYViELIibF8I3_-Xhv6_Z-TtKSA0O5v2wRO2743-F2F7X4_Ir7F0k3-_K4qfeo1jwvZ_fXrIddeFt3f9vd4Pqsy5lEf9hEaoJm5Jkq1p8Ecz78hDQpIddSNKLj-jHeSUaNh_ywZJ8i96MCXe9b52Qa7jrmwzVX3GJxPOxZ_cRYrl4F56VLxDeFX1elnF-zCj-ZXT1xw_cefQqMmqIWzlZ63UyjofP1nunoRrLzbHjrzY8cUeVwd7ULX3uftFIL-ycI53eTOVMLMWjRNagcdTtta9R5MG3VjK-5cebUCgTVwbrdy49m68p4v0XY-Gq4qwNfWiDQbNs3Cae8AUdrS501qGbY-XK8JWQBjD5x5L56G_h1AT63PUo1nwf2Y8WzdGxAhLoeZCYhX_nRJodtT6mwUpviEIR1hqodSdUA3URnfQokFoneutr00Y6ka4dijiUndT_37fxoXURRiOy1C93dTvaGGWHNPlGOkhdIOoPLtB6ekkhKkNFFrLadKnVFyLUIm6hgKlfoJicHCAlp8QlHYgtf6GFXAHHom8rLc8dGgs6Nof75vDsKLgNBQIdjA4ErGwYETTuhg-8JcC4QmBS6tBoadqPYYZYuRTgOLfvJAIznDhSLK_PhAfMKFZiDuKGuGA_2b5oPjd-tNMqAqfoSw2xariNJ2vkzimaZquVowW1byGOaWrxYLQLIqi3444dPpu_P2J6MmCkgWFaJ4Qlvvx2YVVCuMLTzqEJUKVcqg8j7LJ343gcmzmvz49ELY_fB6X077vTPYO5ePx64eg6hfv0g3SO89KaOxZNNTjR8IePh6ckMIJfNWgGXjssz0vMcDcdG-5qknKxkqwXeprsp5qAv1QSFGOF0iwrhphUdhHf-N81DVJ8rv-fmzgwBwKR7j6G12wHL5jjxOXg-fyMf9XLLNCNRLPyKX7AFsbkkwhRwAwjvk1NMJ_AetPYTgW9qAcNp5Z1vNX02z3Q8vT6wbdrcQOVTB9Z76grN-N_j_27N2CqEeGw-feQNCYv6PeX3iHE4vYQMTnAL6rv_fzoyO4XGR54LXrlD4O3dn-hcRugcLIQd9zyIGwtINW9x6Qre5QK4SSKyhb0SEI5b9Snlp0LZoRN2vxHFB4gk_QxgOz7xhtfFVciy8jXnOlg5odmgatp98Y7g9fCZvPqm1SbZINn-F2nq4Xy01C18tZu-X1BotFUizKzTpJCloUm2TBqnpBN3zB6uVMbBllK5qyOaXzNV3GrKQ8oWtkq7Jm6yohC4odFzIO91VtmpmwdsAtm9P1cjmTvEBpz19tZhvorxgaSxZUCuvsq54TTobvuzCxyxu_8jeN5c2FD7AK_Rddet-CM6Jp0ExH468cQsHrFXA2GLn9n1k5pGAJ209ZnLbsPwEAAP__3ayIyQ">