[mlir][spirv] Fix math.ctlz for full zero bit cases

[antiagainst]

Summary

If the integer has all zero bits, GLSL FindUMsb would return -1.
So theoretically (31 - FindUMsb) should still give use the correct
result. However, Adreno GPUshave issues with this:
https://buildkite.com/iree/iree-test-android/builds/6482#01815f05-3926-466f-822a-1e20299e5461
This looks like a driver bug. So handle the corner case explicity
to workaround it.

[antiagainst]

Rebase

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

matchAndRewrite(math::CountLeadingZerosOp countOp, OpAdaptor adaptor,
     if (type.isa<IntegerType>())
       bitwidth = type.getIntOrFloatBitWidth();
     if (auto vectorType = type.dyn_cast<VectorType>())
       bitwidth = vectorType.getElementTypeBitWidth();
     if (bitwidth != 32)
       return failure();
 
     Location loc = countOp.getLoc();
+    Value allOneBits = getScalarOrVectorI32Constant(type, -1, rewriter, loc);
+    Value val32 = getScalarOrVectorI32Constant(type, 32, rewriter, loc);
     Value val31 = getScalarOrVectorI32Constant(type, 31, rewriter, loc);
     Value msb =
         rewriter.create<spirv::GLSLFindUMsbOp>(loc, adaptor.getOperand());
     // We need to subtract from 31 given that the index is from the least
     // significant bit.
-    rewriter.replaceOpWithNewOp<spirv::ISubOp>(countOp, val31, msb);
+    Value sub = rewriter.create<spirv::ISubOp>(loc, val31, msb);
+    // If the integer has all zero bits, GLSL FindUMsb would return -1. So
+    // theoretically (31 - FindUMsb) should still give the correct result.
+    // However, certain Vulkan implementations have driver bugs regarding it.
+    // So handle the corner case explicity to workaround it.
+    Value cmp = rewriter.create<spirv::IEqualOp>(loc, msb, allOneBits);
+    rewriter.replaceOpWithNewOp<spirv::SelectOp>(countOp, cmp, val32, sub);
     return success();
   }
 };
 
 /// Converts math.expm1 to SPIR-V ops.
 ///
 /// SPIR-V does not have a direct operations for exp(x)-1. Explicitly lower to
 /// these operations.

mlir/test/Conversion/MathToSPIRV/math-to-glsl-spirv.mlir

func.func @float32_ternary_vector(%a: vector<4xf32>, %b: vector<4xf32>,
   // CHECK: spv.GLSL.Fma %{{.*}}: vector<4xf32>
   %0 = math.fma %a, %b, %c : vector<4xf32>
   return
 }
 
 // CHECK-LABEL: @ctlz_scalar
 //  CHECK-SAME: (%[[VAL:.+]]: i32)
 func.func @ctlz_scalar(%val: i32) -> i32 {
-  // CHECK: %[[V31:.+]] = spv.Constant 31 : i32
+  // CHECK-DAG: %[[MAX:.+]] = spv.Constant -1 : i32
+  // CHECK-DAG: %[[V32:.+]] = spv.Constant 32 : i32
+  // CHECK-DAG: %[[V31:.+]] = spv.Constant 31 : i32
   // CHECK: %[[MSB:.+]] = spv.GLSL.FindUMsb %[[VAL]] : i32
   // CHECK: %[[SUB:.+]] = spv.ISub %[[V31]], %[[MSB]] : i32
-  // CHECK: return %[[SUB]]
+  // CHECK: %[[CMP:.+]] = spv.IEqual %[[MSB]], %[[MAX]] : i32
+  // CHECK: %[[R:.+]] = spv.Select %[[CMP]], %[[V32]], %[[SUB]] : i1, i32
+  // CHECK: return %[[R]]
   %0 = math.ctlz %val : i32
   return %0 : i32
 }
 
 // CHECK-LABEL: @ctlz_vector1
 func.func @ctlz_vector1(%val: vector<1xi32>) -> vector<1xi32> {
   // CHECK: spv.GLSL.FindUMsb
   // CHECK: spv.ISub
+  // CHECK: spv.IEqual
+  // CHECK: spv.Select
   %0 = math.ctlz %val : vector<1xi32>
   return %0 : vector<1xi32>
 }
 
 // CHECK-LABEL: @ctlz_vector2
 //  CHECK-SAME: (%[[VAL:.+]]: vector<2xi32>)
 func.func @ctlz_vector2(%val: vector<2xi32>) -> vector<2xi32> {
+  // CHECK-DAG: %[[MAX:.+]] = spv.Constant dense<-1> : vector<2xi32>
+  // CHECK-DAG: %[[V32:.+]] = spv.Constant dense<32> : vector<2xi32>
   // CHECK-DAG: %[[V31:.+]] = spv.Constant dense<31> : vector<2xi32>
   // CHECK: %[[MSB:.+]] = spv.GLSL.FindUMsb %[[VAL]] : vector<2xi32>
   // CHECK: %[[SUB:.+]] = spv.ISub %[[V31]], %[[MSB]] : vector<2xi32>
-  // CHECK: return %[[SUB]]
+  // CHECK: %[[CMP:.+]] = spv.IEqual %[[MSB]], %[[MAX]] : vector<2xi32>
+  // CHECK: %[[R:.+]] = spv.Select %[[CMP]], %[[V32]], %[[SUB]] : vector<2xi1>, vector<2xi32>
+  // CHECK: return %[[R]]
   %0 = math.ctlz %val : vector<2xi32>
   return %0 : vector<2xi32>
 }
 
 } // end module
 
 // -----