[mlir] Add clone method to ShapedType

[jpienaar]

Summary

Allow clients to create a new ShapedType of the same "container" type but with
different element or shape. First use case is when refining shape during shape
inference without needing to consider which ShapedType is being refined.

[rriddle]

mlir/test/lib/Transforms/TestUpdateShapedTypes.cpp

9

Honestly this would be much better suited to a directed unit test instead of this.

[jpienaar]

Switch to unittest

[jpienaar]

mlir/test/lib/Transforms/TestUpdateShapedTypes.cpp

9

Yeah I was between those, the unit test is more cumbersome to add tests for but not the most frequent occurrence (also not all linear maps need to be tested :))

mlir/include/mlir/IR/BuiltinTypes.h

public:
 
   // TODO: merge these two special values in a single one used everywhere.
   // Unfortunately, uses of `-1` have crept deep into the codebase now and are
   // hard to track.
   static constexpr int64_t kDynamicSize = -1;
   static constexpr int64_t kDynamicStrideOrOffset =
       std::numeric_limits<int64_t>::min();
 
+  /// Return clone of this type with new shape and element type.
+  ShapedType clone(ArrayRef<int64_t> shape, Type elementType);
+  ShapedType clone(ArrayRef<int64_t> shape);
+  ShapedType clone(Type elementType);
+
   /// Return the element type.
   Type getElementType() const;
 
   /// If an element type is an integer or a float, return its width. Otherwise,
   /// abort.
   unsigned getElementTypeBitWidth() const;
 
   /// If it has static shape, return the number of elements. Otherwise, abort.

mlir/lib/IR/BuiltinTypes.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // ShapedType
 //===----------------------------------------------------------------------===//
 constexpr int64_t ShapedType::kDynamicSize;
 constexpr int64_t ShapedType::kDynamicStrideOrOffset;
 
+ShapedType ShapedType::clone(ArrayRef<int64_t> shape, Type elementType) {
+  if (auto other = dyn_cast<MemRefType>()) {
+    MemRefType::Builder b(other);
+    b.setShape(shape);
+    b.setElementType(elementType);
+    return b;
+  }
+
+  if (isa<TensorType>())
+    return RankedTensorType::get(shape, elementType);
+
+  if (isa<VectorType>())
+    return VectorType::get(shape, elementType);
+
+  llvm_unreachable("Unhandled ShapedType clone case");
+}
+
+ShapedType ShapedType::clone(ArrayRef<int64_t> shape) {
+  if (auto other = dyn_cast<MemRefType>()) {
+    MemRefType::Builder b(other);
+    b.setShape(shape);
+    return b;
+  }
+
+  if (isa<TensorType>())
+    return RankedTensorType::get(shape, getElementType());
+
+  if (isa<VectorType>())
+    return VectorType::get(shape, getElementType());
+
+  llvm_unreachable("Unhandled ShapedType clone case");
+}
+
+ShapedType ShapedType::clone(Type elementType) {
+  if (auto other = dyn_cast<MemRefType>()) {
+    MemRefType::Builder b(other);
+    b.setElementType(elementType);
+    return b;
+  }
+
+  if (isa<TensorType>()) {
+    if (hasRank())
+      return RankedTensorType::get(getShape(), elementType);
+    return UnrankedTensorType::get(elementType);
+  }
+
+  if (isa<VectorType>()) {
+    return VectorType::get(getShape(), elementType);
+  }
+  llvm_unreachable("Unhandled ShapedType clone hit");
+}
+
 Type ShapedType::getElementType() const {
   return static_cast<ImplType *>(impl)->elementType;
 }
 
 unsigned ShapedType::getElementTypeBitWidth() const {
   return getElementType().getIntOrFloatBitWidth();
 }
 

mlir/test/IR/update-shape.mlir

+// RUN: mlir-opt %s -test-update-element-or-type --mlir-print-local-scope | FileCheck %s
+
+func @succeededSameOperandsElementType() {
+  // CHECK: tensor<121xf32>
+  "test.update_variadic_return"() {type_attr = f32} : () -> (tensor<121xi32>)
+  // CHECK: vector<16xf32>
+  "test.update_variadic_return"() {shape_attr = dense<16> : vector<1xi32>} : () -> (vector<121xf32>)
+  // CHECK: memref<64x256xi32, affine_map<(d0, d1) -> (d1, d0)>>
+  "test.update_variadic_return"() {type_attr = i32} : () -> (memref<64x256xf32, affine_map<(d0, d1) -> (d1, d0)>>)
+  // CHECK: memref<10x20xi32, affine_map<(d0, d1) -> (d0, d1 * 2)>>
+  "test.update_variadic_return"() {shape_attr = dense<[10, 20]> : tensor<2xi32>,  type_attr = i32} : () -> (memref<64x256xf32, affine_map<(d0, d1) -> (d0, 2 * d1)>>)
+  return
+}

mlir/test/lib/Dialect/Test/TestOps.td

 }
 
 def MultiTensorRankOf : TEST_Op<"multi_tensor_rank_of"> {
   let arguments = (ins
     TensorRankOf<[I8, I32, F32], [0, 1]>:$arg0
   );
 }
 
+def UpdatedVariadicReturnOp : TEST_Op<"update_variadic_return"> {
+  let arguments = (ins
+      OptionalAttr<TypeAttr>:$type_attr,
+      OptionalAttr<I32ElementsAttr>:$shape_attr
+  );
+  let results = (outs Variadic<AnyShaped>);
+}
+
 def TEST_TestType : DialectType<Test_Dialect,
     CPred<"$_self.isa<::mlir::test::TestType>()">, "test">,
     BuildableType<"$_builder.getType<::mlir::test::TestType>()">;
 
 //===----------------------------------------------------------------------===//
 // Test Symbols
 //===----------------------------------------------------------------------===//
 

mlir/test/lib/Transforms/TestUpdateShapedTypes.cpp

+//===- TestUpdateShapedTypes.cpp - Pass to test updating element & shape --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestDialect.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+/// Pass that changes element or shape of shaped types.
+struct TestChangeElementOrShape
+    : public PassWrapper<TestChangeElementOrShape, OperationPass<ModuleOp>> {
+
+  void runOnOperation() override {
+    getOperation().walk([&](test::UpdatedVariadicReturnOp op) {
+      Optional<Type> elementType = op.type_attr();
+      Optional<DenseIntElementsAttr> shapeAttr = op.shape_attr();
+      if (!elementType && !shapeAttr)
+        return;
+
+      for (auto result : op.getOperation()->getOpResults()) {
+        SmallVector<int64_t> shape;
+        if (shapeAttr) {
+          shape.reserve(shapeAttr->size());
+          for (APInt v : shapeAttr->getIntValues())
+            shape.push_back(v.getSExtValue());
+        }
+
+        if (elementType && shapeAttr) {
+          result.setType(
+              result.getType().cast<ShapedType>().clone(shape, *elementType));
+        } else if (elementType) {
+          result.setType(
+              result.getType().cast<ShapedType>().clone(*elementType));
+        } else if (shapeAttr) {
+          result.setType(result.getType().cast<ShapedType>().clone(shape));
+        }
+      }
+    });
+  }
+};
+
+} // end anonymous namespace
+
+namespace mlir {
+namespace test {
+void registerTestChangeElementOrShape() {
+  PassRegistration<TestChangeElementOrShape> pass(
+      "test-update-element-or-type", "Changes ShapedType's element or shape");
+}
+} // namespace test
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerConvertCallOpPass();
 void registerInliner();
 void registerMemRefBoundCheck();
 void registerPatternsTestPass();
 void registerSimpleParametricTilingPass();
 void registerTestAffineLoopParametricTilingPass();
 void registerTestAliasAnalysisPass();
 void registerTestCallGraphPass();
+void registerTestChangeElementOrShape();
 void registerTestConstantFold();
 void registerTestConvVectorization();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
 void registerTestDecomposeCallGraphTypes();
 void registerTestDialect(DialectRegistry &);
 void registerTestDominancePass();
 void registerTestDynamicPipelinePass();
void registerTestPasses() {
   test::registerConvertCallOpPass();
   test::registerInliner();
   test::registerMemRefBoundCheck();
   test::registerPatternsTestPass();
   test::registerSimpleParametricTilingPass();
   test::registerTestAffineLoopParametricTilingPass();
   test::registerTestAliasAnalysisPass();
   test::registerTestCallGraphPass();
+  test::registerTestChangeElementOrShape();
   test::registerTestConstantFold();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
   test::registerTestConvertGPUKernelToCubinPass();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED
   test::registerTestConvertGPUKernelToHsacoPass();
 #endif
   test::registerTestConvVectorization();