[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

125

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

125

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

[kuhar]

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

125

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
 //===----------------------------------------------------------------------===//
 
-void spirv::BitcastOp::getCanonicalizationPatterns(RewritePatternSet &results,
-                                                   MLIRContext *context) {
-  results.add<ConvertChainedBitcast>(context);
+OpFoldResult spirv::BitcastOp::fold(ArrayRef<Attribute> /*operands*/) {
+  Value arg = getOperand();
+  if (getType() == arg.getType())
+    return arg;
+
+  // Look through nested bitcasts.
+  if (auto bitcast = arg.getDefiningOp<spirv::BitcastOp>()) {
+    Value nestedArg = bitcast.getOperand();
+    if (nestedArg.getType() == getType())
+      return nestedArg;
+
+    getOperandMutable().assign(nestedArg);
+    return getResult();
+  }
+
+  // TODO(kuhar): Consider constant-folding the operand attribute.
+  return getResult();
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.CompositeExtractOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult spirv::CompositeExtractOp::fold(ArrayRef<Attribute> operands) {
   if (auto insertOp =

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

 //==- SPIRVCanonicalization.td - Canonicalization Patterns ---*- tablegen -*==//
 
 // 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 defines SPIR-V canonicalization patterns.
 //
 //===----------------------------------------------------------------------===//
 
 include "mlir/IR/PatternBase.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVOps.td"
 
 //===----------------------------------------------------------------------===//
-// spirv.Bitcast
-//===----------------------------------------------------------------------===//
-
-def ConvertChainedBitcast : Pat<(SPIRV_BitcastOp (SPIRV_BitcastOp $operand)),
-                                (SPIRV_BitcastOp $operand)>;
-
-//===----------------------------------------------------------------------===//
 // spirv.LogicalNot
 //===----------------------------------------------------------------------===//
 
 def ConvertLogicalNotOfIEqual : Pat<
     (SPIRV_LogicalNotOp (SPIRV_IEqualOp $lhs, $rhs)),
     (SPIRV_INotEqualOp $lhs, $rhs)>;
 
 def ConvertLogicalNotOfINotEqual : Pat<

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 vector<2xf32>
+  %3 = spirv.Bitcast %2 : vector<2xf32> to i64
+  spirv.ReturnValue %3 : 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