[mlir][spirv] arith::RemSIOp OpenCL lowering

[Hardcode84]

@antiagainst actually, I have an issue with this Shader/Kernel ops selection approach. I am using AtomicFAddEXT in kernel which depends on SPV_C_AtomicFloat32AddEXT which have implies = [SPV_C_Shader];. Intel L0 drivers can successfully compile and run such kernels but it breaks this entire approach. Do you have any ideas how to handle this properly? As a quick workaround I can manually patch AtomicFloat32AddEXT to remove implies.

[antiagainst]

In D114524#3151394, @Hardcode84 wrote:

@antiagainst actually, I have an issue with this Shader/Kernel ops selection approach. I am using AtomicFAddEXT in kernel which depends on SPV_C_AtomicFloat32AddEXT which have implies = [SPV_C_Shader];. Intel L0 drivers can successfully compile and run such kernels but it breaks this entire approach. Do you have any ideas how to handle this properly? As a quick workaround I can manually patch AtomicFloat32AddEXT to remove implies.

I think it's a bug in the upstream JSON grammar. From the spec (https://github.com/KhronosGroup/SPIRV-Registry/blob/main/extensions/EXT/SPV_EXT_shader_atomic_float_add.asciidoc) AtomicFloat32AddEXT/AtomicFloat64AddEXT has no implicit implied capabilities (the table there is malformed though), but the commit to update JSON grammar has (https://github.com/KhronosGroup/SPIRV-Headers/commit/7083cb52e1812201c2d5641306fcbca639c9560f). I'd suggest you to open a pull request to SPIRV-Headers repro to fix this. In the meanwhile, feel free to manually drop the implies in the td files.

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

struct ConstantScalarOpPattern final
     : public OpConversionPattern<arith::ConstantOp> {
   using OpConversionPattern<arith::ConstantOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::ConstantOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
-/// Converts arith.remsi to SPIR-V ops.
+/// Converts arith.remsi to GLSL SPIR-V ops.
 ///
 /// This cannot be merged into the template unary/binary pattern due to Vulkan
 /// restrictions over spv.SRem and spv.SMod.
-struct RemSIOpPattern final : public OpConversionPattern<arith::RemSIOp> {
+struct RemSIOpGLSLPattern final : public OpConversionPattern<arith::RemSIOp> {
+  using OpConversionPattern<arith::RemSIOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(arith::RemSIOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
+/// Converts arith.remsi to OpenCL SPIR-V ops.
+struct RemSIOpOCLPattern final : public OpConversionPattern<arith::RemSIOp> {
   using OpConversionPattern<arith::RemSIOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::RemSIOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 /// Converts bitwise operations to SPIR-V operations. This is a special pattern
auto dstAttr =
       convertIntegerAttr(srcAttr, dstType.cast<IntegerType>(), rewriter);
   if (!dstAttr)
     return failure();
   rewriter.replaceOpWithNewOp<spirv::ConstantOp>(constOp, dstType, dstAttr);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// RemSIOpPattern
+// RemSIOpGLSLPattern
 //===----------------------------------------------------------------------===//
 
 /// Returns signed remainder for `lhs` and `rhs` and lets the result follow
 /// the sign of `signOperand`.
 ///
 /// Note that this is needed for Vulkan. Per the Vulkan's SPIR-V environment
 /// spec, "for the OpSRem and OpSMod instructions, if either operand is negative
 /// the result is undefined."  So we cannot directly use spv.SRem/spv.SMod
 /// if either operand can be negative. Emulate it via spv.UMod.
+template <typename SignedAbsOp>
 static Value emulateSignedRemainder(Location loc, Value lhs, Value rhs,
                                     Value signOperand, OpBuilder &builder) {
   assert(lhs.getType() == rhs.getType());
   assert(lhs == signOperand || rhs == signOperand);
 
   Type type = lhs.getType();
 
   // Calculate the remainder with spv.UMod.
-  Value lhsAbs = builder.create<spirv::GLSLSAbsOp>(loc, type, lhs);
-  Value rhsAbs = builder.create<spirv::GLSLSAbsOp>(loc, type, rhs);
+  Value lhsAbs = builder.create<SignedAbsOp>(loc, type, lhs);
+  Value rhsAbs = builder.create<SignedAbsOp>(loc, type, rhs);
   Value abs = builder.create<spirv::UModOp>(loc, lhsAbs, rhsAbs);
 
   // Fix the sign.
   Value isPositive;
   if (lhs == signOperand)
     isPositive = builder.create<spirv::IEqualOp>(loc, lhs, lhsAbs);
   else
     isPositive = builder.create<spirv::IEqualOp>(loc, rhs, rhsAbs);
   Value absNegate = builder.create<spirv::SNegateOp>(loc, type, abs);
   return builder.create<spirv::SelectOp>(loc, type, isPositive, abs, absNegate);
 }
 
 LogicalResult
-RemSIOpPattern::matchAndRewrite(arith::RemSIOp op, OpAdaptor adaptor,
+RemSIOpGLSLPattern::matchAndRewrite(arith::RemSIOp op, OpAdaptor adaptor,
+                                    ConversionPatternRewriter &rewriter) const {
+  Value result = emulateSignedRemainder<spirv::GLSLSAbsOp>(
+      op.getLoc(), adaptor.getOperands()[0], adaptor.getOperands()[1],
+      adaptor.getOperands()[0], rewriter);
+  rewriter.replaceOp(op, result);
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// RemSIOpOCLPattern
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+RemSIOpOCLPattern::matchAndRewrite(arith::RemSIOp op, OpAdaptor adaptor,
                                    ConversionPatternRewriter &rewriter) const {
-  Value result = emulateSignedRemainder(op.getLoc(), adaptor.getOperands()[0],
-                                        adaptor.getOperands()[1],
-                                        adaptor.getOperands()[0], rewriter);
+  Value result = emulateSignedRemainder<spirv::OCLSAbsOp>(
+      op.getLoc(), adaptor.getOperands()[0], adaptor.getOperands()[1],
   rewriter.replaceOp(op, result);
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // BitwiseOpPattern
 //===----------------------------------------------------------------------===//
patterns.add<
     ConstantCompositeOpPattern,
     ConstantScalarOpPattern,
     spirv::UnaryAndBinaryOpPattern<arith::AddIOp, spirv::IAddOp>,
     spirv::UnaryAndBinaryOpPattern<arith::SubIOp, spirv::ISubOp>,
     spirv::UnaryAndBinaryOpPattern<arith::MulIOp, spirv::IMulOp>,
     spirv::UnaryAndBinaryOpPattern<arith::DivUIOp, spirv::UDivOp>,
     spirv::UnaryAndBinaryOpPattern<arith::DivSIOp, spirv::SDivOp>,
     spirv::UnaryAndBinaryOpPattern<arith::RemUIOp, spirv::UModOp>,
-    RemSIOpPattern,
+    RemSIOpGLSLPattern, RemSIOpOCLPattern,
     BitwiseOpPattern<arith::AndIOp, spirv::LogicalAndOp, spirv::BitwiseAndOp>,
     BitwiseOpPattern<arith::OrIOp, spirv::LogicalOrOp, spirv::BitwiseOrOp>,
     XOrIOpLogicalPattern, XOrIOpBooleanPattern,
     spirv::UnaryAndBinaryOpPattern<arith::ShLIOp, spirv::ShiftLeftLogicalOp>,
     spirv::UnaryAndBinaryOpPattern<arith::ShRUIOp, spirv::ShiftRightLogicalOp>,
     spirv::UnaryAndBinaryOpPattern<arith::ShRSIOp, spirv::ShiftRightArithmeticOp>,
     spirv::UnaryAndBinaryOpPattern<arith::NegFOp, spirv::FNegateOp>,
     spirv::UnaryAndBinaryOpPattern<arith::AddFOp, spirv::FAddOp>,

mlir/test/Conversion/ArithmeticToSPIRV/arithmetic-to-spirv.mlir

func @scalar_srem(%lhs: i32, %rhs: i32) {
   // CHECK:      %{{.+}} = spv.Select %[[POS]], %[[ABS]], %[[NEG]] : i1, i32
   %0 = arith.remsi %lhs, %rhs: i32
   return
 }
 
 // Check float unary operation conversions.
 // CHECK-LABEL: @float32_unary_scalar
 func @float32_unary_scalar(%arg0: f32) {
-  // CHECK: spv.GLSL.FAbs %{{.*}}: f32
-  %0 = math.abs %arg0 : f32
-  // CHECK: spv.GLSL.Ceil %{{.*}}: f32
-  %1 = math.ceil %arg0 : f32
   // CHECK: spv.FNegate %{{.*}}: f32
-  %5 = arith.negf %arg0 : f32
-  // CHECK: spv.GLSL.Floor %{{.*}}: f32
-  %10 = math.floor %arg0 : f32
+  %0 = arith.negf %arg0 : f32
   return
 }
 
 // Check float binary operation conversions.
 // CHECK-LABEL: @float32_binary_scalar
 func @float32_binary_scalar(%lhs: f32, %rhs: f32) {
   // CHECK: spv.FAdd %{{.*}}, %{{.*}}: f32
   %0 = arith.addf %lhs, %rhs: f32

 // CHECK-LABEL: @sitofp
 func @sitofp(%arg0 : i64) -> f64 {
   // CHECK: spv.ConvertSToF %{{.*}} : i32 to f32
   %0 = arith.sitofp %arg0 : i64 to f64
   return %0: f64
 }
 
 } // end module
+
+// -----
+
+// Check OpenCL lowering of arith.remsi
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Int16, Kernel], []>, {}>
+} {
+
+// CHECK-LABEL: @scalar_srem
+// CHECK-SAME: (%[[LHS:.+]]: i32, %[[RHS:.+]]: i32)
+func @scalar_srem(%lhs: i32, %rhs: i32) {
+  // CHECK: %[[LABS:.+]] = spv.OCL.s_abs %[[LHS]] : i32
+  // CHECK: %[[RABS:.+]] = spv.OCL.s_abs %[[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: i32
+  return
+}
+
+// CHECK-LABEL: @vector_srem
+// CHECK-SAME: (%[[LHS:.+]]: vector<3xi16>, %[[RHS:.+]]: vector<3xi16>)
+func @vector_srem(%arg0: vector<3xi16>, %arg1: vector<3xi16>) {
+  // CHECK: %[[LABS:.+]] = spv.OCL.s_abs %[[LHS]] : vector<3xi16>
+  // CHECK: %[[RABS:.+]] = spv.OCL.s_abs %[[RHS]] : vector<3xi16>
+  // CHECK:  %[[ABS:.+]] = spv.UMod %[[LABS]], %[[RABS]] : vector<3xi16>
+  // CHECK:  %[[POS:.+]] = spv.IEqual %[[LHS]], %[[LABS]] : vector<3xi16>
+  // CHECK:  %[[NEG:.+]] = spv.SNegate %[[ABS]] : vector<3xi16>
+  // CHECK:      %{{.+}} = spv.Select %[[POS]], %[[ABS]], %[[NEG]] : vector<3xi1>, vector<3xi16>
+  %0 = arith.remsi %arg0, %arg1: vector<3xi16>
+  return
+}
+
+} // end module

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

func @float32_unary_scalar(%arg0: f32) {
   // CHECK: spv.GLSL.InverseSqrt %{{.*}}: f32
   %4 = math.rsqrt %arg0 : f32
   // CHECK: spv.GLSL.Sqrt %{{.*}}: f32
   %5 = math.sqrt %arg0 : f32
   // CHECK: spv.GLSL.Tanh %{{.*}}: f32
   %6 = math.tanh %arg0 : f32
   // CHECK: spv.GLSL.Sin %{{.*}}: f32
   %7 = math.sin %arg0 : f32
+  // CHECK: spv.GLSL.FAbs %{{.*}}: f32
+  %8 = math.abs %arg0 : f32
+  // CHECK: spv.GLSL.Ceil %{{.*}}: f32
+  %9 = math.ceil %arg0 : f32
+  // CHECK: spv.GLSL.Floor %{{.*}}: f32
+  %10 = math.floor %arg0 : f32
   return
 }
 
 // CHECK-LABEL: @float32_unary_vector
 func @float32_unary_vector(%arg0: vector<3xf32>) {
   // CHECK: spv.GLSL.Cos %{{.*}}: vector<3xf32>
   %0 = math.cos %arg0 : vector<3xf32>
   // CHECK: spv.GLSL.Exp %{{.*}}: vector<3xf32>

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

func @float32_unary_scalar(%arg0: f32) {
   // CHECK: spv.OCL.rsqrt %{{.*}}: f32
   %4 = math.rsqrt %arg0 : f32
   // CHECK: spv.OCL.sqrt %{{.*}}: f32
   %5 = math.sqrt %arg0 : f32
   // CHECK: spv.OCL.tanh %{{.*}}: f32
   %6 = math.tanh %arg0 : f32
   // CHECK: spv.OCL.sin %{{.*}}: f32
   %7 = math.sin %arg0 : f32
+  // CHECK: spv.OCL.fabs %{{.*}}: f32
+  %8 = math.abs %arg0 : f32
+  // CHECK: spv.OCL.ceil %{{.*}}: f32
+  %9 = math.ceil %arg0 : f32
+  // CHECK: spv.OCL.floor %{{.*}}: f32
+  %10 = math.floor %arg0 : f32
   return
 }
 
 // CHECK-LABEL: @float32_unary_vector
 func @float32_unary_vector(%arg0: vector<3xf32>) {
   // CHECK: spv.OCL.cos %{{.*}}: vector<3xf32>
   %0 = math.cos %arg0 : vector<3xf32>
   // CHECK: spv.OCL.exp %{{.*}}: vector<3xf32>