[mlir][StandardToSPIRV] Fix signedness issue in bitwidth emulation.

[hanchung]

Summary

Previously, after applying the mask, a negative number would convert to a
positive number because the sign flag was forgotten. This patch adds two more
shift operations to do the sign extension. This assumes that we're using two's
complement.

This patch applies sign extension unconditionally when loading a unspported integer width, and it relies the pattern to do the casting because the signedness semantic is carried by operator itself.

[hanchung]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

704

It seems that we should check the signedness semantics of dstType?

This looks like we should apply following sign extension if srcType.isSignedInteger() returns true, and shouldn't apply it if srcType.isUnsignedInteger() returns true. In case of it is signless, ie srcType.isSignlessInteger() returns true, I don't know what should we do. It seems an undef to me.

[hanchung]

Handle bitwidth emulation on unsigned operator.

[hanchung]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

212–213

I'd like to return a failure with this statement, but it would cause an error and I don't know how to fix it in the test.

'std.divi_unsigned' op requires the same type for all operands and results

 note: see current operation: %0 = "std.divi_unsigned"(%arg0, %arg1) : (i32, i32) -> i8

Does anyone have a better thought here?

704

See commit message for the explanation.

[antiagainst]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

158

I think we should have a TreatOperandsAsUnsignedInteger trait or something in the SPIR-V dialect so we can bake this information to the ops themselves. It's better structure: op semantics should belong to the op and it's better for reuse. Fine for me if you put a TODO here for now. But could you follow up on this later?

175

There are other ops like AtomicUMax? Maybe do some regex like def SPV_\w*U\w+Op inside include/Dialect/SPIRV/?

212–213

This is because we are doing partial conversion in Standard to SPIR-V pass: https://github.com/llvm/llvm-project/blob/master/mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.cpp#L43. So the conversion for the wrapper function always succeeds and that will convert the types for the input for the op. For the pattern it bails out but the conversion for the wrapper function remains materialized so we are seeing the error.

You can try to change the pass to full conversion but there may exist other places assuming partial conversion so likely they need to be fixed. (If so it would be nice to have another patch for that.)

BTW, the more conventional way to report matching failures is via return rewriter.notifyMatchFailure(...).

[hanchung]

Add a TODO.

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

158

SG and yes, I'm happy to follow up on this. Let's do it later so we can make some progress (and I also need some time to study it).

175

Yes, there are other ops. I only list the ops I found in this file. I'm not sure if we can do regex on macro definition. I checked in the doc, and couldn't find it. Let's handle it with TreatOperandsAsUnsignedInteger approach?

Would you like me to list other ops that we could find in include/Dialect/SPIRV here?

212–213

The failing tests are std-ops-to-spirv.mlir.test and std-types-to-spirv.mlir.test

I tried full conversion and also add mlir::ModuleOp and
mlir::ModuleTerminatorOp to legal op, but it still failing on legalizing func op. If I add func op to legal op, it crashes somewhere. I can fix it in a later change if I get a proper fix about it.

[hanchung]

Add check on more unsigned ops.

[hanchung]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

175

I list other ops as well.

[antiagainst]

LGTM after addressing the last two comments. :)

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

158

Yup! You can find examples for dialect specific traits in Linalg dialect: https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h#L29 and the wrapper in TableGen like https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td#L24. We can define something similar for the SPIR-V dialect.

159

Nit: have a ... trait

212–213

Ah, yes. Thanks for looking into this! It seems we should introduce a bool option to -convert-std-to-spirv to indicate whether to use partial or full conversion. Then we can split those tests (basically negative tests) that require a full conversion into a new file and test with -convert-std-to-spirv=use-full-conversion or something. For now I'm fine to return operation.emitError(...) here and check that the error message is 'std.divi_unsigned' op requires the same type for all operands and results. Returning directly is safer than what's implemented at the moment, which we still carry on the op replacement. Could you change here?

[hanchung]

address comments

[hanchung]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

158

Sounds good, I'll work on it. :)

212–213

This needs to pass -verify-diagnostics in the command and split the test out to another module. Does it look okay to you?

[rriddle]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

162

nit: static

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

static Value adjustAccessChainForBitwidth(SPIRVTypeConverter &typeConverter,
   // There are two elements if this is a 1-D tensor.
   assert(indices.size() == 2);
   indices.back() = builder.create<spirv::SDivOp>(loc, lastDim, idx);
   Type t = typeConverter.convertType(op.component_ptr().getType());
   return builder.create<spirv::AccessChainOp>(loc, t, op.base_ptr(), indices);
 }
 
 /// Returns the shifted `targetBits`-bit value with the given offset.
-Value shiftValue(Location loc, Value value, Value offset, Value mask,
+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
+
 //===----------------------------------------------------------------------===//
 // Operation conversion
 //===----------------------------------------------------------------------===//
 
 // Note that DRR cannot be used for the patterns in this file: we may need to
 // convert type along the way, which requires ConversionPattern. DRR generates
 // normal RewritePattern.
 
 namespace {
 
 /// Converts unary and binary standard operations to SPIR-V operations.
 template <typename StdOp, typename SPIRVOp>
 class UnaryAndBinaryOpPattern final : public SPIRVOpLowering<StdOp> {
 public:
   using SPIRVOpLowering<StdOp>::SPIRVOpLowering;
 
   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()) {
+      return operation.emitError(
+          "bitwidth emulation is not implemented yet on unsigned op");
+    }
     rewriter.template replaceOpWithNewOp<SPIRVOp>(operation, dstType, operands,
                                                   ArrayRef<NamedAttribute>());
     return success();
   }
 };
 
 /// Converts bitwise standard operations to SPIR-V operations. This is a special
 /// pattern other than the BinaryOpPatternPattern because if the operands are
CmpIOpPattern::matchAndRewrite(CmpIOp cmpIOp, ArrayRef<Value> operands,
   Type operandType = cmpIOp.lhs().getType();
   if (operandType.isa<IntegerType>() &&
       operandType.cast<IntegerType>().getWidth() == 1)
     return failure();
 
   switch (cmpIOp.getPredicate()) {
 #define DISPATCH(cmpPredicate, spirvOp)                                        \
   case cmpPredicate:                                                           \
+    if (isUnsignedOp<spirvOp>() &&                                             \
+        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();
 
     DISPATCH(CmpIPredicate::eq, spirv::IEqualOp);
     DISPATCH(CmpIPredicate::ne, spirv::INotEqualOp);
     DISPATCH(CmpIPredicate::slt, spirv::SLessThanOp);
IntLoadOpPattern::matchAndRewrite(LoadOp loadOp, ArrayRef<Value> operands,
   Value offset = getOffsetForBitwidth(loc, lastDim, srcBits, dstBits, rewriter);
   Value result = rewriter.create<spirv::ShiftRightArithmeticOp>(
       loc, spvLoadOp.getType(), spvLoadOp, offset);
 
   // Apply the mask to extract corresponding bits.
   Value mask = rewriter.create<spirv::ConstantOp>(
       loc, dstType, rewriter.getIntegerAttr(dstType, (1 << srcBits) - 1));
   result = rewriter.create<spirv::BitwiseAndOp>(loc, dstType, result, mask);
+
+  // Apply sign extension on the loading value unconditionally. The signedness
+  // semantic is carried in the operator itself, we relies other pattern to
+  // handle the casting.
+  IntegerAttr shiftValueAttr =
+      rewriter.getIntegerAttr(dstType, dstBits - srcBits);
+  Value shiftValue =
+      rewriter.create<spirv::ConstantOp>(loc, dstType, shiftValueAttr);
+  result = rewriter.create<spirv::ShiftLeftLogicalOp>(loc, dstType, result,
+                                                      shiftValue);
+  result = rewriter.create<spirv::ShiftRightArithmeticOp>(loc, dstType, result,
+                                                          shiftValue);
   rewriter.replaceOp(loadOp, result);
 
   assert(accessChainOp.use_empty());
   rewriter.eraseOp(accessChainOp);
 
   return success();
 }
 

mlir/test/Conversion/StandardToSPIRV/std-ops-to-spirv.mlir

-// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -convert-std-to-spirv %s -o - | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -convert-std-to-spirv -verify-diagnostics %s -o - | FileCheck %s
 
 //===----------------------------------------------------------------------===//
 // std arithmetic ops
 //===----------------------------------------------------------------------===//
 
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Int8, Int16, Int64, Float16, Float64], []>,

 // Check that types are converted to 32-bit when no special capabilities.
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [], []>,
     {max_compute_workgroup_invocations = 128 : i32,
      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
 } {
 
-// CHECK-LABEL: @int_vector234
-func @int_vector234(%arg0: vector<2xi8>, %arg1: vector<3xi16>, %arg2: vector<4xi64>) {
+// CHECK-LABEL: @int_vector23
+func @int_vector23(%arg0: vector<2xi8>, %arg1: vector<3xi16>) {
   // CHECK: spv.SDiv %{{.*}}, %{{.*}}: vector<2xi32>
   %0 = divi_signed %arg0, %arg0: vector<2xi8>
   // CHECK: spv.SRem %{{.*}}, %{{.*}}: vector<3xi32>
   %1 = remi_signed %arg1, %arg1: vector<3xi16>
-  // CHECK: spv.UDiv %{{.*}}, %{{.*}}: vector<4xi32>
-  %2 = divi_unsigned %arg2, %arg2: vector<4xi64>
   return
 }
 
 // CHECK-LABEL: @float_scalar
 func @float_scalar(%arg0: f16, %arg1: f64) {
   // CHECK: spv.FAdd %{{.*}}, %{{.*}}: f32
   %0 = addf %arg0, %arg0: f16
   // CHECK: spv.FMul %{{.*}}, %{{.*}}: f32
   %1 = mulf %arg1, %arg1: f64
   return
 }
 
 } // end module
 
 // -----
 
+// Check that types are converted to 32-bit when no special capabilities that
+// are not supported.
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [], []>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+} {
+
+// CHECK-LEBEL: @int_vector4_invalid
+func @int_vector4_invalid(%arg0: vector<4xi64>) {
+  // expected-error @+2 {{bitwidth emulation is not implemented yet on unsigned op}}
+  // expected-error @+1 {{op requires the same type for all operands and results}}
+  %0 = divi_unsigned %arg0, %arg0: vector<4xi64>
+  return
+}
+
+} // end module
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // std bit ops
 //===----------------------------------------------------------------------===//
 
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [], []>,
     {max_compute_workgroup_invocations = 128 : i32,
func @load_i8(%arg0: memref<i8>) {
   //     CHECK: %[[PTR:.+]] = spv.AccessChain %{{.+}}[%[[ZERO]], %[[QUOTIENT]]]
   //     CHECK: %[[LOAD:.+]] = spv.Load  "StorageBuffer" %[[PTR]]
   //     CHECK: %[[FOUR2:.+]] = spv.constant 4 : i32
   //     CHECK: %[[EIGHT:.+]] = spv.constant 8 : i32
   //     CHECK: %[[IDX:.+]] = spv.SMod %[[ZERO]], %[[FOUR2]] : i32
   //     CHECK: %[[BITS:.+]] = spv.IMul %[[IDX]], %[[EIGHT]] : i32
   //     CHECK: %[[VALUE:.+]] = spv.ShiftRightArithmetic %[[LOAD]], %[[BITS]] : i32, i32
   //     CHECK: %[[MASK:.+]] = spv.constant 255 : i32
-  //     CHECK: spv.BitwiseAnd %[[VALUE]], %[[MASK]] : i32
+  //     CHECK: %[[T1:.+]] = spv.BitwiseAnd %[[VALUE]], %[[MASK]] : i32
+  //     CHECK: %[[T2:.+]] = spv.constant 24 : i32
+  //     CHECK: %[[T3:.+]] = spv.ShiftLeftLogical %[[T1]], %[[T2]] : i32, i32
+  //     CHECK: spv.ShiftRightArithmetic %[[T3]], %[[T2]] : i32, i32
   %0 = load %arg0[] : memref<i8>
   return
 }
 
 // CHECK-LABEL: @load_i16
 //       CHECK: (%[[ARG0:.+]]: {{.*}}, %[[ARG1:.+]]: i32)
 func @load_i16(%arg0: memref<10xi16>, %index : index) {
   //     CHECK: %[[ONE:.+]] = spv.constant 1 : i32
   //     CHECK: %[[FLAT_IDX:.+]] = spv.IMul %[[ONE]], %[[ARG1]] : i32
   //     CHECK: %[[ZERO:.+]] = spv.constant 0 : i32
   //     CHECK: %[[TWO1:.+]] = spv.constant 2 : i32
   //     CHECK: %[[QUOTIENT:.+]] = spv.SDiv %[[FLAT_IDX]], %[[TWO1]] : i32
   //     CHECK: %[[PTR:.+]] = spv.AccessChain %{{.+}}[%[[ZERO]], %[[QUOTIENT]]]
   //     CHECK: %[[LOAD:.+]] = spv.Load  "StorageBuffer" %[[PTR]]
   //     CHECK: %[[TWO2:.+]] = spv.constant 2 : i32
   //     CHECK: %[[SIXTEEN:.+]] = spv.constant 16 : i32
   //     CHECK: %[[IDX:.+]] = spv.SMod %[[FLAT_IDX]], %[[TWO2]] : i32
   //     CHECK: %[[BITS:.+]] = spv.IMul %[[IDX]], %[[SIXTEEN]] : i32
   //     CHECK: %[[VALUE:.+]] = spv.ShiftRightArithmetic %[[LOAD]], %[[BITS]] : i32, i32
   //     CHECK: %[[MASK:.+]] = spv.constant 65535 : i32
-  //     CHECK: spv.BitwiseAnd %[[VALUE]], %[[MASK]] : i32
+  //     CHECK: %[[T1:.+]] = spv.BitwiseAnd %[[VALUE]], %[[MASK]] : i32
+  //     CHECK: %[[T2:.+]] = spv.constant 16 : i32
+  //     CHECK: %[[T3:.+]] = spv.ShiftLeftLogical %[[T1]], %[[T2]] : i32, i32
+  //     CHECK: spv.ShiftRightArithmetic %[[T3]], %[[T2]] : i32, i32
   %0 = load %arg0[%index] : memref<10xi16>
   return
 }
 
 // CHECK-LABEL: @load_i32
 func @load_i32(%arg0: memref<i32>) {
   // CHECK-NOT: spv.SDiv
   //     CHECK: spv.Load
func @load_i8(%arg0: memref<i8>) {
   //     CHECK: %[[PTR:.+]] = spv.AccessChain %{{.+}}[%[[ZERO]], %[[QUOTIENT]]]
   //     CHECK: %[[LOAD:.+]] = spv.Load  "StorageBuffer" %[[PTR]]
   //     CHECK: %[[FOUR2:.+]] = spv.constant 4 : i32
   //     CHECK: %[[EIGHT:.+]] = spv.constant 8 : i32
   //     CHECK: %[[IDX:.+]] = spv.SMod %[[ZERO]], %[[FOUR2]] : i32
   //     CHECK: %[[BITS:.+]] = spv.IMul %[[IDX]], %[[EIGHT]] : i32
   //     CHECK: %[[VALUE:.+]] = spv.ShiftRightArithmetic %[[LOAD]], %[[BITS]] : i32, i32
   //     CHECK: %[[MASK:.+]] = spv.constant 255 : i32
-  //     CHECK: spv.BitwiseAnd %[[VALUE]], %[[MASK]] : i32
+  //     CHECK: %[[T1:.+]] = spv.BitwiseAnd %[[VALUE]], %[[MASK]] : i32
+  //     CHECK: %[[T2:.+]] = spv.constant 24 : i32
+  //     CHECK: %[[T3:.+]] = spv.ShiftLeftLogical %[[T1]], %[[T2]] : i32, i32
+  //     CHECK: spv.ShiftRightArithmetic %[[T3]], %[[T2]] : i32, i32
   %0 = load %arg0[] : memref<i8>
   return
 }
 
 // CHECK-LABEL: @load_i16
 func @load_i16(%arg0: memref<i16>) {
   // CHECK-NOT: spv.SDiv
   //     CHECK: spv.Load