[Mlir-commits] [mlir] 38e802a - [mlir][spirv] Allow converting from index type in unsigned ops

[Lei Zhang]

Author: Lei Zhang
Date: 2022-04-27T10:13:50-04:00
New Revision: 38e802a09de645b54c3fc983d692a1436bfc7d13

URL: https://github.com/llvm/llvm-project/commit/38e802a09de645b54c3fc983d692a1436bfc7d13
DIFF: https://github.com/llvm/llvm-project/commit/38e802a09de645b54c3fc983d692a1436bfc7d13.diff

LOG: [mlir][spirv] Allow converting from index type in unsigned ops

`index` type is converted to `i32` in SPIR-V. This is fine to
support for all signed/unsigned ops.

Reviewed By: hanchung

Differential Revision: https://reviews.llvm.org/D124451

Added: 
    

Modified: 
    mlir/lib/Conversion/SPIRVCommon/Pattern.h
    mlir/test/Conversion/ArithmeticToSPIRV/arithmetic-to-spirv.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/SPIRVCommon/Pattern.h b/mlir/lib/Conversion/SPIRVCommon/Pattern.h
index c70009fcc23e..5d32fa81cc6b 100644
--- a/mlir/lib/Conversion/SPIRVCommon/Pattern.h
+++ b/mlir/lib/Conversion/SPIRVCommon/Pattern.h
@@ -30,7 +30,7 @@ class ElementwiseOpPattern final : public OpConversionPattern<Op> {
     if (!dstType)
       return failure();
     if (SPIRVOp::template hasTrait<OpTrait::spirv::UnsignedOp>() &&
-        dstType != op.getType()) {
+        !op.getType().isIndex() && dstType != op.getType()) {
       return op.emitError(
           "bitwidth emulation is not implemented yet on unsigned op");
     }

diff  --git a/mlir/test/Conversion/ArithmeticToSPIRV/arithmetic-to-spirv.mlir b/mlir/test/Conversion/ArithmeticToSPIRV/arithmetic-to-spirv.mlir
index 8da0ae19bb38..b40d01cb0b8f 100644
--- a/mlir/test/Conversion/ArithmeticToSPIRV/arithmetic-to-spirv.mlir
+++ b/mlir/test/Conversion/ArithmeticToSPIRV/arithmetic-to-spirv.mlir
@@ -27,9 +27,9 @@ func.func @int32_scalar(%lhs: i32, %rhs: i32) {
   return
 }
 
-// CHECK-LABEL: @scalar_srem
+// CHECK-LABEL: @int32_scalar_srem
 // CHECK-SAME: (%[[LHS:.+]]: i32, %[[RHS:.+]]: i32)
-func.func @scalar_srem(%lhs: i32, %rhs: i32) {
+func.func @int32_scalar_srem(%lhs: i32, %rhs: i32) {
   // CHECK: %[[LABS:.+]] = spv.GLSL.SAbs %[[LHS]] : i32
   // CHECK: %[[RABS:.+]] = spv.GLSL.SAbs %[[RHS]] : i32
   // CHECK:  %[[ABS:.+]] = spv.UMod %[[LABS]], %[[RABS]] : i32
@@ -40,6 +40,36 @@ func.func @scalar_srem(%lhs: i32, %rhs: i32) {
   return
 }
 
+// CHECK-LABEL: @index_scalar
+func.func @index_scalar(%lhs: index, %rhs: index) {
+  // CHECK: spv.IAdd %{{.*}}, %{{.*}}: i32
+  %0 = arith.addi %lhs, %rhs: index
+  // CHECK: spv.ISub %{{.*}}, %{{.*}}: i32
+  %1 = arith.subi %lhs, %rhs: index
+  // CHECK: spv.IMul %{{.*}}, %{{.*}}: i32
+  %2 = arith.muli %lhs, %rhs: index
+  // CHECK: spv.SDiv %{{.*}}, %{{.*}}: i32
+  %3 = arith.divsi %lhs, %rhs: index
+  // CHECK: spv.UDiv %{{.*}}, %{{.*}}: i32
+  %4 = arith.divui %lhs, %rhs: index
+  // CHECK: spv.UMod %{{.*}}, %{{.*}}: i32
+  %5 = arith.remui %lhs, %rhs: index
+  return
+}
+
+// CHECK-LABEL: @index_scalar_srem
+// CHECK-SAME: (%[[LHS:.+]]: i32, %[[RHS:.+]]: i32)
+func.func @index_scalar_srem(%lhs: index, %rhs: index) {
+  // CHECK: %[[LABS:.+]] = spv.GLSL.SAbs %[[LHS]] : i32
+  // CHECK: %[[RABS:.+]] = spv.GLSL.SAbs %[[RHS]] : i32
+  // CHECK:  %[[ABS:.+]] = spv.UMod %[[LABS]], %[[RABS]] : i32
+  // CHECK:  %[[POS:.+]] = spv.IEqual %[[LHS]], %[[LABS]] : i32
+  // CHECK:  %[[NEG:.+]] = spv.SNegate %[[ABS]] : i32
+  // CHECK:      %{{.+}} = spv.Select %[[POS]], %[[ABS]], %[[NEG]] : i1, i32
+  %0 = arith.remsi %lhs, %rhs: index
+  return
+}
+
 // Check float unary operation conversions.
 // CHECK-LABEL: @float32_unary_scalar
 func.func @float32_unary_scalar(%arg0: f32) {