[MLIR][SPIRV] Add `UnsignedOp` trait.

[ergawy]

Summary

This commit adds a new trait that can be attached to ops that have
unsigned semantics.

TODO:

• Check if other places in code can use the new attribute (possibly in this patch).
• Add a similar SignedOp attribute (in a new patch).

[antiagainst]

This looks awesome! Thanks so much! Just one request.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.h

26

Could you create another header file for hosting this trait? (I expect we will have more traits later so it's better to have a proper structure from the beginning.) :D

[ergawy]

Add a new header for Op traits.

[antiagainst]

Just two remaining nits. Feel free to land after fixing them.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOpTraits.h

2

This should be updated to match the filename.

19

Please use an empty line to separate namespace and declarations enclosed inside.

[ergawy]

Handle some formatting comments.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVArithmeticOps.td

let description = [{
     %5 = spv.SRem %2, %3 : vector<4xi32>
 
     ```
   }];
 }
 
 // -----
 
-def SPV_UDivOp : SPV_ArithmeticBinaryOp<"UDiv", SPV_Integer, []> {
+def SPV_UDivOp : SPV_ArithmeticBinaryOp<"UDiv", SPV_Integer, [UnsignedOp]> {
   let summary = "Unsigned-integer division of Operand 1 divided by Operand 2.";
 
   let description = [{
     Result Type must be a scalar or vector of integer type, whose Signedness
     operand is 0.
 
      The types of Operand 1 and Operand 2 both must be the same as Result
     Type.
let description = [{
     %5 = spv.UDiv %2, %3 : vector<4xi32>
 
     ```
   }];
 }
 
 // -----
 
-def SPV_UModOp : SPV_ArithmeticBinaryOp<"UMod", SPV_Integer> {
+def SPV_UModOp : SPV_ArithmeticBinaryOp<"UMod", SPV_Integer, [UnsignedOp]> {
   let summary = "Unsigned modulo operation of Operand 1 modulo Operand 2.";
 
   let description = [{
     Result Type must be a scalar or vector of integer type, whose Signedness
     operand is 0.
 
      The types of Operand 1 and Operand 2 both must be the same as Result
     Type.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVAtomicOps.td

let description = [{
     %0 = spv.AtomicSMin "Device" "None" %pointer, %value :
                         !spv.ptr<i32, StorageBuffer>
     ```
   }];
 }
 
 // -----
 
-def SPV_AtomicUMaxOp : SPV_AtomicUpdateWithValueOp<"AtomicUMax", []> {
+def SPV_AtomicUMaxOp : SPV_AtomicUpdateWithValueOp<"AtomicUMax", [UnsignedOp]> {
   let summary = [{
     Perform the following steps atomically with respect to any other atomic
     accesses within Scope to the same location:
   }];
 
   let description = [{
     1) load through Pointer to get an Original Value,
 
let description = [{
     %0 = spv.AtomicUMax "Device" "None" %pointer, %value :
                         !spv.ptr<i32, StorageBuffer>
     ```
   }];
 }
 
 // -----
 
-def SPV_AtomicUMinOp : SPV_AtomicUpdateWithValueOp<"AtomicUMin", []> {
+def SPV_AtomicUMinOp : SPV_AtomicUpdateWithValueOp<"AtomicUMin", [UnsignedOp]> {
   let summary = [{
     Perform the following steps atomically with respect to any other atomic
     accesses within Scope to the same location:
   }];
 
   let description = [{
     1) load through Pointer to get an Original Value,
 

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td

def InFunctionScope : PredOpTrait<
   "op must appear in a function-like op's block",
   CPred<"isNestedInFunctionLikeOp($_op.getParentOp())">>;
 
 // Check that an op can only be used within the scope of a module-like op.
 def InModuleScope : PredOpTrait<
   "op must appear in a module-like op's block",
   CPred<"isDirectInModuleLikeOp($_op.getParentOp())">>;
 
+def UnsignedOp : NativeOpTrait<"spirv::UnsignedOp">;
+
 //===----------------------------------------------------------------------===//
 // SPIR-V opcode specification
 //===----------------------------------------------------------------------===//
 
 class SPV_OpCode<string name, int val> {
   // Name used as reference to retrieve the opcode
   string opname = name;
 

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBitOps.td

let availability = [
     MaxVersion<SPV_V_1_5>,
     Extension<[]>,
     Capability<[SPV_C_Shader]>
   ];
 }
 
 // -----
 
-def SPV_BitFieldUExtractOp : SPV_BitFieldExtractOp<"BitFieldUExtract", []> {
+def SPV_BitFieldUExtractOp : SPV_BitFieldExtractOp<"BitFieldUExtract",
+                                                   [UnsignedOp]> {
   let summary = "Extract a bit field from an object, without sign extension.";
 
   let description = [{
     The semantics are the same as with OpBitFieldSExtract with the exception
     that there is no sign extension. The remaining bits of the result will
     all be 0.
 
     <!-- End of AutoGen section -->

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVCastOps.td

let description = [{
     ```
   }];
 
   let verifier = [{ return verifyCastOp(this->getOperation(), false, true); }];
 }
 
 // -----
 
-def SPV_ConvertUToFOp : SPV_CastOp<"ConvertUToF", SPV_Float, SPV_Integer, []> {
+def SPV_ConvertUToFOp : SPV_CastOp<"ConvertUToF",
+                                   SPV_Float,
+                                   SPV_Integer,
+                                   [UnsignedOp]> {
   let summary = [{
     Convert value numerically from unsigned integer to floating point.
   }];
 
   let description = [{
     Result Type must be a scalar or vector of floating-point type.
 
     Unsigned Value must be a scalar or vector of integer type.  It must have
let description = [{
     ```
   }];
 
   let verifier = [{ return verifyCastOp(this->getOperation(), false); }];
 }
 
 // -----
 
-def SPV_UConvertOp : SPV_CastOp<"UConvert", SPV_Integer, SPV_Integer, []> {
+def SPV_UConvertOp : SPV_CastOp<"UConvert",
+                                SPV_Integer,
+                                SPV_Integer,
+                                [UnsignedOp]> {
   let summary = [{
     Convert unsigned width. This is either a truncate or a zero extend.
   }];
 
   let description = [{
     Result Type must be a scalar or vector of integer type, whose Signedness
     operand is 0.
 

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td

def SPV_SelectOp : SPV_Op<"Select",
 
   let assemblyFormat = [{
     operands attr-dict `:` type($condition) `,` type($result)
   }];
 }
 
 // -----
 
-def SPV_UGreaterThanOp : SPV_LogicalBinaryOp<"UGreaterThan", SPV_Integer, []> {
+def SPV_UGreaterThanOp : SPV_LogicalBinaryOp<"UGreaterThan",
+                                             SPV_Integer,
+                                             [UnsignedOp]> {
   let summary = [{
     Unsigned-integer comparison if Operand 1 is greater than  Operand 2.
   }];
 
   let description = [{
     Result Type must be a scalar or vector of Boolean type.
 
      The type of Operand 1 and Operand 2  must be a scalar or vector of
let description = [{
     %5 = spv.UGreaterThan %2, %3 : vector<4xi32>
 
     ```
   }];
 }
 
 // -----
 
-def SPV_UGreaterThanEqualOp : SPV_LogicalBinaryOp<"UGreaterThanEqual", SPV_Integer, []> {
+def SPV_UGreaterThanEqualOp : SPV_LogicalBinaryOp<"UGreaterThanEqual",
+                                                  SPV_Integer,
+                                                  [UnsignedOp]> {
   let summary = [{
     Unsigned-integer comparison if Operand 1 is greater than or equal to
     Operand 2.
   }];
 
   let description = [{
     Result Type must be a scalar or vector of Boolean type.
 
let description = [{
     %5 = spv.UGreaterThanEqual %2, %3 : vector<4xi32>
 
     ```
   }];
 }
 
 // -----
 
-def SPV_ULessThanOp : SPV_LogicalBinaryOp<"ULessThan", SPV_Integer, []> {
+def SPV_ULessThanOp : SPV_LogicalBinaryOp<"ULessThan",
+                                          SPV_Integer,
+                                          [UnsignedOp]> {
   let summary = [{
     Unsigned-integer comparison if Operand 1 is less than Operand 2.
   }];
 
   let description = [{
     Result Type must be a scalar or vector of Boolean type.
 
      The type of Operand 1 and Operand 2  must be a scalar or vector of
let description = [{
 
     ```
   }];
 }
 
 // -----
 
 def SPV_ULessThanEqualOp :
-  SPV_LogicalBinaryOp<"ULessThanEqual", SPV_Integer, []> {
+  SPV_LogicalBinaryOp<"ULessThanEqual", SPV_Integer, [UnsignedOp]> {
   let summary = [{
     Unsigned-integer comparison if Operand 1 is less than or equal to
     Operand 2.
   }];
 
   let description = [{
     Result Type must be a scalar or vector of Boolean type.
 

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVNonUniformOps.td

let description = [{
     %1 = spv.GroupNonUniformSMin "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xi32>
     ```
   }];
 }
 
 // -----
 
 def SPV_GroupNonUniformUMaxOp :
-    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformUMax", SPV_Integer, []> {
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformUMax",
+                                    SPV_Integer,
+                                    [UnsignedOp]> {
   let summary = [{
     An unsigned integer maximum group operation of all Value operands
     contributed by active invocations in the group.
   }];
 
   let description = [{
     Result Type  must be a scalar or vector of integer type, whose
     Signedness operand is 0.
let description = [{
     %1 = spv.GroupNonUniformUMax "Subgroup" "ClusteredReduce" %vector cluster_size(%four) : vector<4xi32>
     ```
   }];
 }
 
 // -----
 
 def SPV_GroupNonUniformUMinOp :
-    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformUMin", SPV_Integer, []> {
+    SPV_GroupNonUniformArithmeticOp<"GroupNonUniformUMin",
+                                    SPV_Integer,
+                                    [UnsignedOp]> {
   let summary = [{
     An unsigned integer minimum group operation of all Value operands
     contributed by active invocations in the group.
   }];
 
   let description = [{
     Result Type  must be a scalar or vector of integer type, whose
     Signedness operand is 0.

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOpTraits.h

+//===- SPIRVOps.h - MLIR SPIR-V operation traits ----------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares C++ classes for some of operation traits in the SPIR-V
+// dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPIRV_IR_SPIRVOPTRAITS_H_
+#define MLIR_DIALECT_SPIRV_IR_SPIRVOPTRAITS_H_
+
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+namespace OpTrait {
+namespace spirv {
+
+template <typename ConcreteType>
+class UnsignedOp : public TraitBase<ConcreteType, UnsignedOp> {};
+
+} // namespace spirv
+} // namespace OpTrait
+} // namespace mlir
+
+#endif // MLIR_DIALECT_SPIRV_IR_SPIRVOPTRAITS_H_

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.h

 //===- SPIRVOps.h - MLIR SPIR-V operations ----------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the operations in the SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPIRV_IR_SPIRVOPS_H_
 #define MLIR_DIALECT_SPIRV_IR_SPIRVOPS_H_
 
+#include "mlir/Dialect/SPIRV/IR/SPIRVOpTraits.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
 class OpBuilder;
 
 namespace spirv {
 class VerCapExtAttr;
 
 // TableGen'erated operation interfaces for querying versions, extensions, and
 // capabilities.
 #include "mlir/Dialect/SPIRV/IR/SPIRVAvailability.h.inc"
 } // namespace spirv
 } // namespace mlir
 

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

 static Value shiftValue(Location loc, Value value, Value offset, Value mask,
                         int targetBits, OpBuilder &builder) {
   Type targetType = builder.getIntegerType(targetBits);
   Value result = builder.create<spirv::BitwiseAndOp>(loc, value, mask);
   return builder.create<spirv::ShiftLeftLogicalOp>(loc, targetType, result,
                                                    offset);
 }
 
-/// Returns true if the operator is operating on unsigned integers.
-/// TODO: Have a TreatOperandsAsUnsignedInteger trait and bake the information
-/// to the ops themselves.
-template <typename SPIRVOp>
-bool isUnsignedOp() {
-  return false;
-}
-
-#define CHECK_UNSIGNED_OP(SPIRVOp)                                             \
-  template <>                                                                  \
-  bool isUnsignedOp<SPIRVOp>() {                                               \
-    return true;                                                               \
-  }
-
-CHECK_UNSIGNED_OP(spirv::AtomicUMaxOp)
-CHECK_UNSIGNED_OP(spirv::AtomicUMinOp)
-CHECK_UNSIGNED_OP(spirv::BitFieldUExtractOp)
-CHECK_UNSIGNED_OP(spirv::ConvertUToFOp)
-CHECK_UNSIGNED_OP(spirv::GroupNonUniformUMaxOp)
-CHECK_UNSIGNED_OP(spirv::GroupNonUniformUMinOp)
-CHECK_UNSIGNED_OP(spirv::UConvertOp)
-CHECK_UNSIGNED_OP(spirv::UDivOp)
-CHECK_UNSIGNED_OP(spirv::UGreaterThanEqualOp)
-CHECK_UNSIGNED_OP(spirv::UGreaterThanOp)
-CHECK_UNSIGNED_OP(spirv::ULessThanEqualOp)
-CHECK_UNSIGNED_OP(spirv::ULessThanOp)
-CHECK_UNSIGNED_OP(spirv::UModOp)
-
-#undef CHECK_UNSIGNED_OP
-
 /// Returns true if the allocations of type `t` can be lowered to SPIR-V.
 static bool isAllocationSupported(MemRefType t) {
   // Currently only support workgroup local memory allocations with static
   // shape and int or float or vector of int or float element type.
   if (!(t.hasStaticShape() &&
         SPIRVTypeConverter::getMemorySpaceForStorageClass(
             spirv::StorageClass::Workgroup) == t.getMemorySpace()))
     return false;
public:
 
   LogicalResult
   matchAndRewrite(StdOp operation, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     assert(operands.size() <= 2);
     auto dstType = this->typeConverter.convertType(operation.getType());
     if (!dstType)
       return failure();
-    if (isUnsignedOp<SPIRVOp>() && dstType != operation.getType()) {
+    if (SPIRVOp::template hasTrait<OpTrait::spirv::UnsignedOp>() &&
+        dstType != operation.getType()) {
       return operation.emitError(
           "bitwidth emulation is not implemented yet on unsigned op");
     }
     rewriter.template replaceOpWithNewOp<SPIRVOp>(operation, dstType, operands);
     return success();
   }
 };
 
CmpIOpPattern::matchAndRewrite(CmpIOp cmpIOp, ArrayRef<Value> operands,
 
   Type operandType = cmpIOp.lhs().getType();
   if (isBoolScalarOrVector(operandType))
     return failure();
 
   switch (cmpIOp.getPredicate()) {
 #define DISPATCH(cmpPredicate, spirvOp)                                        \
   case cmpPredicate:                                                           \
-    if (isUnsignedOp<spirvOp>() &&                                             \
+    if (spirvOp::template hasTrait<OpTrait::spirv::UnsignedOp>() &&            \
         operandType != this->typeConverter.convertType(operandType)) {         \
       return cmpIOp.emitError(                                                 \
           "bitwidth emulation is not implemented yet on unsigned op");         \
     }                                                                          \
     rewriter.replaceOpWithNewOp<spirvOp>(cmpIOp, cmpIOp.getResult().getType(), \
                                          cmpIOpOperands.lhs(),                 \
                                          cmpIOpOperands.rhs());                \
     return success();