[mlir][spirv] Improve integer cast during type conversion

[antiagainst]

Summary

In SPIR-V, the capabilities for storage and compute are separate.
We have good handling of the storage side in general via MemRef
type conversion and various memref dialect ops.

Once the value was loaded properly, if the compute capability is
supported directly, we don't need to emulate like the storage side
with int32. However, we do need to make sure casting ops are
properly inserted to chain the flow to go back to the original
bitwidth.

Right now that is done in the each individual pattern directly,
which put lots of pressure that shouldn't be on the patterns and
causes duplication and trickiness w.r.t. capability check and such.

Instead, we should handle such casting within the SPIR-V conversion
framework using addSourceMaterialization, where we can check with
the target environment to make sure the corresponding compute
capability is allowed and then we can materialize and SPIR-V casting
op.

Along the way, we can drop all the duplicated cast materialization
registration in various places.

[kuhar]

LGTM

mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp

void runOnOperation() override {
 
     SPIRVConversionOptions options;
     options.emulateLT32BitScalarTypes = this->emulateLT32BitScalarTypes;
     options.enableFastMathMode = this->enableFastMath;
     SPIRVTypeConverter typeConverter(targetAttr, options);
 
     // Use UnrealizedConversionCast as the bridge so that we don't need to pull
     // in patterns for other dialects.
-    auto addUnrealizedCast = [](OpBuilder &builder, Type type,
-                                ValueRange inputs, Location loc) {
-      auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
-      return std::optional<Value>(cast.getResult(0));
-    };
-    typeConverter.addSourceMaterialization(addUnrealizedCast);
-    typeConverter.addTargetMaterialization(addUnrealizedCast);
     target->addLegalOp<UnrealizedConversionCastOp>();
 
     // Fail hard when there are any remaining 'arith' ops.
     target->addIllegalDialect<arith::ArithDialect>();
 
     RewritePatternSet patterns(&getContext());
     arith::populateArithToSPIRVPatterns(typeConverter, patterns);
 

mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRVPass.cpp

void runOnOperation() override {
     std::unique_ptr<ConversionTarget> target =
         SPIRVConversionTarget::get(targetAttr);
 
     SPIRVConversionOptions options;
     SPIRVTypeConverter typeConverter(targetAttr, options);
 
     // Use UnrealizedConversionCast as the bridge so that we don't need to pull
     // in patterns for other dialects.
-    auto addUnrealizedCast = [](OpBuilder &builder, Type type,
-                                ValueRange inputs, Location loc) {
-      auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
-      return std::optional<Value>(cast.getResult(0));
-    };
-    typeConverter.addSourceMaterialization(addUnrealizedCast);
-    typeConverter.addTargetMaterialization(addUnrealizedCast);
     target->addLegalOp<UnrealizedConversionCastOp>();
 
     RewritePatternSet patterns(context);
     populateComplexToSPIRVPatterns(typeConverter, patterns);
 
     if (failed(applyPartialConversion(op, *target, std::move(patterns))))
       return signalPassFailure();
   }
 };
 } // namespace

mlir/lib/Conversion/MathToSPIRV/MathToSPIRVPass.cpp

void ConvertMathToSPIRVPass::runOnOperation() {
   auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
   std::unique_ptr<ConversionTarget> target =
       SPIRVConversionTarget::get(targetAttr);
 
   SPIRVTypeConverter typeConverter(targetAttr);
 
   // Use UnrealizedConversionCast as the bridge so that we don't need to pull
   // in patterns for other dialects.
-  auto addUnrealizedCast = [](OpBuilder &builder, Type type, ValueRange inputs,
-                              Location loc) {
-    auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
-    return std::optional<Value>(cast.getResult(0));
-  };
-  typeConverter.addSourceMaterialization(addUnrealizedCast);
-  typeConverter.addTargetMaterialization(addUnrealizedCast);
   target->addLegalOp<UnrealizedConversionCastOp>();
 
   RewritePatternSet patterns(context);
   populateMathToSPIRVPatterns(typeConverter, patterns);
 
   if (failed(applyPartialConversion(op, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<>> mlir::createConvertMathToSPIRVPass() {
   return std::make_unique<ConvertMathToSPIRVPass>();
 }

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

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);
 
-  if (isBool) {
-    dstType = typeConverter.convertType(loadOp.getType());
-    mask = spirv::ConstantOp::getOne(result.getType(), loc, rewriter);
-    result = rewriter.create<spirv::IEqualOp>(loc, result, mask);
-  } else if (result.getType().getIntOrFloatBitWidth() !=
-             static_cast<unsigned>(dstBits)) {
-    result = rewriter.create<spirv::SConvertOp>(loc, dstType, result);
-  }
   rewriter.replaceOp(loadOp, result);
 
   assert(accessChainOp.use_empty());
   rewriter.eraseOp(accessChainOp);
 
   return success();
 }
 

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.cpp

std::unique_ptr<ConversionTarget> target =
       SPIRVConversionTarget::get(targetAttr);
 
   SPIRVConversionOptions options;
   options.boolNumBits = this->boolNumBits;
   SPIRVTypeConverter typeConverter(targetAttr, options);
 
   // Use UnrealizedConversionCast as the bridge so that we don't need to pull in
   // patterns for other dialects.
-  auto addUnrealizedCast = [](OpBuilder &builder, Type type, ValueRange inputs,
-                              Location loc) {
-    auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
-    return std::optional<Value>(cast.getResult(0));
-  };
-  typeConverter.addSourceMaterialization(addUnrealizedCast);
-  typeConverter.addTargetMaterialization(addUnrealizedCast);
   target->addLegalOp<UnrealizedConversionCastOp>();
 
   RewritePatternSet patterns(context);
   populateMemRefToSPIRVPatterns(typeConverter, patterns);
 
   if (failed(applyPartialConversion(op, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<>> mlir::createConvertMemRefToSPIRVPass() {
   return std::make_unique<ConvertMemRefToSPIRVPass>();
 }

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRVPass.cpp

void ConvertVectorToSPIRVPass::runOnOperation() {
   auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
   std::unique_ptr<ConversionTarget> target =
       SPIRVConversionTarget::get(targetAttr);
 
   SPIRVTypeConverter typeConverter(targetAttr);
 
   // Use UnrealizedConversionCast as the bridge so that we don't need to pull in
   // patterns for other dialects.
-  auto addUnrealizedCast = [](OpBuilder &builder, Type type, ValueRange inputs,
-                              Location loc) {
-    auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
-    return std::optional<Value>(cast.getResult(0));
-  };
-  typeConverter.addSourceMaterialization(addUnrealizedCast);
-  typeConverter.addTargetMaterialization(addUnrealizedCast);
   target->addLegalOp<UnrealizedConversionCastOp>();
 
   RewritePatternSet patterns(context);
   populateVectorToSPIRVPatterns(typeConverter, patterns);
 
   if (failed(applyPartialConversion(op, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<>> mlir::createConvertVectorToSPIRVPass() {
   return std::make_unique<ConvertVectorToSPIRVPass>();
 }

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

static Type convertMemrefType(const spirv::TargetEnv &targetEnv,
   int64_t arrayElemCount = llvm::divideCeil(*memrefSize, *arrayElemSize);
   int64_t stride = needsExplicitLayout(storageClass) ? *arrayElemSize : 0;
   auto arrayType = spirv::ArrayType::get(arrayElemType, arrayElemCount, stride);
   if (targetEnv.allows(spirv::Capability::Kernel))
     return spirv::PointerType::get(arrayType, storageClass);
   return wrapInStructAndGetPointer(arrayType, storageClass);
 }
 
+//===----------------------------------------------------------------------===//
+// Type casting materialization
+//===----------------------------------------------------------------------===//
+
+/// Converts the given `inputs` to the original source `type` considering the
+/// `targetEnv`'s capabilities.
+///
+/// This function is meant to be used for source materialization in type
+/// converters. When the type converter needs to materialize a cast op back
+/// to some original source type, we need to check whether the original source
+/// type is supported in the target environment. If so, we can insert legal
+/// SPIR-V cast ops accordingly.
+///
+/// Note that in SPIR-V the capabilities for storage and compute are separate.
+/// This function is meant to handle the **compute** side; so it does not
+/// involve storage classes in its logic. The storage side is expected to be
+/// handled by MemRef conversion logic.
+std::optional<Value> castToSourceType(const spirv::TargetEnv &targetEnv,
+                                      OpBuilder &builder, Type type,
+                                      ValueRange inputs, Location loc) {
+  // We can only cast one value in SPIR-V.
+  if (inputs.size() != 1) {
+    auto castOp = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
+    return castOp.getResult(0);
+  }
+  Value input = inputs.front();
+
+  // Only support integer types for now. Floating point types to be implemented.
+  if (!isa<IntegerType>(type)) {
+    auto castOp = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
+    return castOp.getResult(0);
+  }
+  auto inputType = cast<IntegerType>(input.getType());
+
+  auto scalarType = dyn_cast<spirv::ScalarType>(type);
+  if (!scalarType) {
+    auto castOp = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
+    return castOp.getResult(0);
+  }
+
+  // Only support source type with a smaller bitwidth. This would mean we are
+  // truncating to go back so we don't need to worry about the signedness.
+  // For extension, we cannot have enough signal here to decide which op to use.
+  if (inputType.getIntOrFloatBitWidth() < scalarType.getIntOrFloatBitWidth()) {
+    auto castOp = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
+    return castOp.getResult(0);
+  }
+
+  // Boolean values would need to use different ops than normal integer values.
+  if (type.isInteger(1)) {
+    Value one = spirv::ConstantOp::getOne(inputType, loc, builder);
+    return builder.create<spirv::IEqualOp>(loc, input, one);
+  }
+
+  // Check that the source integer type is supported by the environment.
+  SmallVector<ArrayRef<spirv::Extension>, 1> exts;
+  SmallVector<ArrayRef<spirv::Capability>, 2> caps;
+  scalarType.getExtensions(exts);
+  scalarType.getCapabilities(caps);
+  if (failed(checkCapabilityRequirements(type, targetEnv, caps)) ||
+      failed(checkExtensionRequirements(type, targetEnv, exts))) {
+    auto castOp = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
+    return castOp.getResult(0);
+  }
+
+  // We've already made sure this is truncating previously, so we don't need to
+  // care about signedness here. Still try to use a corresponding op for better
+  // consistency though.
+  if (type.isSignedInteger()) {
+    return builder.create<spirv::SConvertOp>(loc, type, input);
+  }
+  return builder.create<spirv::UConvertOp>(loc, type, input);
+}
+
+//===----------------------------------------------------------------------===//
+// SPIRVTypeConverter
+//===----------------------------------------------------------------------===//
+
 SPIRVTypeConverter::SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr,
                                        const SPIRVConversionOptions &options)
     : targetEnv(targetAttr), options(options) {
   // Add conversions. The order matters here: later ones will be tried earlier.
 
   // Allow all SPIR-V dialect specific types. This assumes all builtin types
   // adopted in the SPIR-V dialect (i.e., IntegerType, FloatType, VectorType)
   // were tried before.
SPIRVTypeConverter::SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr,
 
   addConversion([this](TensorType tensorType) {
     return convertTensorType(this->targetEnv, this->options, tensorType);
   });
 
   addConversion([this](MemRefType memRefType) {
     return convertMemrefType(this->targetEnv, this->options, memRefType);
   });
+
+  // Register some last line of defense casting logic.
+  addSourceMaterialization(
+      [this](OpBuilder &builder, Type type, ValueRange inputs, Location loc) {
+        return castToSourceType(this->targetEnv, builder, type, inputs, loc);
+      });
+  addTargetMaterialization([](OpBuilder &builder, Type type, ValueRange inputs,
+                              Location loc) {
+    auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
+    return std::optional<Value>(cast.getResult(0));
+  });
 }
 
 //===----------------------------------------------------------------------===//
 // func::FuncOp Conversion Patterns
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// A pattern for rewriting function signature to convert arguments of functions

mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir

func.func @load_i8(%arg0: memref<i8, #spirv.storage_class<StorageBuffer>>) -> i8 {
   //     CHECK: %[[EIGHT:.+]] = spirv.Constant 8 : i32
   //     CHECK: %[[IDX:.+]] = spirv.UMod %[[ZERO]], %[[FOUR]] : i32
   //     CHECK: %[[BITS:.+]] = spirv.IMul %[[IDX]], %[[EIGHT]] : i32
   //     CHECK: %[[VALUE:.+]] = spirv.ShiftRightArithmetic %[[LOAD]], %[[BITS]] : i32, i32
   //     CHECK: %[[MASK:.+]] = spirv.Constant 255 : i32
   //     CHECK: %[[T1:.+]] = spirv.BitwiseAnd %[[VALUE]], %[[MASK]] : i32
   //     CHECK: %[[T2:.+]] = spirv.Constant 24 : i32
   //     CHECK: %[[T3:.+]] = spirv.ShiftLeftLogical %[[T1]], %[[T2]] : i32, i32
-  //     CHECK: spirv.ShiftRightArithmetic %[[T3]], %[[T2]] : i32, i32
+  //     CHECK: %[[SR:.+]] = spirv.ShiftRightArithmetic %[[T3]], %[[T2]] : i32, i32
+  //     CHECK: builtin.unrealized_conversion_cast %[[SR]]
   %0 = memref.load %arg0[] : memref<i8, #spirv.storage_class<StorageBuffer>>
   return %0 : i8
 }
 
 // CHECK-LABEL: @load_i16
 //       CHECK: (%[[ARG0:.+]]: {{.*}}, %[[ARG1:.+]]: index)
 func.func @load_i16(%arg0: memref<10xi16, #spirv.storage_class<StorageBuffer>>, %index : index) -> i16 {
   //     CHECK: %[[ARG1_CAST:.+]] = builtin.unrealized_conversion_cast %[[ARG1]] : index to i32
   //     CHECK: %[[ZERO:.+]] = spirv.Constant 0 : i32
   //     CHECK: %[[ONE:.+]] = spirv.Constant 1 : i32
   //     CHECK: %[[UPDATE:.+]] = spirv.IMul %[[ONE]], %[[ARG1_CAST]] : i32
   //     CHECK: %[[FLAT_IDX:.+]] = spirv.IAdd %[[ZERO]], %[[UPDATE]] : i32
   //     CHECK: %[[TWO:.+]] = spirv.Constant 2 : i32
   //     CHECK: %[[QUOTIENT:.+]] = spirv.SDiv %[[FLAT_IDX]], %[[TWO]] : i32
   //     CHECK: %[[PTR:.+]] = spirv.AccessChain %{{.+}}[%[[ZERO]], %[[QUOTIENT]]]
   //     CHECK: %[[LOAD:.+]] = spirv.Load  "StorageBuffer" %[[PTR]]
   //     CHECK: %[[SIXTEEN:.+]] = spirv.Constant 16 : i32
   //     CHECK: %[[IDX:.+]] = spirv.UMod %[[FLAT_IDX]], %[[TWO]] : i32
   //     CHECK: %[[BITS:.+]] = spirv.IMul %[[IDX]], %[[SIXTEEN]] : i32
   //     CHECK: %[[VALUE:.+]] = spirv.ShiftRightArithmetic %[[LOAD]], %[[BITS]] : i32, i32
   //     CHECK: %[[MASK:.+]] = spirv.Constant 65535 : i32
   //     CHECK: %[[T1:.+]] = spirv.BitwiseAnd %[[VALUE]], %[[MASK]] : i32
   //     CHECK: %[[T3:.+]] = spirv.ShiftLeftLogical %[[T1]], %[[SIXTEEN]] : i32, i32
-  //     CHECK: spirv.ShiftRightArithmetic %[[T3]], %[[SIXTEEN]] : i32, i32
+  //     CHECK: %[[SR:.+]] = spirv.ShiftRightArithmetic %[[T3]], %[[SIXTEEN]] : i32, i32
+  //     CHECK: builtin.unrealized_conversion_cast %[[SR]]
   %0 = memref.load %arg0[%index] : memref<10xi16, #spirv.storage_class<StorageBuffer>>
   return %0: i16
 }
 
 // CHECK-LABEL: @load_f32
 func.func @load_f32(%arg0: memref<f32, #spirv.storage_class<StorageBuffer>>) {
   // CHECK-NOT: spirv.SDiv
   //     CHECK: spirv.Load
func.func @load_i4(%arg0: memref<?xi4, #spirv.storage_class<StorageBuffer>>, %i: index) -> i4 {
   // CHECK: %[[FOUR:.+]] = spirv.Constant 4 : i32
   // CHECK: %[[IDX:.+]] = spirv.UMod %[[OFFSET]], %[[EIGHT]] : i32
   // CHECK: %[[BITS:.+]] = spirv.IMul %[[IDX]], %[[FOUR]] : i32
   // CHECK: %[[VALUE:.+]] = spirv.ShiftRightArithmetic %[[LOAD]], %[[BITS]] : i32, i32
   // CHECK: %[[MASK:.+]] = spirv.Constant 15 : i32
   // CHECK: %[[AND:.+]] = spirv.BitwiseAnd %[[VALUE]], %[[MASK]] : i32
   // CHECK: %[[C28:.+]] = spirv.Constant 28 : i32
   // CHECK: %[[SL:.+]] = spirv.ShiftLeftLogical %[[AND]], %[[C28]] : i32, i32
-  // CHECK: spirv.ShiftRightArithmetic %[[SL]], %[[C28]] : i32, i32
+  // CHECK: %[[SR:.+]] = spirv.ShiftRightArithmetic %[[SL]], %[[C28]] : i32, i32
+  // CHECK: builtin.unrealized_conversion_cast %[[SR]]
   %0 = memref.load %arg0[%i] : memref<?xi4, #spirv.storage_class<StorageBuffer>>
   return %0 : i4
 }
 
 // CHECK-LABEL: @store_i4
 func.func @store_i4(%arg0: memref<?xi4, #spirv.storage_class<StorageBuffer>>, %value: i4, %i: index) {
   // CHECK: %[[VAL:.+]] = builtin.unrealized_conversion_cast %{{.+}} : i4 to i32
   // CHECK: %[[INDEX:.+]] = builtin.unrealized_conversion_cast %{{.+}} : index to i32
func.func @store_i4(%arg0: memref<?xi4, #spirv.storage_class<StorageBuffer>>, %value: i4, %i: index) {
   // CHECK: %[[PTR:.+]] = spirv.AccessChain %{{.+}}[%[[ZERO]], %[[ACCESS_INDEX]]]
   // CHECK: spirv.AtomicAnd "Device" "AcquireRelease" %[[PTR]], %[[MASK2]]
   // CHECK: spirv.AtomicOr "Device" "AcquireRelease" %[[PTR]], %[[STORE_VAL]]
   memref.store %value, %arg0[%i] : memref<?xi4, #spirv.storage_class<StorageBuffer>>
   return
 }
 
 } // end module
+
+// -----
+
+// Check that casts are properly inserted if the corresponding **compute**
+// capability is allowed.
+module attributes {
+  spirv.target_env = #spirv.target_env<
+    #spirv.vce<v1.0, [Shader, Int8, Int16], [
+      SPV_KHR_8bit_storage, SPV_KHR_16bit_storage, SPV_KHR_storage_buffer_storage_class
+      ]>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: @load_i1
+func.func @load_i1(%arg0: memref<i1, #spirv.storage_class<StorageBuffer>>) -> i1 {
+  //     CHECK: %[[ONE:.+]] = spirv.Constant 1 : i32
+  //     CHECK: %[[RES:.+]]  = spirv.IEqual %{{.+}}, %[[ONE]] : i32
+  //     CHECK: return %[[RES]]
+  %0 = memref.load %arg0[] : memref<i1, #spirv.storage_class<StorageBuffer>>
+  return %0 : i1
+}
+
+// CHECK-LABEL: @load_i8
+func.func @load_i8(%arg0: memref<i8, #spirv.storage_class<StorageBuffer>>) -> i8 {
+  //     CHECK: %[[RES:.+]] = spirv.UConvert %{{.+}} : i32 to i8
+  //     CHECK: return %[[RES]]
+  %0 = memref.load %arg0[] : memref<i8, #spirv.storage_class<StorageBuffer>>
+  return %0 : i8
+}
+
+// CHECK-LABEL: @load_i16
+func.func @load_i16(%arg0: memref<10xi16, #spirv.storage_class<StorageBuffer>>, %index : index) -> i16 {
+  //     CHECK: %[[RES:.+]] = spirv.UConvert %{{.+}} : i32 to i16
+  //     CHECK: return %[[RES]]
+  %0 = memref.load %arg0[%index] : memref<10xi16, #spirv.storage_class<StorageBuffer>>
+  return %0: i16
+}
+
+} // end module