[mlir] Add a pattern to bufferize std.index_cast.

[bkramer]

looks good

mlir/lib/Dialect/StandardOps/Transforms/Bufferize.cpp

public:
   matchAndRewrite(memref::DimOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     memref::DimOp::Adaptor adaptor(operands);
     rewriter.replaceOpWithNewOp<memref::DimOp>(op, adaptor.memrefOrTensor(),
                                                adaptor.index());
     return success();
   }
 };
-} // namespace
 
-namespace {
+class BufferizeIndexCastOp : public OpConversionPattern<IndexCastOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(IndexCastOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    IndexCastOp::Adaptor adaptor(operands);
+    auto tensorType = op.getType().cast<RankedTensorType>();
+    rewriter.replaceOpWithNewOp<IndexCastOp>(
+        op, adaptor.in(),
+        MemRefType::get(tensorType.getShape(), tensorType.getElementType()));
+    return success();
+  }
+};
+
 class BufferizeSelectOp : public OpConversionPattern<SelectOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(SelectOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     if (!op.condition().getType().isa<IntegerType>())
       return rewriter.notifyMatchFailure(op, "requires scalar condition");
 
     SelectOp::Adaptor adaptor(operands);
     rewriter.replaceOpWithNewOp<SelectOp>(
         op, adaptor.condition(), adaptor.true_value(), adaptor.false_value());
     return success();
   }
 };
 } // namespace
 
 void mlir::populateStdBufferizePatterns(BufferizeTypeConverter &typeConverter,
                                         RewritePatternSet &patterns) {
-  patterns.add<BufferizeDimOp, BufferizeSelectOp>(typeConverter,
-                                                  patterns.getContext());
+  patterns.add<BufferizeDimOp, BufferizeSelectOp, BufferizeIndexCastOp>(
+      typeConverter, patterns.getContext());
 }
 
 namespace {
 struct StdBufferizePass : public StdBufferizeBase<StdBufferizePass> {
   void runOnFunction() override {
     auto *context = &getContext();
     BufferizeTypeConverter typeConverter;
     RewritePatternSet patterns(context);
     ConversionTarget target(*context);
 
-    target.addLegalDialect<memref::MemRefDialect>();
-    target.addLegalDialect<StandardOpsDialect>();
-    target.addLegalDialect<scf::SCFDialect>();
+    target.addLegalDialect<scf::SCFDialect, StandardOpsDialect,
+                           memref::MemRefDialect>();
 
     populateStdBufferizePatterns(typeConverter, patterns);
     // We only bufferize the case of tensor selected type and scalar condition,
     // as that boils down to a select over memref descriptors (don't need to
     // touch the data).
+    target.addDynamicallyLegalOp<IndexCastOp>(
+        [&](IndexCastOp op) { return typeConverter.isLegal(op.getType()); });
     target.addDynamicallyLegalOp<SelectOp>([&](SelectOp op) {
       return typeConverter.isLegal(op.getType()) ||
              !op.condition().getType().isa<IntegerType>();
     });
     target.addDynamicallyLegalOp<memref::DimOp>(
         [&](memref::DimOp op) { return typeConverter.isLegal(op); });
     if (failed(
             applyPartialConversion(getFunction(), target, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createStdBufferizePass() {
   return std::make_unique<StdBufferizePass>();
 }

mlir/test/Dialect/Standard/bufferize.mlir

 // CHECK:           %[[FALSE_VAL_MEMREF:.*]] = memref.buffer_cast %[[FALSE_VAL]] : memref<f32>
 // CHECK:           %[[RET_MEMREF:.*]] = select %[[PRED]], %[[TRUE_VAL_MEMREF]], %[[FALSE_VAL_MEMREF]] : memref<f32>
 // CHECK:           %[[RET:.*]] = memref.tensor_load %[[RET_MEMREF]] : memref<f32>
 // CHECK:           return %[[RET]] : tensor<f32>
 func @select(%arg0: i1, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
   %0 = select %arg0, %arg1, %arg2 : tensor<f32>
   return %0 : tensor<f32>
 }
+
+// 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 = index_cast %tensor : tensor<i32> to tensor<index>
+  %index_scalar = index_cast %scalar : i32 to index
+  return %index_tensor, %index_scalar : tensor<index>, index
+}
+// CHECK:  %[[MEMREF:.*]] = memref.buffer_cast %[[TENSOR]] : memref<i32>
+// CHECK-NEXT: %[[INDEX_MEMREF:.*]] = index_cast %[[MEMREF]]
+// CHECK-SAME:   memref<i32> to memref<index>
+// CHECK-NEXT: %[[INDEX_TENSOR:.*]] = memref.tensor_load %[[INDEX_MEMREF]]
+// CHECK: return %[[INDEX_TENSOR]]