[mlir][spirv] Allow vectors of index types in elementwise conversions

[qedawkins]

Summary

Currently the conversion of elementwise ops only checks for scalar index
types when checking for bitwidth emulation.

mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp

CmpIOpPattern::matchAndRewrite(arith::CmpIOp op, OpAdaptor adaptor,
   Type dstType = getTypeConverter()->convertType(srcType);
   if (!dstType)
     return getTypeConversionFailure(rewriter, op, srcType);
 
   switch (op.getPredicate()) {
 #define DISPATCH(cmpPredicate, spirvOp)                                        \
   case cmpPredicate:                                                           \
     if (spirvOp::template hasTrait<OpTrait::spirv::UnsignedOp>() &&            \
-        srcType != dstType && !hasSameBitwidth(srcType, dstType)) {            \
+        !getElementTypeOrSelf(srcType).isIndex() && srcType != dstType &&      \
+        !hasSameBitwidth(srcType, dstType)) {                                  \
       return op.emitError(                                                     \
           "bitwidth emulation is not implemented yet on unsigned op");         \
     }                                                                          \
     rewriter.replaceOpWithNewOp<spirvOp>(op, adaptor.getLhs(),                 \
                                          adaptor.getRhs());                    \
     return success();
 
     DISPATCH(arith::CmpIPredicate::eq, spirv::IEqualOp);

mlir/lib/Conversion/SPIRVCommon/Pattern.h

 //===- Pattern.h - SPIRV Common Conversion Patterns -----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_SPIRVCOMMON_PATTERN_H
 #define MLIR_CONVERSION_SPIRVCOMMON_PATTERN_H
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVOpTraits.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/Support/FormatVariadic.h"
 
 namespace mlir {
 namespace spirv {
 
 /// Converts elementwise unary, binary and ternary standard operations to SPIR-V
 /// operations.
 template <typename Op, typename SPIRVOp>
 struct ElementwiseOpPattern : public OpConversionPattern<Op> {
   using OpConversionPattern<Op>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(Op op, typename Op::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     assert(adaptor.getOperands().size() <= 3);
     Type dstType = this->getTypeConverter()->convertType(op.getType());
     if (!dstType) {
       return rewriter.notifyMatchFailure(
           op->getLoc(),
           llvm::formatv("failed to convert type {0} for SPIR-V", op.getType()));
     }
 
     if (SPIRVOp::template hasTrait<OpTrait::spirv::UnsignedOp>() &&
-        !op.getType().isIndex() && dstType != op.getType()) {
-      return op.emitError(
-          "bitwidth emulation is not implemented yet on unsigned op");
+        !getElementTypeOrSelf(op.getType()).isIndex() &&
+        dstType != op.getType()) {
+      op.dump();
+      return op.emitError("bitwidth emulation is not implemented yet on "
+                          "unsigned op pattern version");
     }
     rewriter.template replaceOpWithNewOp<SPIRVOp>(op, dstType,
                                                   adaptor.getOperands());
     return success();
   }
 };
 
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_SPIRVCOMMON_PATTERN_H

mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir

 func.func @int_vector234(%arg0: vector<2xi8>, %arg1: vector<4xi64>) {
   // CHECK: spirv.SDiv %{{.*}}, %{{.*}}: vector<2xi8>
   %0 = arith.divsi %arg0, %arg0: vector<2xi8>
   // CHECK: spirv.UDiv %{{.*}}, %{{.*}}: vector<4xi64>
   %1 = arith.divui %arg1, %arg1: vector<4xi64>
   return
 }
 
+// CHECK-LABEL: @index_vector
+func.func @index_vector(%arg0: vector<4xindex>) {
+  // CHECK: spirv.UMod %{{.*}}, %{{.*}}: vector<4xi32>
+  %0 = arith.remui %arg0, %arg0: vector<4xindex>
+  return
+}
+
 // CHECK-LABEL: @vector_srem
 // CHECK-SAME: (%[[LHS:.+]]: vector<3xi16>, %[[RHS:.+]]: vector<3xi16>)
 func.func @vector_srem(%arg0: vector<3xi16>, %arg1: vector<3xi16>) {
   // CHECK: %[[LABS:.+]] = spirv.GL.SAbs %[[LHS]] : vector<3xi16>
   // CHECK: %[[RABS:.+]] = spirv.GL.SAbs %[[RHS]] : vector<3xi16>
   // CHECK:  %[[ABS:.+]] = spirv.UMod %[[LABS]], %[[RABS]] : vector<3xi16>
   // CHECK:  %[[POS:.+]] = spirv.IEqual %[[LHS]], %[[LABS]] : vector<3xi16>
   // CHECK:  %[[NEG:.+]] = spirv.SNegate %[[ABS]] : vector<3xi16>
func.func @cmpi(%arg0 : i32, %arg1 : i32) {
   %7 = arith.cmpi ule, %arg0, %arg1 : i32
   // CHECK: spirv.UGreaterThan
   %8 = arith.cmpi ugt, %arg0, %arg1 : i32
   // CHECK: spirv.UGreaterThanEqual
   %9 = arith.cmpi uge, %arg0, %arg1 : i32
   return
 }
 
+// CHECK-LABEL: @indexcmpi
+func.func @indexcmpi(%arg0 : index, %arg1 : index) {
+  // CHECK: spirv.IEqual
+  %0 = arith.cmpi eq, %arg0, %arg1 : index
+  // CHECK: spirv.INotEqual
+  %1 = arith.cmpi ne, %arg0, %arg1 : index
+  // CHECK: spirv.SLessThan
+  %2 = arith.cmpi slt, %arg0, %arg1 : index
+  // CHECK: spirv.SLessThanEqual
+  %3 = arith.cmpi sle, %arg0, %arg1 : index
+  // CHECK: spirv.SGreaterThan
+  %4 = arith.cmpi sgt, %arg0, %arg1 : index
+  // CHECK: spirv.SGreaterThanEqual
+  %5 = arith.cmpi sge, %arg0, %arg1 : index
+  // CHECK: spirv.ULessThan
+  %6 = arith.cmpi ult, %arg0, %arg1 : index
+  // CHECK: spirv.ULessThanEqual
+  %7 = arith.cmpi ule, %arg0, %arg1 : index
+  // CHECK: spirv.UGreaterThan
+  %8 = arith.cmpi ugt, %arg0, %arg1 : index
+  // CHECK: spirv.UGreaterThanEqual
+  %9 = arith.cmpi uge, %arg0, %arg1 : index
+  return
+}
+
 // CHECK-LABEL: @vec1cmpi
 func.func @vec1cmpi(%arg0 : vector<1xi32>, %arg1 : vector<1xi32>) {
   // CHECK: spirv.ULessThan
   %0 = arith.cmpi ult, %arg0, %arg1 : vector<1xi32>
   // CHECK: spirv.SGreaterThan
   %1 = arith.cmpi sgt, %arg0, %arg1 : vector<1xi32>
   return
 }