[mlir][linalg][bufferize] Support arith.index_cast bufferization

[springerm]

Summary

This is in preparation of switching -tensor-constant-bufferize and -arith-bufferize to BufferizableOpInterface-based implementations.

Depends On D118246

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ArithInterfaceImpl.cpp

struct ConstantOpInterface
   bool isWritable(Operation *op, Value value,
                   const BufferizationState &state) const {
     // Memory locations returned by memref::GetGlobalOp may not be written to.
     assert(value.isa<OpResult>());
     return false;
   }
 };
 
+struct IndexCastOpInterface
+    : public BufferizableOpInterface::ExternalModel<IndexCastOpInterface,
+                                                    arith::IndexCastOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const BufferizationState &state) const {
+    return false;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const BufferizationState &state) const {
+    return false;
+  }
+
+  OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                               const BufferizationState &state) const {
+    return op->getResult(0);
+  }
+
+  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+                                const BufferizationState &state) const {
+    return BufferRelation::Equivalent;
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationState &state) const {
+    auto castOp = cast<arith::IndexCastOp>(op);
+
+    Value source = *state.getBuffer(rewriter, op->getOpOperand(0) /*in*/);
+    auto sourceType = source.getType().cast<BaseMemRefType>();
+
+    // Result type should have same layout and address space as the source type.
+    MemRefLayoutAttrInterface layout = {};
+    if (auto rankedMemRefType = sourceType.dyn_cast<MemRefType>())
+      layout = rankedMemRefType.getLayout();
+    Type resultType =
+        getMemRefType(castOp.getType().cast<TensorType>(), state.getOptions(),
+                      layout, sourceType.getMemorySpace());
+
+    replaceOpWithNewBufferizedOp<arith::IndexCastOp>(rewriter, op, source,
+                                                     resultType);
+    return success();
+  }
+};
 } // namespace arith_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 void mlir::linalg::comprehensive_bufferize::arith_ext::
     registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
   registry.addOpInterface<arith::ConstantOp, arith_ext::ConstantOpInterface>();
+  registry
+      .addOpInterface<arith::IndexCastOp, arith_ext::IndexCastOpInterface>();
 }

mlir/test/Dialect/Linalg/comprehensive-function-bufferize.mlir

%9 = scf.for %arg9 = %c0 to %c6 step %c1 iter_args(%arg10 = %2) -> (tensor<1x6x8xf32>) {
       %12 = tensor.insert_slice %11 into %arg10[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x1x8xf32> into tensor<1x6x8xf32>
       scf.yield %12 : tensor<1x6x8xf32>
     }
     %10 = tensor.insert_slice %9 into %arg8[%7, 0, 0, 0] [1, 1, 6, 8] [1, 1, 1, 1] : tensor<1x6x8xf32> into tensor<?x1x6x8xf32>
     scf.yield %10 : tensor<?x1x6x8xf32>
   }
   return %5: tensor<?x1x6x8xf32>
 }
+
+// -----
+
+// CHECK: #[[$MAP:.*]] = affine_map<()[s0] -> (s0)>
+// CHECK-LABEL:   func @index_cast(
+// CHECK-SAME:  %[[TENSOR:.*]]: tensor<i32>, %[[SCALAR:.*]]: i32
+func @index_cast(%tensor: tensor<i32>, %scalar: i32) -> (tensor<index>, index) {
+  %index_tensor = arith.index_cast %tensor : tensor<i32> to tensor<index>
+  %index_scalar = arith.index_cast %scalar : i32 to index
+  return %index_tensor, %index_scalar : tensor<index>, index
+}
+// CHECK:  %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<i32, #[[$MAP]]>
+// CHECK-NEXT: %[[INDEX_MEMREF:.*]] = arith.index_cast %[[MEMREF]]
+// CHECK-SAME:   memref<i32, #[[$MAP]]> to memref<index, #[[$MAP]]>
+// CHECK-NEXT: %[[INDEX_TENSOR:.*]] = bufferization.to_tensor %[[INDEX_MEMREF]]
+// CHECK: return %[[INDEX_TENSOR]]