[MLIR][Shape] Lower `shape_of` for unranked tensors

[frgossen]

Summary

Lower shape_of for unranked tensors.
Materializes shape in stack-allocated memory.

[herhut]

mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

73

nit: maybe to for loop so it reads better?

mlir/lib/IR/StandardTypes.cpp

401

This is intentional. Memrefs may not have index type, as index has no size. @ftynse Has this changed?

So I guess the only way out is to allocate this as int64 and put casts around it.

[frgossen]

Materialize as memref<i64>

[herhut]

With the missing cast, this looks fine to me. Maybe @ftynse has an opinion though.

mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

116

You now need to cast the result back to index, as it is a ?xi64 now. All a bit suboptimal but yeah...

[frgossen]

mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

116

I could not see any operation in std that supports this.
tensor_cast requires the element type to be equal, index_cast does not have a lowering to LLVM.
In fact, I could not find any lowering or legalization pass for tensors.
Is that generally missing or is there a trick?

[frgossen]

Add missing cast

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

def DimOp : Std_Op<"dim", [NoSideEffect]> {
   let results = (outs Index:$result);
 
   let assemblyFormat = [{
     attr-dict $memrefOrTensor `,` $index `:` type($memrefOrTensor)
   }];
 
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
-              "Value memrefOrTensor, int64_t index">
+              "Value memrefOrTensor, int64_t index">,
+    OpBuilder<"OpBuilder &builder, OperationState &result, "
+              "Value memrefOrTensor, Value index">
   ];
 
   let extraClassDeclaration = [{
     /// Helper function to get the index as a simple integer if it is constant.
     Optional<int64_t> getConstantIndex();
   }];
 
   let hasFolder = 1;

mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

auto loop = rewriter.create<scf::ForOp>(
         b.create<scf::YieldOp>(loc, mappedResults);
       });
 
   rewriter.replaceOp(reduceOp, loop.getResults());
   return success();
 }
 
 namespace {
+/// Converts `shape_of` to for unranked tensors.
+class ShapeOfOpConverter : public OpConversionPattern<ShapeOfOp> {
+public:
+  using OpConversionPattern<ShapeOfOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ShapeOfOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+} // namespace
+
+LogicalResult
+ShapeOfOpConverter::matchAndRewrite(ShapeOfOp op, ArrayRef<Value> operands,
+                                    ConversionPatternRewriter &rewriter) const {
+  ShapeOfOp::Adaptor transformed(operands);
+  auto tensorVal = transformed.arg();
+  auto tensorTy = tensorVal.getType();
+
+  // For ranked tensors `shape_of` lowers to `std` and the pattern can be
+  // found in the corresponding pass.
+  if (tensorTy.isa<RankedTensorType>())
+    return failure();
+
+  // Allocate stack memory.
+  auto loc = op.getLoc();
+  auto rankVal = rewriter.create<RankOp>(loc, tensorVal);
+  auto indexTy = rewriter.getIndexType();
+  auto memTy = MemRefType::get({ShapedType::kDynamicSize}, indexTy);
+  auto memVal = rewriter.create<AllocaOp>(loc, memTy, ValueRange({rankVal}));
+
+  // Copy shape extents to stack-allocated memory.
+  auto zeroVal = rewriter.create<ConstantIndexOp>(loc, 0);
+  auto oneVal = rewriter.create<ConstantIndexOp>(loc, 1);
+  rewriter.create<scf::ForOp>(
+      loc, zeroVal, rankVal, oneVal, ValueRange(),
+      [&](OpBuilder &b, Location loc, Value iVal, ValueRange args) {
+        auto dimVal = b.create<DimOp>(loc, tensorVal, iVal);
+        b.create<StoreOp>(loc, dimVal, memVal, ValueRange({iVal}));
+        b.create<scf::YieldOp>(loc);
+      });
+
+  // Load extents to tensor value.
+  rewriter.replaceOpWithNewOp<TensorLoadOp>(op.getOperation(), memVal);
+  return success();
+}
+
+namespace {
 struct ConvertShapeToSCFPass
     : public ConvertShapeToSCFBase<ConvertShapeToSCFPass> {
   void runOnFunction() override;
 };
 } // namespace
 
 void ConvertShapeToSCFPass::runOnFunction() {
   MLIRContext &ctx = getContext();
 
+  // Populate conversion patterns.
   OwningRewritePatternList patterns;
   populateShapeToSCFConversionPatterns(patterns, &ctx);
 
+  // Setup target legality.
   ConversionTarget target(getContext());
   target.addLegalDialect<ShapeDialect, scf::SCFDialect, StandardOpsDialect>();
-  target.addIllegalOp<ReduceOp>();
-  if (failed(mlir::applyPartialConversion(getFunction(), target, patterns)))
+  target.addIllegalOp<ReduceOp, ShapeOfOp>();
+
+  // Apply conversion.
+  if (failed(applyPartialConversion(getFunction(), target, patterns)))
     signalPassFailure();
 }
 
 void mlir::populateShapeToSCFConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
-  patterns.insert<ReduceOpConverter>(ctx);
+  patterns.insert<ReduceOpConverter, ShapeOfOpConverter>(ctx);
 }
 
 std::unique_ptr<FunctionPass> mlir::createConvertShapeToSCFPass() {
   return std::make_unique<ConvertShapeToSCFPass>();
 }

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 //===----------------------------------------------------------------------===//
 // DimOp
 //===----------------------------------------------------------------------===//
 
 void DimOp::build(OpBuilder &builder, OperationState &result,
                   Value memrefOrTensor, int64_t index) {
   auto loc = result.location;
   Value indexValue = builder.create<ConstantIndexOp>(loc, index);
+  build(builder, result, memrefOrTensor, indexValue);
+}
+
+void DimOp::build(OpBuilder &builder, OperationState &result,
+                  Value memrefOrTensor, Value index) {
   auto indexTy = builder.getIndexType();
-  build(builder, result, indexTy, memrefOrTensor, indexValue);
+  build(builder, result, indexTy, memrefOrTensor, index);
 }
 
 Optional<int64_t> DimOp::getConstantIndex() {
   if (auto constantOp = index().getDefiningOp<ConstantOp>())
     return constantOp.getValue().cast<IntegerAttr>().getInt();
   return {};
 }
 

mlir/lib/IR/StandardTypes.cpp

 /// pass in an instance of UnknownLoc.
 MemRefType MemRefType::getImpl(ArrayRef<int64_t> shape, Type elementType,
                                ArrayRef<AffineMap> affineMapComposition,
                                unsigned memorySpace,
                                Optional<Location> location) {
   auto *context = elementType.getContext();
 
   // Check that memref is formed from allowed types.
-  if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>() &&
+  if (!elementType.isIntOrIndexOrFloat() && !elementType.isa<VectorType>() &&
       !elementType.isa<ComplexType>())
     return emitOptionalError(location, "invalid memref element type"),
            MemRefType();
-
   for (int64_t s : shape) {
     // Negative sizes are not allowed except for `-1` that means dynamic size.
     if (s < -1)
       return emitOptionalError(location, "invalid memref size"), MemRefType();
   }
 
   // Check that the structure of the composition is valid, i.e. that each
   // subsequent affine map has as many inputs as the previous map has results.

mlir/test/Conversion/ShapeToSCF/shape-to-scf.mlir

 // CHECK-NEXT: [[RESULT:%.*]] = scf.for [[I:%.*]] = [[C0]] to [[SIZE]]
 // CHECK-SAME:       step [[C1]] iter_args([[ACC:%.*]] = [[SHAPE_C1]])
 // CHECK-NEXT:   [[EXTENT_INDEX:%.*]] = extract_element [[EXTENTS]]{{\[}}[[I]]]
 // CHECK-NEXT:   [[EXTENT:%.*]] = shape.index_to_size [[EXTENT_INDEX]]
 // CHECK-NEXT:   [[NEW_ACC:%.*]] = shape.mul [[ACC]], [[EXTENT]]
 // CHECK-NEXT:   scf.yield [[NEW_ACC]] : !shape.size
 // CHECK-NEXT: }
 // CHECK-NEXT: return [[RESULT]] : !shape.size
+
+// -----
+
+// Lower `shape_of` for unranked tensors.
+// CHECK-LABEL: @shape_of_unranked
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>)
+func @shape_of_unranked(%arg : tensor<*xf32>) {
+  // CHECK-DAG: %[[RANK:.*]] = rank %[[ARG]] : tensor<*xf32>
+  // CHECK-DAG: %[[SHAPE_MEM:.*]] = alloca(%[[RANK]]) : memref<?xindex>
+  // CHECK-DAG: %[[C0:.*]] = constant 0 : index
+  // CHECK-DAG: %[[C1:.*]] = constant 1 : index
+  // CHECK:     scf.for %[[I:.*]] = %[[C0]] to %[[RANK]] step %[[C1]] {
+  // CHECK-DAG:   %[[DIM:.]] = dim %[[ARG]], %[[I]] : tensor<*xf32>
+  // CHECK-DAG:   store %[[DIM]], %[[SHAPE_MEM]][%[[I]]] : memref<?xindex>
+  // CHECK:     }
+  // CHECK-DAG: %[[SHAPE:.*]] = tensor_load %[[SHAPE_MEM]] : memref<?xindex>
+  %shape = shape.shape_of %arg : tensor<*xf32>
+  return
+}
+