Add tensor.bitcast op to Tensor dialect

[trevor-m]

Summary

Add tensor.bitcast operator to bitcast between two tensors of compatible shape and same bit width. This can be use to reinterpret an unsigned integer as a signed integer or vice versa.

[rsuderman]

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

173

LLVM is migrating to a different cast syntax: https://mlir.llvm.org/deprecation/

So this should be dyn_cast<TensorType>(a)

198

Ditto on the cast issue.

[trevor-m]

Updated to use free function casts.

Thanks @rsuderman for the review!

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

code extraBaseClassDeclaration = [{
     ::mlir::SmallVector<::mlir::Range, 8> getOrCreateRanges(
         ::mlir::OpBuilder &b, ::mlir::Location loc) {
       return ::mlir::getOrCreateRanges(*this, b, loc);
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// BitcastOp
+//===----------------------------------------------------------------------===//
+
+def Tensor_BitcastOp : Tensor_Op<"bitcast", [
+    DeclareOpInterfaceMethods<CastOpInterface>,
+    Pure
+  ]> {
+  let summary = "tensor bitcast operation";
+  let description = [{
+    Bitcast a tensor from one type to another type of equivalent element width.
+    If both are ranked, then the rank should be the same and static dimensions
+    should match. 
+
+    Example:
+
+    ```mlir
+    // Bitcast from unsigned to signed or signless integer.
+    %2 = tensor.bitcast %1 : tensor<4xui32> to tensor<4xi32>
+    ```
+  }];
+
+  let arguments = (ins AnyTensor:$source);
+  let results = (outs AnyTensor:$dest);
+  let assemblyFormat = "$source attr-dict `:` type($source) `to` type($dest)";
+
+  let hasCanonicalizer = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_CastOp : Tensor_Op<"cast", [
     DeclareOpInterfaceMethods<CastOpInterface>,
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
     Pure
   ]> {

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

static llvm::SmallBitVector getDroppedDims(ArrayRef<int64_t> reducedShape,
   }
 
   assert(shapePos == static_cast<int64_t>(reducedShape.size()) &&
          "dimension mismatch");
   return droppedDims;
 }
 
 //===----------------------------------------------------------------------===//
+// BitcastOp
+//===----------------------------------------------------------------------===//
+
+bool BitcastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
+  if (inputs.size() != 1 || outputs.size() != 1)
+    return false;
+  Type a = inputs.front(), b = outputs.front();
+  auto aT = dyn_cast<TensorType>(a);
+  auto bT = dyn_cast<TensorType>(b);
+  if (!aT || !bT)
+    return false;
+
+  if (aT.getElementTypeBitWidth() != bT.getElementTypeBitWidth())
+    return false;
+
+  return succeeded(verifyCompatibleShape(aT, bT));
+}
+
+namespace {
+
+/// Replaces chains of two tensor.bitcast operations by a single tensor.bitcast
+/// operation.
+struct ChainedTensorBitcast : public OpRewritePattern<BitcastOp> {
+  using OpRewritePattern<BitcastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BitcastOp tensorBitcast,
+                                PatternRewriter &rewriter) const final {
+    auto tensorBitcastOperand =
+        tensorBitcast.getOperand().getDefiningOp<BitcastOp>();
+    if (!tensorBitcastOperand)
+      return failure();
+
+    auto resultType = cast<TensorType>(tensorBitcast.getType());
+    rewriter.replaceOpWithNewOp<BitcastOp>(tensorBitcast, resultType,
+                                           tensorBitcastOperand.getOperand());
+    return success();
+  }
+};
+
+} // namespace
+
+void BitcastOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                            MLIRContext *context) {
+  results.add<ChainedTensorBitcast>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 void CastOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "cast");
 }
 
 /// Returns true if `target` is a ranked tensor type that preserves static

mlir/test/Dialect/Tensor/canonicalize.mlir

 // RUN: mlir-opt %s -split-input-file -canonicalize="test-convergence" | FileCheck %s
 
+// CHECK-LABEL: @tensor_bitcast_chain_ok
+// CHECK-SAME: %[[IN:.*]]: tensor<2xi32>
+func.func @tensor_bitcast_chain_ok(%input: tensor<2xi32>) -> tensor<2xf32> {
+  // CHECK-NEXT: %[[RES:.*]] = tensor.bitcast %[[IN]] : tensor<2xi32> to tensor<2xf32>
+  %0 = tensor.bitcast %input : tensor<2xi32> to tensor<2xui32>
+  %1 = tensor.bitcast %0 : tensor<2xui32> to tensor<2xf32>
+  // CHECK-NEXT: return %[[RES]]
+  return %1 : tensor<2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @tensor_bitcast_chain_nop
+// CHECK-SAME: %[[IN:.*]]: tensor<4xi32>
+func.func @tensor_bitcast_chain_nop(%input: tensor<4xi32>) -> tensor<4xi32> {
+  %0 = tensor.bitcast %input : tensor<4xi32> to tensor<4xui32>
+  %1 = tensor.bitcast %0 : tensor<4xui32> to tensor<4xi32>
+  // CHECK-NEXT: return %[[IN]]
+  return %1 : tensor<4xi32>
+}
+
+// -----
+
 // Checks that NOP casts are removed.
 // CHECK-LABEL: cast_values
 func.func @cast_values(%arg0: tensor<*xi32>) -> tensor<2xi32> {
   // NOP cast
   %0 = tensor.cast %arg0 : tensor<*xi32> to tensor<*xi32>
   // CHECK-NEXT: %[[RET:.*]] = tensor.cast %arg0 : tensor<*xi32> to tensor<2xi32>
   %2 = tensor.cast %0 : tensor<*xi32> to tensor<2xi32>
   // NOP cast