[mlir][Linalg] Lower linalg.reshape to LLVM for the static case

[nicolasvasilache]

Summary

This diff adds lowering of the linalg.reshape op to LLVM.

A new descriptor is created with fields initialized as follows:

1. allocatedPTr, alignedPtr and offset are copied from the source descriptor
2. sizes are copied from the static destination shape
3. strides are copied from the static strides collected with getStridesAndOffset

Only the static case in which the target view conforms to strided memref
semantics is supported. Other cases are left for future work and will be added on
a per-need basis.

[nicolasvasilache]

Remove debug spew.

[merge_guards_bot]

Unit tests: pass. 61273 tests passed, 0 failed and 736 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[merge_guards_bot]

Unit tests: pass. 61273 tests passed, 0 failed and 736 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[ftynse]

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

185

I wonder whether hasStaticShape is the right name here. It only checks the sizes, but not the strides, which is a bit counter-intuitive. Not for this commit certainly...

195

Can reshape do permutations using affine maps? It does not seem supported here, neither it is guarded against in the matching.

204

MemRefDescriptor has setConstantSize for this purposes, consider replicating it in BaseViewConversionHelper ?

[nicolasvasilache]

Address review comments.

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

185

I am not sure I follow. If the destination does not have static shape we can early exit and no need to compute the strides.
If it does, we still need to check that it has static strides, which happens below.

195

No, reshape can only collapse contiguous dims and this is enforced in the op verifier.

[merge_guards_bot]

Unit tests: pass. 61305 tests passed, 0 failed and 736 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[ftynse]

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

185

I find it confusing that strides are not included in "shape", but it's a separate discussion.

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

public:
   Value allocatedPtr() { return d.allocatedPtr(rewriter(), loc()); }
   void setAllocatedPtr(Value v) { d.setAllocatedPtr(rewriter(), loc(), v); }
   Value alignedPtr() { return d.alignedPtr(rewriter(), loc()); }
   void setAlignedPtr(Value v) { d.setAlignedPtr(rewriter(), loc(), v); }
   Value offset() { return d.offset(rewriter(), loc()); }
   void setOffset(Value v) { d.setOffset(rewriter(), loc(), v); }
   Value size(unsigned i) { return d.size(rewriter(), loc(), i); }
   void setSize(unsigned i, Value v) { d.setSize(rewriter(), loc(), i, v); }
+  void setConstantSize(unsigned i, int64_t v) {
+    d.setConstantSize(rewriter(), loc(), i, v);
+  }
   Value stride(unsigned i) { return d.stride(rewriter(), loc(), i); }
   void setStride(unsigned i, Value v) { d.setStride(rewriter(), loc(), i, v); }
+  void setConstantStride(unsigned i, int64_t v) {
+    d.setConstantStride(rewriter(), loc(), i, v);
+  }
 
   operator Value() { return d; }
 
 private:
   OpBuilder &rewriter() { return ScopedContext::getBuilder(); }
   Location loc() { return ScopedContext::getLocation(); }
 
   MemRefDescriptor d;
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     desc = insertvalue(desc, adaptor.min(), rewriter.getI64ArrayAttr(0));
     desc = insertvalue(desc, adaptor.max(), rewriter.getI64ArrayAttr(1));
     desc = insertvalue(desc, adaptor.step(), rewriter.getI64ArrayAttr(2));
     rewriter.replaceOp(op, desc);
     return matchSuccess();
   }
 };
 
+// ReshapeOp creates a new view descriptor of the proper rank.
+// For now, the only conversion supported is for target MemRef with static sizes
+// and strides.
+class ReshapeOpConversion : public LLVMOpLowering {
+public:
+  explicit ReshapeOpConversion(MLIRContext *context,
+                               LLVMTypeConverter &lowering_)
+      : LLVMOpLowering(ReshapeOp::getOperationName(), context, lowering_) {}
+
+  PatternMatchResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto reshapeOp = cast<ReshapeOp>(op);
+    MemRefType dstType = reshapeOp.getResult().getType().cast<MemRefType>();
+
+    if (!dstType.hasStaticShape())
+      return matchFailure();
+
+    int64_t offset;
+    SmallVector<int64_t, 4> strides;
+    auto res = getStridesAndOffset(dstType, strides, offset);
+    if (failed(res) || llvm::any_of(strides, [](int64_t val) {
+          return ShapedType::isDynamicStrideOrOffset(val);
+        }))
+      return matchFailure();
+
+    edsc::ScopedContext context(rewriter, op->getLoc());
+    ReshapeOpOperandAdaptor adaptor(operands);
+    BaseViewConversionHelper baseDesc(adaptor.view());
+    BaseViewConversionHelper desc(lowering.convertType(dstType));
+    desc.setAllocatedPtr(baseDesc.allocatedPtr());
+    desc.setAlignedPtr(baseDesc.alignedPtr());
+    desc.setOffset(baseDesc.offset());
+    for (auto en : llvm::enumerate(dstType.getShape()))
+      desc.setConstantSize(en.index(), en.value());
+    for (auto en : llvm::enumerate(strides))
+      desc.setConstantStride(en.index(), en.value());
+    rewriter.replaceOp(op, {desc});
+    return matchSuccess();
+  }
+};
+
 /// Conversion pattern that transforms a linalg.slice op into:
 ///   1. A function entry `alloca` operation to allocate a ViewDescriptor.
 ///   2. A load of the ViewDescriptor from the pointer allocated in 1.
 ///   3. Updates to the ViewDescriptor to introduce the data ptr, offset, size
 ///      and stride corresponding to the region of memory within the bounds of
 ///      the parent view.
 ///   4. A store of the resulting ViewDescriptor to the alloca'ed pointer.
 /// The linalg.slice op is replaced by the alloca'ed pointer.
patterns.insert<CopyTransposeConversion, LinalgOpConversion<ConvOp>,
                   LinalgOpConversion<MatmulOp>, LinalgOpConversion<MatvecOp>>(
       ctx);
 }
 
 /// Populate the given list with patterns that convert from Linalg to LLVM.
 void mlir::populateLinalgToLLVMConversionPatterns(
     LinalgTypeConverter &converter, OwningRewritePatternList &patterns,
     MLIRContext *ctx) {
-  patterns.insert<RangeOpConversion, SliceOpConversion, TransposeOpConversion,
-                  YieldOpConversion>(ctx, converter);
+  patterns.insert<RangeOpConversion, ReshapeOpConversion, SliceOpConversion,
+                  TransposeOpConversion, YieldOpConversion>(ctx, converter);
 }
 
 namespace {
 struct ConvertLinalgToLLVMPass : public ModulePass<ConvertLinalgToLLVMPass> {
   void runOnModule() override;
 };
 } // namespace
 

mlir/test/Dialect/Linalg/llvm.mlir

^bb0(%i: index, %j: index, %k: index,
       %d = vector.outerproduct %a, %b, %c: !vector_type_A, !vector_type_B
       linalg.yield %d: !vector_type_C
   } : !matrix_type_A, !matrix_type_B, !matrix_type_C
   return
 }
 // CHECK-LABEL: func @matmul_vec_indexed(
 //   CHECK: %[[ZERO:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
 //   CHECK: llvm.call @external_indexed_outerproduct_matmul(%[[ZERO]], %[[ZERO]], %[[ZERO]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.i64, !llvm.i64, !llvm.i64, !llvm<"{ <4 x float>*, <4 x float>*, i64, [2 x i64], [2 x i64] }*">, !llvm<"{ <4 x float>*, <4 x float>*, i64, [2 x i64], [2 x i64] }*">, !llvm<"{ [4 x <4 x float>]*, [4 x <4 x float>]*, i64, [2 x i64], [2 x i64] }*">) -> ()
+
+func @reshape_static(%arg0: memref<3x4x5xf32>) {
+  // Reshapes that expand and collapse back a contiguous tensor with some 1's.
+  %0 = linalg.reshape %arg0 [(i, j, k, l, m) -> (i, j),
+                             (i, j, k, l, m) -> (k),
+                             (i, j, k, l, m) -> (l, m)] :
+    memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
+  %r0 = linalg.reshape %0 [(i, j, k, l, m) -> (i, j),
+                           (i, j, k, l, m) -> (k),
+                           (i, j, k, l, m) -> (l, m)] :
+    memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
+  return
+}
+// CHECK-LABEL: func @reshape_static(
+//       CHECK: llvm.mlir.undef : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.extractvalue {{.*}}[0] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.insertvalue {{.*}}[0] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.extractvalue {{.*}}[1] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.insertvalue {{.*}}[1] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.extractvalue {{.*}}[2] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.insertvalue {{.*}}[2] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.mlir.constant(1 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[3, 0] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.mlir.constant(3 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[3, 1] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.mlir.constant(4 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[3, 2] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.mlir.constant(1 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[3, 3] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.mlir.constant(5 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[3, 4] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.mlir.constant(60 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[4, 0] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.mlir.constant(20 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[4, 1] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.mlir.constant(5 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[4, 2] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.mlir.constant(5 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[4, 3] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.mlir.constant(1 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[4, 4] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.mlir.undef : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.extractvalue {{.*}}[0] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.insertvalue {{.*}}[0] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.extractvalue {{.*}}[1] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.insertvalue {{.*}}[1] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.extractvalue {{.*}}[2] : !llvm<"{ float*, float*, i64, [5 x i64], [5 x i64] }">
+//       CHECK: llvm.insertvalue {{.*}}[2] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.mlir.constant(3 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[3, 0] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.mlir.constant(4 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[3, 1] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.mlir.constant(5 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[3, 2] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.mlir.constant(20 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[4, 0] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.mlir.constant(5 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[4, 1] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">
+//       CHECK: llvm.mlir.constant(1 : index) : !llvm.i64
+//       CHECK: llvm.insertvalue {{.*}}[4, 2] : !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }">