[mlir][tensor] Add tensor.dim operation

* Split memref.dim into two operations: memref.dim and tensor.dim. Both ops have the same builder interface and op argument names, so that they can be used with templates in patterns that apply to both tensors and memrefs (e.g., some patterns in Linalg).
* Add constant materializer to TensorDialect (needed for folding in affine.apply etc.).
* Remove some MemRefDialect dependencies, make some explicit.

Differential Revision: https://reviews.llvm.org/D105165

Author: matthias-springer

mlir/include/mlir/Conversion/Passes.td

@@ -365,7 +365,6 @@ def ConvertShapeToStandard : Pass<"convert-shape-to-std", "ModuleOp"> {
                 "dialect";
   let constructor = "mlir::createConvertShapeToStandardPass()";
   let dependentDialects = [
-    "memref::MemRefDialect",
     "StandardOpsDialect",
     "scf::SCFDialect",
     "tensor::TensorDialect"

mlir/include/mlir/Dialect/MemRef/IR/MemRef.h

@@ -34,6 +34,10 @@ SmallVector<Range, 8> getOrCreateRanges(OffsetSizeAndStrideOpInterface op,
 /// Given an operation, retrieves the value of each dynamic dimension through
 /// constructing the necessary DimOp operators.
 SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b);
+
+// Helper function that creates a memref::DimOp or tensor::DimOp depending on
+// the type of `source`.
+Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

@@ -512,7 +512,7 @@ def MemRef_DeallocOp : MemRef_Op<"dealloc", [MemRefsNormalizable]> {
 // DimOp
 //===----------------------------------------------------------------------===//
 
-def DimOp : MemRef_Op<"dim", [NoSideEffect, MemRefsNormalizable]> {
+def MemRef_DimOp : MemRef_Op<"dim", [NoSideEffect, MemRefsNormalizable]> {
   let summary = "dimension index operation";
   let description = [{
     The `dim` operation takes a memref and a dimension operand of type `index`.
@@ -538,18 +538,17 @@ def DimOp : MemRef_Op<"dim", [NoSideEffect, MemRefsNormalizable]> {
     ```
   }];
 
-  let arguments = (ins AnyTypeOf<[AnyTensor, AnyRankedOrUnrankedMemRef],
-                                 "any memref or tensor type">:$memrefOrTensor,
+  let arguments = (ins AnyRankedOrUnrankedMemRef:$source,
                        Index:$index);
   let results = (outs Index:$result);
 
   let assemblyFormat = [{
-    attr-dict $memrefOrTensor `,` $index `:` type($memrefOrTensor)
+    attr-dict $source `,` $index `:` type($source)
   }];
 
   let builders = [
-    OpBuilder<(ins "Value":$memrefOrTensor, "int64_t":$index)>,
-    OpBuilder<(ins "Value":$memrefOrTensor, "Value":$index)>
+    OpBuilder<(ins "Value":$source, "int64_t":$index)>,
+    OpBuilder<(ins "Value":$source, "Value":$index)>
   ];
 
   let extraClassDeclaration = [{
@@ -1288,6 +1287,7 @@ def TensorLoadOp : MemRef_Op<"tensor_load",
 
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
 
+  let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td

@@ -14,7 +14,7 @@ include "mlir/Pass/PassBase.td"
 def StdBufferize : FunctionPass<"std-bufferize"> {
   let summary = "Bufferize the std dialect";
   let constructor = "mlir::createStdBufferizePass()";
-  let dependentDialects = ["scf::SCFDialect"];
+  let dependentDialects = ["memref::MemRefDialect", "scf::SCFDialect"];
 }
 
 def StdExpandOps : FunctionPass<"std-expand"> {

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

@@ -14,6 +14,7 @@ include "mlir/IR/OpBase.td"
 def Tensor_Dialect : Dialect {
   let name = "tensor";
   let cppNamespace = "::mlir::tensor";
+
   let description = [{
     The `tensor` dialect is intended to hold core tensor creation and
     manipulation ops, which are not strongly associated with any particular
@@ -43,6 +44,8 @@ def Tensor_Dialect : Dialect {
     dialect), and does not live in this dialect.
 
   }];
+
+  let hasConstantMaterializer = 1;
 }
 
 #endif // TENSOR_BASE

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

@@ -60,6 +60,58 @@ def Tensor_CastOp : Tensor_Op<"cast", [
   let verifier = ?;
 }
 
+//===----------------------------------------------------------------------===//
+// DimOp
+//===----------------------------------------------------------------------===//
+
+def Tensor_DimOp : Tensor_Op<"dim", [NoSideEffect]> {
+  let summary = "dimension index operation";
+  let description = [{
+    The `dim` operation takes a tensor and a dimension operand of type `index`.
+    It returns the size of the requested dimension of the given tensor.
+    If the dimension index is out of bounds, the behavior is undefined.
+
+    The specified tensor type is that of the first operand.
+
+    Example:
+
+    ```mlir
+    // Always returns 4, can be constant folded:
+    %c0 = constant 0 : index
+    %x = tensor.dim %A, %c0 : tensor<4x?xf32>
+
+    // Returns the dynamic dimension of %A.
+    %c1 = constant 1 : index
+    %y = tensor.dim %A, %c1 : memref<4x?xf32>
+
+    // Equivalent generic form:
+    %x = "tensor.dim"(%A, %c0) : (memref<4x?xf32>, index) -> index
+    %y = "tensor.dim"(%A, %c1) : (memref<4x?xf32>, index) -> index
+    ```
+  }];
+
+  let arguments = (ins AnyTensor:$source,
+                       Index:$index);
+  let results = (outs Index:$result);
+
+  let assemblyFormat = [{
+    attr-dict $source `,` $index `:` type($source)
+  }];
+
+  let builders = [
+    OpBuilder<(ins "Value":$source, "int64_t":$index)>,
+    OpBuilder<(ins "Value":$source, "Value":$index)>
+  ];
+
+  let extraClassDeclaration = [{
+    /// Helper function to get the index as a simple integer if it is constant.
+    Optional<int64_t> getConstantIndex();
+  }];
+
+  let hasCanonicalizer = 1;
+  let hasFolder = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // ExtractOp
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td

@@ -14,7 +14,7 @@ include "mlir/Pass/PassBase.td"
 def TensorBufferize : FunctionPass<"tensor-bufferize"> {
   let summary = "Bufferize the `tensor` dialect";
   let constructor = "mlir::createTensorBufferizePass()";
-  let dependentDialects = ["scf::SCFDialect"];
+  let dependentDialects = ["scf::SCFDialect", "memref::MemRefDialect"];
 }
 
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES

mlir/lib/Conversion/ShapeToStandard/CMakeLists.txt

@@ -18,7 +18,6 @@ add_mlir_conversion_library(MLIRShapeToStandard
 
   LINK_LIBS PUBLIC
   MLIRIR
-  MLIRMemRef
   MLIRShape
   MLIRTensor
   MLIRPass

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

@@ -9,7 +9,6 @@
 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 
 #include "../PassDetail.h"
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
@@ -140,7 +139,7 @@ LogicalResult BroadcastOpConverter::matchAndRewrite(
   // dimension in the tensor.
   SmallVector<Value> ranks, rankDiffs;
   llvm::append_range(ranks, llvm::map_range(transformed.shapes(), [&](Value v) {
-                       return lb.create<memref::DimOp>(v, zero);
+                       return lb.create<tensor::DimOp>(v, zero);
                      }));
 
   // Find the maximum rank
@@ -254,7 +253,7 @@ LogicalResult IsBroadcastableOpConverter::matchAndRewrite(
   // dimension in the tensor.
   SmallVector<Value> ranks, rankDiffs;
   llvm::append_range(ranks, llvm::map_range(transformed.shapes(), [&](Value v) {
-                       return lb.create<memref::DimOp>(v, zero);
+                       return lb.create<tensor::DimOp>(v, zero);
                      }));
 
   // Find the maximum rank
@@ -346,7 +345,7 @@ LogicalResult GetExtentOpConverter::matchAndRewrite(
   // circumvents the necessity to materialize the shape in memory.
   if (auto shapeOfOp = op.shape().getDefiningOp<ShapeOfOp>()) {
     if (shapeOfOp.arg().getType().isa<ShapedType>()) {
-      rewriter.replaceOpWithNewOp<memref::DimOp>(op, shapeOfOp.arg(),
+      rewriter.replaceOpWithNewOp<tensor::DimOp>(op, shapeOfOp.arg(),
                                                  transformed.dim());
       return success();
     }
@@ -377,7 +376,7 @@ RankOpConverter::matchAndRewrite(shape::RankOp op, ArrayRef<Value> operands,
     return failure();
 
   shape::RankOp::Adaptor transformed(operands);
-  rewriter.replaceOpWithNewOp<memref::DimOp>(op, transformed.shape(), 0);
+  rewriter.replaceOpWithNewOp<tensor::DimOp>(op, transformed.shape(), 0);
   return success();
 }
 
@@ -407,7 +406,7 @@ ReduceOpConverter::matchAndRewrite(shape::ReduceOp op, ArrayRef<Value> operands,
   Value one = rewriter.create<ConstantIndexOp>(loc, 1);
   Type indexTy = rewriter.getIndexType();
   Value rank =
-      rewriter.create<memref::DimOp>(loc, indexTy, transformed.shape(), zero);
+      rewriter.create<tensor::DimOp>(loc, indexTy, transformed.shape(), zero);
 
   auto loop = rewriter.create<scf::ForOp>(
       loc, zero, rank, one, op.initVals(),
@@ -494,11 +493,11 @@ ShapeEqOpConverter::matchAndRewrite(ShapeEqOp op, ArrayRef<Value> operands,
   Value zero = rewriter.create<ConstantIndexOp>(loc, 0);
   Value firstShape = transformed.shapes().front();
   Value firstRank =
-      rewriter.create<memref::DimOp>(loc, indexTy, firstShape, zero);
+      rewriter.create<tensor::DimOp>(loc, indexTy, firstShape, zero);
   Value result = nullptr;
   // Generate a linear sequence of compares, all with firstShape as lhs.
   for (Value shape : transformed.shapes().drop_front(1)) {
-    Value rank = rewriter.create<memref::DimOp>(loc, indexTy, shape, zero);
+    Value rank = rewriter.create<tensor::DimOp>(loc, indexTy, shape, zero);
     Value eqRank =
         rewriter.create<CmpIOp>(loc, CmpIPredicate::eq, firstRank, rank);
     auto same = rewriter.create<IfOp>(
@@ -563,7 +562,7 @@ LogicalResult ShapeOfOpConversion::matchAndRewrite(
     int64_t rank = rankedTensorTy.getRank();
     for (int64_t i = 0; i < rank; i++) {
       if (rankedTensorTy.isDynamicDim(i)) {
-        Value extent = rewriter.create<memref::DimOp>(loc, tensor, i);
+        Value extent = rewriter.create<tensor::DimOp>(loc, tensor, i);
         extentValues.push_back(extent);
       } else {
         Value extent =
@@ -587,7 +586,7 @@ LogicalResult ShapeOfOpConversion::matchAndRewrite(
       op, getExtentTensorType(ctx), ValueRange{rank},
       [&](OpBuilder &b, Location loc, ValueRange args) {
         Value dim = args.front();
-        Value extent = b.create<memref::DimOp>(loc, tensor, dim);
+        Value extent = b.create<tensor::DimOp>(loc, tensor, dim);
         b.create<tensor::YieldOp>(loc, extent);
       });
 
@@ -617,7 +616,7 @@ LogicalResult SplitAtOpConversion::matchAndRewrite(
   SplitAtOp::Adaptor transformed(op);
   ImplicitLocOpBuilder b(op.getLoc(), rewriter);
   Value zero = b.create<ConstantIndexOp>(0);
-  Value rank = b.create<memref::DimOp>(transformed.operand(), zero);
+  Value rank = b.create<tensor::DimOp>(transformed.operand(), zero);
 
   // index < 0 ? index + rank : index
   Value originalIndex = transformed.index();
@@ -675,8 +674,8 @@ void ConvertShapeToStandardPass::runOnOperation() {
   // Setup target legality.
   MLIRContext &ctx = getContext();
   ConversionTarget target(ctx);
-  target.addLegalDialect<memref::MemRefDialect, StandardOpsDialect, SCFDialect,
-                         tensor::TensorDialect>();
+  target
+      .addLegalDialect<StandardOpsDialect, SCFDialect, tensor::TensorDialect>();
   target.addLegalOp<CstrRequireOp, FuncOp, ModuleOp>();
 
   // Setup conversion patterns.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

@@ -2965,7 +2965,7 @@ struct DimOpLowering : public ConvertOpToLLVMPattern<memref::DimOp> {
   LogicalResult
   matchAndRewrite(memref::DimOp dimOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    Type operandType = dimOp.memrefOrTensor().getType();
+    Type operandType = dimOp.source().getType();
     if (operandType.isa<UnrankedMemRefType>()) {
       rewriter.replaceOp(dimOp, {extractSizeOfUnrankedMemRef(
                                     operandType, dimOp, operands, rewriter)});
@@ -2977,7 +2977,7 @@ struct DimOpLowering : public ConvertOpToLLVMPattern<memref::DimOp> {
                                     operandType, dimOp, operands, rewriter)});
       return success();
     }
-    return failure();
+    llvm_unreachable("expected MemRefType or UnrankedMemRefType");
   }
 
 private:
@@ -2995,7 +2995,7 @@ struct DimOpLowering : public ConvertOpToLLVMPattern<memref::DimOp> {
     // Extract pointer to the underlying ranked descriptor and bitcast it to a
     // memref<element_type> descriptor pointer to minimize the number of GEP
     // operations.
-    UnrankedMemRefDescriptor unrankedDesc(transformed.memrefOrTensor());
+    UnrankedMemRefDescriptor unrankedDesc(transformed.source());
     Value underlyingRankedDesc = unrankedDesc.memRefDescPtr(rewriter, loc);
     Value scalarMemRefDescPtr = rewriter.create<LLVM::BitcastOp>(
         loc,
@@ -3033,7 +3033,7 @@ struct DimOpLowering : public ConvertOpToLLVMPattern<memref::DimOp> {
       int64_t i = index.getValue();
       if (memRefType.isDynamicDim(i)) {
         // extract dynamic size from the memref descriptor.
-        MemRefDescriptor descriptor(transformed.memrefOrTensor());
+        MemRefDescriptor descriptor(transformed.source());
         return descriptor.size(rewriter, loc, i);
       }
       // Use constant for static size.
@@ -3042,7 +3042,7 @@ struct DimOpLowering : public ConvertOpToLLVMPattern<memref::DimOp> {
     }
     Value index = dimOp.index();
     int64_t rank = memRefType.getRank();
-    MemRefDescriptor memrefDescriptor(transformed.memrefOrTensor());
+    MemRefDescriptor memrefDescriptor(transformed.source());
     return memrefDescriptor.size(rewriter, loc, index, rank);
   }
 };