[mlir][spirv] Use UnrealizedConversionCast in ArithmeticToSPIRV

[antiagainst]

Summary

This avoids pulling in function converion patterns, which is not
part of what we want to test in ArithmeticToSPIRV. It also allows
using ConvertArithmeticToSPIRVPass as a standalone step.

[ThomasRaoux]

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

881–889

should this go in the type converter so that all the convertToSPIRV pass would have it?

[antiagainst]

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

881–889

That would mean to force all downstream passes to accept the unrealized_conversion_cast as a legal op and allow generating them. I'd prefer to actually give the option to downstream passes cause that's where we pull all patterns together, so we have the flexibility to either use one pass to convert all or use multiple passes to bridge the conversion step by step. These few lines of code will be duplicated for sure but I think it's okay.

[ThomasRaoux]

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

881–889

ok, I don't have a strong opinion but since this is how convert to llvm does it I thought it would make sense to match it. In my experience it is convenient to see the unrealized cast that didn't get eliminated when looking for ops that have not been lowered but that's subjective.

[antiagainst]

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

881–889

I see your points. Yeah that's reasonable too. Anyway, not a big issue; we can change it later. :)

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

void runOnOperation() override {
     auto module = getOperation();
     auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
     auto target = SPIRVConversionTarget::get(targetAttr);
 
     SPIRVTypeConverter::Options options;
     options.emulateNon32BitScalarTypes = this->emulateNon32BitScalarTypes;
     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 Optional<Value>(cast.getResult(0));
+    };
+    typeConverter.addSourceMaterialization(addUnrealizedCast);
+    typeConverter.addTargetMaterialization(addUnrealizedCast);
+    target->addLegalOp<UnrealizedConversionCastOp>();
+
     RewritePatternSet patterns(&getContext());
     arith::populateArithmeticToSPIRVPatterns(typeConverter, patterns);
-    populateFuncToSPIRVPatterns(typeConverter, patterns);
-    populateBuiltinFuncToSPIRVPatterns(typeConverter, patterns);
 
     if (failed(applyPartialConversion(module, *target, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::arith::createConvertArithmeticToSPIRVPass() {
   return std::make_unique<ConvertArithmeticToSPIRVPass>();
 }

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

func.func @index_scalar(%lhs: index, %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)
+// CHECK-SAME: (%[[A:.+]]: index, %[[B:.+]]: index)
 func.func @index_scalar_srem(%lhs: index, %rhs: index) {
+  // CHECK: %[[LHS:.+]] = builtin.unrealized_conversion_cast %[[A]] : index to i32
+  // CHECK: %[[RHS:.+]] = builtin.unrealized_conversion_cast %[[B]] : index to i32
   // 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 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, [], []>, #spv.resource_limits<>>
 } {
 
-// expected-error @+1 {{failed to materialize conversion for block argument #0 that remained live after conversion, type was 'vector<4xi64>', with target type 'vector<4xi32>'}}
 func.func @int_vector4_invalid(%arg0: vector<4xi64>) {
-  // expected-error @+2 {{bitwidth emulation is not implemented yet on unsigned op}}
-  // expected-note @+1 {{see existing live user here}}
+  // expected-error @+1 {{bitwidth emulation is not implemented yet on unsigned op}}
   %0 = arith.divui %arg0, %arg0: vector<4xi64>
   return
 }
 
 } // end module
 
 // -----
 

 
 // Checks that cast types will be adjusted when missing special capabilities for
 // certain non-32-bit scalar types.
 module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [Float64], []>, #spv.resource_limits<>>
 } {
 
 // CHECK-LABEL: @fpext1
-// CHECK-SAME: %[[ARG:.*]]: f32
+// CHECK-SAME: %[[A:.*]]: f16
 func.func @fpext1(%arg0: f16) -> f64 {
+  // CHECK: %[[ARG:.+]] = builtin.unrealized_conversion_cast %[[A]] : f16 to f32
   // CHECK-NEXT: spv.FConvert %[[ARG]] : f32 to f64
   %0 = arith.extf %arg0 : f16 to f64
   return %0: f64
 }
 
 // CHECK-LABEL: @fpext2
 // CHECK-SAME: %[[ARG:.*]]: f32
 func.func @fpext2(%arg0 : f32) -> f64 {
   // CHECK-NEXT: spv.FConvert %[[ARG]] : f32 to f64
   %0 = arith.extf %arg0 : f32 to f64
   return %0: f64
 }
 
 } // end module
 
 // -----
 
 // Checks that cast types will be adjusted when missing special capabilities for
 // certain non-32-bit scalar types.
 module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [Float16], []>, #spv.resource_limits<>>
 } {
 
 // CHECK-LABEL: @fptrunc1
-// CHECK-SAME: %[[ARG:.*]]: f32
+// CHECK-SAME: %[[A:.*]]: f64
 func.func @fptrunc1(%arg0 : f64) -> f16 {
+  // CHECK: %[[ARG:.+]] = builtin.unrealized_conversion_cast %[[A]] : f64 to f32
   // CHECK-NEXT: spv.FConvert %[[ARG]] : f32 to f16
   %0 = arith.truncf %arg0 : f64 to f16
   return %0: f16
 }
 
 // CHECK-LABEL: @fptrunc2
 // CHECK-SAME: %[[ARG:.*]]: f32
 func.func @fptrunc2(%arg0: f32) -> f16 {

 // -----
 
 // 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, [], []>, #spv.resource_limits<>>
 } {
 
-// expected-error@below {{failed to materialize conversion for block argument #0 that remained live after conversion}}
 func.func @int_vector4_invalid(%arg0: vector<4xi64>) {
-  // expected-error@below {{bitwidth emulation is not implemented yet on unsigned op}}
-  // expected-note@below {{see existing live user here}}
+  // expected-error@+1 {{bitwidth emulation is not implemented yet on unsigned op}}
   %0 = arith.divui %arg0, %arg0: vector<4xi64>
   return
 }
 
 } // end module
 
 // -----
 

 
 // Checks that cast types will be adjusted when missing special capabilities for
 // certain non-32-bit scalar types.
 module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [Float64], []>, #spv.resource_limits<>>
 } {
 
 // CHECK-LABEL: @fpext1
-// CHECK-SAME: %[[ARG:.*]]: f32
+// CHECK-SAME: %[[A:.*]]: f16
 func.func @fpext1(%arg0: f16) -> f64 {
+  // CHECK: %[[ARG:.+]] = builtin.unrealized_conversion_cast %[[A]] : f16 to f32
   // CHECK-NEXT: spv.FConvert %[[ARG]] : f32 to f64
   %0 = arith.extf %arg0 : f16 to f64
   return %0: f64
 }
 
 // CHECK-LABEL: @fpext2
 // CHECK-SAME: %[[ARG:.*]]: f32
 func.func @fpext2(%arg0 : f32) -> f64 {
   // CHECK-NEXT: spv.FConvert %[[ARG]] : f32 to f64
   %0 = arith.extf %arg0 : f32 to f64
   return %0: f64
 }
 
 } // end module
 
 // -----
 
 // Checks that cast types will be adjusted when missing special capabilities for
 // certain non-32-bit scalar types.
 module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [Float16], []>, #spv.resource_limits<>>
 } {
 
 // CHECK-LABEL: @fptrunc1
-// CHECK-SAME: %[[ARG:.*]]: f32
+// CHECK-SAME: %[[A:.*]]: f64
 func.func @fptrunc1(%arg0 : f64) -> f16 {
+  // CHECK: %[[ARG:.+]] = builtin.unrealized_conversion_cast %[[A]] : f64 to f32
   // CHECK-NEXT: spv.FConvert %[[ARG]] : f32 to f16
   %0 = arith.truncf %arg0 : f64 to f16
   return %0: f16
 }
 
 // CHECK-LABEL: @fptrunc2
 // CHECK-SAME: %[[ARG:.*]]: f32
 func.func @fptrunc2(%arg0: f32) -> f16 {