Add a shape function library op

[jpienaar]

Summary

Op with mapping from ops to corresponding shape functions for those op in the library and mechanism to associate shape functions to functions. The mapping of operand to shape function is kept separate from the shape functions themselves as the operation is associated to the shape function and not vice versa, and one could have a common library of shape functions that can be used in different contexts.

[rriddle]

mlir/include/mlir/Dialect/Shape/IR/Shape.h

17

Where is the dependency on standard coming from?

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

820

Operation& -> Operation &

Also I'd probably use Operation * here instead given that this is the overwhelming common way of passing an operation, and the one that has the most nice implicit conversions.

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp

25

nit: Move this out of the anonymous namespace.

35

nit: Drop all of the mlir::.

[jpienaar]

Switching to Operation* and moving out anon namespace.

mlir/include/mlir/Dialect/Shape/IR/Shape.h

17

Removed, I needed BuiltinDialect instead for FuncOp.

[mehdi_amini]

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp

32

Walk is recursive, if the assumption is that the library is at the top-level this would be more efficient:

for (auto libraryOp : module.getOps<shape::FunctionLibraryOp>()) {
   ...
}

[mehdi_amini]

mlir/test/Analysis/test-shape-fn-report.mlir

6

This attribute seems redundant with the mapping of the library?

More importantly, this SymbolRef does not seem correct to me: it'll lookup in the current symbol table (current module) and not in the function library.

[jpienaar]

Don't do recursive walk when finding shape fn lib in test pass.

mlir/test/Analysis/test-shape-fn-report.mlir

6

Good question.

The mapping in the library is on ops rather than functions. This is associated with the function and hence attached to the function. The ones below is mapping between operation and associated shape function. I see them connected differently/non-redundantly:

• shape function library function is a shape transfer function/utility function, they are neither op nor computation specific
• mapping in shape.function_library is between an op type and a shape function library function, but associated with the op
• a function can have a shape function attached and that shape function refers to a shape function in the shape function library

But the shape function library I see as independent of the computations, it could be in a separate module even (and that may be more common from other feedback I got :))

What do you mean wrong lookup? It is up to the user to do the lookup. In this case the consumer of the op/attribute would have to know that the shape.function is in the shape.function_library. So it is a convention. Now that would fail if the library was in a different context than this one I believe, but beyond that I think it would work. Now, of course if something blindly were to check symbolrefs then yes this would fail, and I'd need to switch to a string. But I don't think we've generally enforced that all symbolrefattr have to be defined in the same module or even in the same symbol table that the op is defined in.

[tpopp]

Only nits.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

813

Maybe use the OneRegion trait instead to catch any code expecting that trait instead.

828

I think this should be without shape in the op name as that's redundant (shape.shape_fn_lib_terminator)

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp

13–14

Does automatic formatting consider actual ascii values? IR vs Interfaces

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

138

operand without the 's'

[rriddle]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

813

You should never be adding that yourself, it's automatic.

[mehdi_amini]

mlir/test/Analysis/test-shape-fn-report.mlir

6

What do you mean wrong lookup? It is up to the user to do the lookup. In this case the consumer of the op/attribute would have to know that the shape.function is in the shape.function_library. So it is a convention. Now that would fail if the library was in a different context than this one I believe, but beyond that I think it would work. Now, of course if something blindly were to check symbolrefs then yes this would fail, and I'd need to switch to a string. But I don't think we've generally enforced that all symbolrefattr have to be defined in the same module or even in the same symbol table that the op is defined in.

I don't think this is correct, basically my take is that this is not allowed / breaking MLIR symbol resolution assumptions I believe. We could have a symbol with the same name in the module and this would count as a use of this symbol, this would be walked by the symbol use/def chains, etc.
You shouldn't use a SymbolRef but a StringAttr if it does not refer to a symbol in the current module.

[mehdi_amini]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

813

SizedRegion<1> is to express that the region has a single block, not that there is a single region.
(it is also redundant with the SingleBlockImplicitTerminator trait, I wonder if this trait shouldn't be instead placed on the region declaration here)

[jpienaar]

Switched to stringattr & AnyRegion

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp

13–14

I think so, I just rely on clang-format for these.

[herhut]

I am fine either way, just a suggestion.

mlir/test/Analysis/test-shape-fn-report.mlir

6

Now that the shape function library is a symbol table itself and has an (optional) name, you could use fully qualified symbol names here instead of strings. The advantage would be that this composes nicely (multiple shape fun libraries, etc.) and one can use standard means to drop unused shape functions when they are no longer referenced. Could come in handy when model sizes become an issue.

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp

52

You would not need to search for the single shapeFnLib with qualified symbols here but instead could just do the lookup at the scope of the function.

[mehdi_amini]

mlir/test/Analysis/test-shape-fn-report.mlir

6

That's a good point: you marked the shape function library as "Symbol" but the name is optional? I'm even surprised this is allowed.

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp

52

You can't inspect siblings operations though, this can't be used in function passes.

[herhut]

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp

52

That is an interesting restriction and I wonder whether it should be lifted. If I have an operation pass for OpA, I should be allowed to inspect OpB in the parent, because those are not impacted by the operation pass on OpA. With restrictions lifted that way, it would be feasible to lookup shape functions and associations even from a function pass.

[mehdi_amini]

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp

52

It isn't clear to me that function pass and a "FunctionLibraryOp" pass can't run in parallel?

[jpienaar]

Switched to fully qualified names

mlir/test/Analysis/test-shape-fn-report.mlir

6

Module also follows this form, by setting the flag that it optionally defines a symbol. I've removed that for now as Stephan made a good point about allowing multiple.

mlir/include/mlir/Dialect/Shape/IR/Shape.h

 //===- Shape.h - MLIR Shape dialect -----------------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the shape dialect that is used to describe and solve shape
 // relations of MLIR operations using ShapedType.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_SHAPE_IR_SHAPE_H
 #define MLIR_SHAPE_IR_SHAPE_H
 
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

 #ifndef SHAPE_OPS
 #define SHAPE_OPS
 
 include "mlir/Dialect/Shape/IR/ShapeBase.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
+include "mlir/IR/SymbolInterfaces.td"
 
 //===----------------------------------------------------------------------===//
 // Shape op definitions
 //===----------------------------------------------------------------------===//
 
 // Base class for the operation in this dialect
 class Shape_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<ShapeDialect, mnemonic, traits>;
def Shape_WithOp : Shape_Op<"with_shape", [NoSideEffect]> {
   let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueShapeType]>:$operand,
                        Shape_ShapeType:$shape);
   let results = (outs Shape_ValueShapeType:$result);
 
   let assemblyFormat = "operands attr-dict `:` type($operand) `,` type($shape)";
 }
 
 def Shape_YieldOp : Shape_Op<"yield",
-    [HasParent<"ReduceOp">,
+    [HasParent<"ReduceOp, FunctionLibraryOp">,
      NoSideEffect,
      ReturnLike,
      Terminator]> {
   let summary = "Returns the value to parent op";
 
   let arguments = (ins Variadic<AnyType>:$operands);
 
   let builders = [OpBuilderDAG<(ins),
def Shape_CstrRequireOp : Shape_Op<"cstr_require", []> {
   let arguments = (ins I1:$pred, StrAttr:$msg);
   let results = (outs Shape_WitnessType:$result);
 
   let assemblyFormat = "$pred `,` $msg attr-dict";
 
   let hasFolder = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// Shape collection ops.
+//===----------------------------------------------------------------------===//
+
+def Shape_FunctionLibraryOp : Shape_Op<"function_library",
+    [AffineScope, IsolatedFromAbove, NoRegionArguments, SymbolTable,
+     SingleBlockImplicitTerminator<"ShapeFunctionLibraryTerminatorOp">]> {
+  let summary = "Represents shape functions and corresponding ops";
+  let description = [{
+    Represents a list of shape functions and the ops whose shape transfer
+    functions they represent.
+
+    Example:
+
+    ```mlir
+    shape.function_library {
+      func @same_result_shape(%arg: !shape.value_shape) -> !shape.shape {
+        %0 = shape.shape_of %arg : !shape.value_shape -> !shape.shape
+        return %0 : !shape.shape
+      }
+    } mapping {
+      std.atan = @same_result_shape
+    }
+    ```
+  }];
+
+  let arguments = (ins OptionalAttr<SymbolNameAttr>:$sym_name,
+                       OptionalAttr<StrAttr>:$sym_visibility);
+  let arguments = (ins DictionaryAttr:$mapping);
+  let regions = (region SizedRegion<1>:$body);
+  let assemblyFormat = [{
+    attr-dict-with-keyword $body `mapping` $mapping
+  }];
+
+  let extraClassDeclaration = [{
+    /// Returns an associated shape function for an operation if defined.
+    FuncOp getShapeFunction(Operation& op);
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// ShapeFunctionLibraryTerminatorOp
+//===----------------------------------------------------------------------===//
+
+def ShapeFunctionLibraryTerminatorOp : Shape_Op<"shape_fn_lib_terminator",
+    [Terminator, HasParent<"FunctionLibraryOp">]> {
+  let summary = "A pseudo op that marks the end of a shape function library";
+  let description = [{
+    `shape_fn_lib_terminator` is a special pseudo terminator operation for the
+    shape function library. It has no semantic meaning beyond keeping the body
+    well-formed.
+  }];
+  let assemblyFormat = "attr-dict";
+}
+
 #endif // SHAPE_OPS

mlir/lib/Dialect/Shape/IR/Shape.cpp

 //===- Shape.cpp - MLIR Shape Operations ----------------------------------===//
 //
 // 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 "mlir/Dialect/Shape/IR/Shape.h"
 
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/Function.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
OpFoldResult FromExtentsOp::fold(ArrayRef<Attribute> operands) {
   SmallVector<int64_t, 6> extents;
   for (auto attr : operands)
     extents.push_back(attr.cast<IntegerAttr>().getInt());
   Builder builder(getContext());
   return builder.getIndexTensorAttr(extents);
 }
 
 //===----------------------------------------------------------------------===//
+// FunctionLibraryOp
+//===----------------------------------------------------------------------===//
+
+FuncOp FunctionLibraryOp::getShapeFunction(Operation &op) {
+  auto attr = mapping()
+                  .get(op.getName().getIdentifier())
+                  .dyn_cast_or_null<FlatSymbolRefAttr>();
+  if (!attr)
+    return nullptr;
+  return lookupSymbol<FuncOp>(attr);
+}
+
+//===----------------------------------------------------------------------===//
 // GetExtentOp
 //===----------------------------------------------------------------------===//
 
 Optional<int64_t> GetExtentOp::getConstantDim() {
   if (auto constSizeOp = dim().getDefiningOp<ConstSizeOp>())
     return constSizeOp.value().getLimitedValue();
   if (auto constantOp = dim().getDefiningOp<ConstantOp>())
     return constantOp.value().cast<IntegerAttr>().getInt();

mlir/test/Analysis/test-shape-fn-report.mlir

+// RUN: mlir-opt %s --test-shape-function-report -verify-diagnostics
+
+// expected-remark@+1 {{associated shape function: same_result_shape}}
+func @tanh(%arg: tensor<10x20xf32>) -> tensor<10x20xf32>
+    attributes {shape.function = @same_result_shape} {
+  // expected-remark@+1 {{no associated way}}
+  %0 = tanh %arg : tensor<10x20xf32>
+  // expected-remark@+1 {{associated shape function: same_result_shape}}
+  %1 = "test.same_operand_result_type"(%0) : (tensor<10x20xf32>) -> tensor<10x20xf32>
+  return %1 : tensor<10x20xf32>
+}
+
+// The shape function library with some local functions.
+shape.function_library {
+  // Test shape function that returns the shape of input arg as result shape.
+  func @same_result_shape(%arg: !shape.value_shape) -> !shape.shape {
+    %0 = shape.shape_of %arg : !shape.value_shape -> !shape.shape
+    return %0 : !shape.shape
+  }
+} mapping {
+  test.same_operand_result_type = @same_result_shape
+}

mlir/test/lib/Dialect/CMakeLists.txt

 add_subdirectory(Affine)
+add_subdirectory(Shape)
 add_subdirectory(SPIRV)
 add_subdirectory(Test)
 add_subdirectory(Tosa)

mlir/test/lib/Dialect/Shape/CMakeLists.txt

+# Exclude tests from libMLIR.so
+add_mlir_library(MLIRShapeTestPasses
+  TestShapeFunctions.cpp
+
+  EXCLUDE_FROM_LIBMLIR
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Shape
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRPass
+  MLIRShape
+  MLIRSupport
+  )

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp

+//===- TestShapeFunctions.cpp - Passes to test shape function  ------------===//
+//
+// 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 <queue>
+
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/IR/Function.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+/// This is a pass that reports shape functions associated with ops.
+struct ReportShapeFnPass
+    : public PassWrapper<ReportShapeFnPass, OperationPass<ModuleOp>> {
+  void runOnOperation() override;
+};
+
+void ReportShapeFnPass::runOnOperation() {
+  auto module = getOperation();
+
+  // Lookup shape function library.
+  shape::FunctionLibraryOp shapeFnLib = nullptr;
+  auto res = module.getBodyRegion().walk([&](shape::FunctionLibraryOp lib) {
+    if (shapeFnLib) {
+      lib.emitError("duplicate shape library op")
+              .attachNote(shapeFnLib.getLoc())
+          << "previous mapping";
+      return mlir::WalkResult::interrupt();
+    }
+    shapeFnLib = lib;
+    return mlir::WalkResult::advance();
+  });
+  if (res.wasInterrupted())
+    return signalPassFailure();
+
+  // Report the shape function available to refine the op.
+  auto shapeFnId = Identifier::get("shape.function", &getContext());
+  auto remarkShapeFn = [&](Operation &op) {
+    if (op.isKnownTerminator())
+      return;
+    if (auto typeInterface = dyn_cast<mlir::InferTypeOpInterface>(&op)) {
+      op.emitRemark() << "implements InferType op interface";
+    } else if (auto fn = shapeFnLib.getShapeFunction(op)) {
+      op.emitRemark() << "associated shape function: " << fn.getName();
+    } else if (auto attr = op.getAttrOfType<FlatSymbolRefAttr>(shapeFnId)) {
+      auto fn = shapeFnLib.lookupSymbol<FuncOp>(attr);
+      op.emitRemark() << "associated shape function: " << fn.getName();
+    } else {
+      op.emitRemark() << "no associated way to refine shape";
+    }
+  };
+
+  module.getBodyRegion().walk([&](FuncOp func) {
+    // Skip ops in the shape function library.
+    if (isa<shape::FunctionLibraryOp>(func.getParentOp()))
+      return;
+
+    func.walk([&](Operation *op) { remarkShapeFn(*op); });
+  });
+}
+
+} // end anonymous namespace
+
+namespace mlir {
+void registerShapeFunctionTestPasses() {
+  PassRegistration<ReportShapeFnPass>(
+      "test-shape-function-report",
+      "Test pass to report associated shape functions");
+}
+} // namespace mlir

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

 }
 def VariadicWithSameOperandsResult :
       TEST_Op<"variadic_with_same_operand_results",
               [SameOperandsAndResultType]> {
   let arguments = (ins Variadic<AnySignlessInteger>:$operands);
   let results = (outs AnySignlessInteger:$result);
 }
 
+def SameOperandsResultType : TEST_Op<
+    "same_operand_result_type", [SameOperandsAndResultType]> {
+  let arguments = (ins AnyTensor:$operands);
+  let results = (outs AnyTensor:$result);
+}
+
 //===----------------------------------------------------------------------===//
 // Test Results
 //===----------------------------------------------------------------------===//
 
 def MixedNormalVariadicResults : TEST_Op<
     "mixed_normal_variadic_result", [SameVariadicResultSize]> {
   let results = (outs
     Variadic<AnyTensor>:$output1,

mlir/tools/mlir-opt/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   null.cpp
 )
 
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
 set(LLVM_LINK_COMPONENTS
   Core
   Support
   AsmParser
   )
 
 if(MLIR_INCLUDE_TESTS)
   set(test_libs
     MLIRAffineTransformsTestPasses
+    MLIRShapeTestPasses
     MLIRSPIRVTestPasses
     MLIRTestDialect
     MLIRTestIR
     MLIRTestPass
     MLIRTestReducer
     MLIRTestTransforms
     )
 endif()

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

 using namespace llvm;
 using namespace mlir;
 
 // Defined in the test directory, no public header.
 namespace mlir {
 void registerConvertToTargetEnvPass();
 void registerPassManagerTestPass();
 void registerPrintOpAvailabilityPass();
+void registerShapeFunctionTestPasses();
 void registerSideEffectTestPasses();
 void registerSliceAnalysisTestPass();
 void registerSymbolTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAffineLoopUnswitchingPass();
 void registerTestAllReduceLoweringPass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();

 } // namespace test
 } // namespace mlir
 
 #ifdef MLIR_INCLUDE_TESTS
 void registerTestPasses() {
   registerConvertToTargetEnvPass();
   registerPassManagerTestPass();
   registerPrintOpAvailabilityPass();
+  registerShapeFunctionTestPasses();
   registerSideEffectTestPasses();
   registerSliceAnalysisTestPass();
   registerSymbolTestPasses();
   registerTestAffineDataCopyPass();
   registerTestAffineLoopUnswitchingPass();
   registerTestAllReduceLoweringPass();
   registerTestFunc();
   registerTestGpuMemoryPromotionPass();