[mlir][spirv] Fold noop `BitcastsOp`s

[kuhar]

Summary

This allows for bitcast conversion to roundtrip.

Fixes: https://github.com/llvm/llvm-project/issues/58801

[Hardcode84]

Should we also update verifier to allow bitcasts to have same src and dst type? Is this allowed in spirv?

upd: original issue was failing with "result type must be different from operand type" but I cannot even find where it is checked.

[kuhar]

In D137459#3909079, @Hardcode84 wrote:

Should we also update verifier to allow bitcasts to have same src and dst type? Is this allowed in spirv?

It's not allowed: https://registry.khronos.org/SPIR-V/specs/1.2/SPIRV.html#OpBitcast

Operand must have a type of OpTypePointer, or a scalar or vector of numerical-type. It must be a different type than Result Type.

[Hardcode84]

I'm ok to merge it as is to fix the bug, but we are basically abusing the fact the folder will run immediately after pattern, 'fixing' invalid IR that was produced.

[mravishankar]

In D137459#3909091, @Hardcode84 wrote:

I'm ok to merge it as is to fix the bug, but we are basically abusing the fact the folder will run immediately after pattern, 'fixing' invalid IR that was produced.

Good point. Stamped without thinking through this. Maybe we can just change the folder to look for

%0 = spirv.Bitcast %a : <typea> to <typeb>
%1 = spirv.Bitcast %b : <typeb> to <typec>

and then fold to either

%0 = spirv.BItcast %a : <typea> to <typec>

or just

%a

if <typea> == <typeb>.

Marking as request changes for now.

[kuhar]

Replace the canonicalizer with the folder to avoid generating invalid bitcasts.

[kuhar]

Add early return

[Hardcode84]

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

134

It is only checks for limited depth, you can do smth like this to check into arbitrary deoth

Value current = getOperand();
while (auto parent = current.getDefiningOp<BitcastOp>()) {
  auto parentSource = parent.getOperand();
  if (parentSource.getType() == getType())
    return parentSource;

  current = parentSource;
}

[Hardcode84]

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

134

Wait, nevermind, it will still check for arbitrary depth after multiple iterations, right?

[kuhar]

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

134

It works with longer chains because the innermost argument is always pushed through

[kuhar]

Update test to demonstrate that long chains of bitcasts are also folded

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

def SPIRV_BitcastOp : SPIRV_Op<"Bitcast", [Pure]> {
   let results = (outs
     SPIRV_ScalarOrVectorOrPtr:$result
   );
 
   let assemblyFormat = [{
     $operand attr-dict `:` type($operand) `to` type($result)
   }];
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
 }
 
 // -----
 
 def SPIRV_ConvertFToSOp : SPIRV_CastOp<"ConvertFToS", SPIRV_Integer, SPIRV_Float, []> {
   let summary = [{
     Convert value numerically from floating point to signed integer, with
     round toward 0.0.

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

void spirv::AccessChainOp::getCanonicalizationPatterns(
     RewritePatternSet &results, MLIRContext *context) {
   results.add<CombineChainedAccessChain>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.BitcastOp
 //===----------------------------------------------------------------------===//
 
+OpFoldResult spirv::BitcastOp::fold(ArrayRef<Attribute> /*operands*/) {
+  if (getType() == getOperand().getType())
+    return getOperand();
+
+  return Value{};
+}
+
 void spirv::BitcastOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                    MLIRContext *context) {
   results.add<ConvertChainedBitcast>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.CompositeExtractOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult spirv::CompositeExtractOp::fold(ArrayRef<Attribute> operands) {
   if (auto insertOp =
           getComposite().getDefiningOp<spirv::CompositeInsertOp>()) {
     if (getIndices() == insertOp.getIndices())
       return insertOp.getObject();
   }
 
   if (auto constructOp =

mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir

func.func @convert_bitcast_multi_use(%arg0 : vector<2xf32>, %arg1 : !spirv.ptr<i64, Uniform>) -> f64 {
   %0 = spirv.Bitcast %arg0 : vector<2xf32> to i64
   %1 = spirv.Bitcast %0 : i64 to f64
   spirv.Store "Uniform" %arg1, %0 : i64
   spirv.ReturnValue %1 : f64
 }
 
 // -----
 
+// CHECK-LABEL: @convert_bitcast_roundtip
+// CHECK-SAME:    %[[ARG:.+]]: i64
+func.func @convert_bitcast_roundtip(%arg0 : i64) -> i64 {
+  // CHECK: spirv.ReturnValue %[[ARG]]
+  %0 = spirv.Bitcast %arg0 : i64 to f64
+  %1 = spirv.Bitcast %0 : f64 to i64
+  spirv.ReturnValue %1 : i64
+}
+
+// -----
+
+// CHECK-LABEL: @convert_bitcast_chained_roundtip
+// CHECK-SAME:    %[[ARG:.+]]: i64
+func.func @convert_bitcast_chained_roundtip(%arg0 : i64) -> i64 {
+  // CHECK: spirv.ReturnValue %[[ARG]]
+  %0 = spirv.Bitcast %arg0 : i64 to f64
+  %1 = spirv.Bitcast %0 : f64 to vector<2xi32>
+  %2 = spirv.Bitcast %1 : vector<2xi32> to i64
+  spirv.ReturnValue %2 : i64
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // spirv.CompositeExtract
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: extract_vector
 func.func @extract_vector() -> (i32, i32, i32) {
   // CHECK-DAG: spirv.Constant 6 : i32
   // CHECK-DAG: spirv.Constant -33 : i32