[llvm-branch-commits] [mlir] 01659f3 - Generalized op transformation logic for weight tensor.
Wen-Heng Chung via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Thu Oct 22 13:20:31 PDT 2020
Author: Wen-Heng (Jack) Chung
Date: 2020-06-05T22:18:20-05:00
New Revision: 01659f382d0ab554d59bfd793a0f9153613c7e20
URL: https://github.com/llvm/llvm-project/commit/01659f382d0ab554d59bfd793a0f9153613c7e20
DIFF: https://github.com/llvm/llvm-project/commit/01659f382d0ab554d59bfd793a0f9153613c7e20.diff
LOG: Generalized op transformation logic for weight tensor.
Add test cases.
Added:
mlir/test/Dialect/MIOpen/lowering_kcyx_nchw_nkhw.mlir
mlir/test/Dialect/MIOpen/lowering_kyxc_nhwc_nhwk.mlir
Modified:
mlir/lib/Dialect/MIOpenOps/LowerMIOpenOps.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/MIOpenOps/LowerMIOpenOps.cpp b/mlir/lib/Dialect/MIOpenOps/LowerMIOpenOps.cpp
index 0378a4c113bf..46083be58a35 100644
--- a/mlir/lib/Dialect/MIOpenOps/LowerMIOpenOps.cpp
+++ b/mlir/lib/Dialect/MIOpenOps/LowerMIOpenOps.cpp
@@ -67,44 +67,53 @@ struct Conv2DOpRewritePattern : public OpRewritePattern<miopen::Conv2DOp> {
llvm::SmallVector<NamedAttribute, 3> transformedFilterAttrs;
// TBD: set layout attribute.
- // TBD: Merge part.
- llvm::SmallVector<NamedAttribute, 5> transformedFilterLayoutPart1Specs;
- transformedFilterLayoutPart1Specs.push_back(rewriter.getNamedAttr("dimensions", ArrayAttr::get({IntegerAttr::get(IntegerType::get(32, op.getContext()), 0)}, op.getContext())));
- transformedFilterLayoutPart1Specs.push_back(rewriter.getNamedAttr("names", ArrayAttr::get({StringAttr::get("gemmK", op.getContext())}, op.getContext())));
- transformedFilterLayoutPart1Specs.push_back(rewriter.getNamedAttr("transformation", StringAttr::get("Merge", op.getContext())));
- transformedFilterLayoutPart1Specs.push_back(rewriter.getNamedAttr("source_dimensions",
- ArrayAttr::get({
- IntegerAttr::get(IntegerType::get(32, op.getContext()), 1),
- IntegerAttr::get(IntegerType::get(32, op.getContext()), 2),
- IntegerAttr::get(IntegerType::get(32, op.getContext()), 3),
- }, op.getContext())));
- transformedFilterLayoutPart1Specs.push_back(rewriter.getNamedAttr("source_names",
- ArrayAttr::get({
- StringAttr::get("c", op.getContext()),
- StringAttr::get("y", op.getContext()),
- StringAttr::get("x", op.getContext())
- }, op.getContext())));
-
- // TBD: Passthrough part.
- llvm::SmallVector<NamedAttribute, 5> transformedFilterLayoutPart2Specs;
- transformedFilterLayoutPart2Specs.push_back(rewriter.getNamedAttr("dimensions", ArrayAttr::get({IntegerAttr::get(IntegerType::get(32, op.getContext()), 1)}, op.getContext())));
- transformedFilterLayoutPart2Specs.push_back(rewriter.getNamedAttr("names", ArrayAttr::get({StringAttr::get("gemmM", op.getContext())}, op.getContext())));
- transformedFilterLayoutPart2Specs.push_back(rewriter.getNamedAttr("transformation", StringAttr::get("PassThrough", op.getContext())));
- transformedFilterLayoutPart2Specs.push_back(rewriter.getNamedAttr("source_dimensions",
- ArrayAttr::get({
- IntegerAttr::get(IntegerType::get(32, op.getContext()), 0),
- }, op.getContext())));
- transformedFilterLayoutPart2Specs.push_back(rewriter.getNamedAttr("source_names",
- ArrayAttr::get({
- StringAttr::get("k", op.getContext())
- }, op.getContext())));
-
- auto transformedFilterLayoutAttr = rewriter.getNamedAttr("layout",
- ArrayAttr::get({
- DictionaryAttr::get(transformedFilterLayoutPart1Specs, op.getContext()),
- DictionaryAttr::get(transformedFilterLayoutPart2Specs, op.getContext())
- }, op.getContext()));
- transformedFilterAttrs.push_back(transformedFilterLayoutAttr);
+ // Weight tensor transformation:
+ // - Part 1: Merge non-K dimensions to dimension 0, name it as gemmK.
+ // - Part 2: PassThrough K dimension to dimension 1, name it as gemmM.
+ {
+ llvm::SmallVector<IntegerAttr, 3> nonKDims;
+ IntegerAttr kDim;
+ llvm::SmallVector<StringAttr, 3> nonKDimNames;
+ StringAttr kDimName;
+ for (unsigned i = 0; i < filterLayoutAttr.size(); ++i) {
+ if (auto strAttr = filterLayoutAttr.getValue()[i].dyn_cast<StringAttr>()) {
+ if (strAttr.getValue() == "k") {
+ kDim = IntegerAttr::get(IntegerType::get(32, op.getContext()), i);
+ kDimName = StringAttr::get(strAttr.getValue(), op.getContext());
+ } else {
+ nonKDims.push_back(IntegerAttr::get(IntegerType::get(32, op.getContext()), i));
+ nonKDimNames.push_back(StringAttr::get(strAttr.getValue(), op.getContext()));
+ }
+ }
+ }
+
+ // Part 1: Merge part.
+ llvm::SmallVector<NamedAttribute, 5> transformedFilterLayoutPart1Specs;
+ transformedFilterLayoutPart1Specs.push_back(rewriter.getNamedAttr("dimensions", ArrayAttr::get({IntegerAttr::get(IntegerType::get(32, op.getContext()), 0)}, op.getContext())));
+ transformedFilterLayoutPart1Specs.push_back(rewriter.getNamedAttr("names", ArrayAttr::get({StringAttr::get("gemmK", op.getContext())}, op.getContext())));
+ transformedFilterLayoutPart1Specs.push_back(rewriter.getNamedAttr("transformation", StringAttr::get("Merge", op.getContext())));
+ transformedFilterLayoutPart1Specs.push_back(rewriter.getNamedAttr("source_dimensions",
+ ArrayAttr::get(ArrayRef<Attribute>(nonKDims.begin(), nonKDims.end()), op.getContext())));
+ transformedFilterLayoutPart1Specs.push_back(rewriter.getNamedAttr("source_names",
+ ArrayAttr::get(ArrayRef<Attribute>(nonKDimNames.begin(), nonKDimNames.end()), op.getContext())));
+
+ // Part 2: Passthrough part.
+ llvm::SmallVector<NamedAttribute, 5> transformedFilterLayoutPart2Specs;
+ transformedFilterLayoutPart2Specs.push_back(rewriter.getNamedAttr("dimensions", ArrayAttr::get({IntegerAttr::get(IntegerType::get(32, op.getContext()), 1)}, op.getContext())));
+ transformedFilterLayoutPart2Specs.push_back(rewriter.getNamedAttr("names", ArrayAttr::get({StringAttr::get("gemmM", op.getContext())}, op.getContext())));
+ transformedFilterLayoutPart2Specs.push_back(rewriter.getNamedAttr("transformation", StringAttr::get("PassThrough", op.getContext())));
+ transformedFilterLayoutPart2Specs.push_back(rewriter.getNamedAttr("source_dimensions",
+ ArrayAttr::get({kDim}, op.getContext())));
+ transformedFilterLayoutPart2Specs.push_back(rewriter.getNamedAttr("source_names",
+ ArrayAttr::get({kDimName}, op.getContext())));
+
+ auto transformedFilterLayoutAttr = rewriter.getNamedAttr("layout",
+ ArrayAttr::get({
+ DictionaryAttr::get(transformedFilterLayoutPart1Specs, op.getContext()),
+ DictionaryAttr::get(transformedFilterLayoutPart2Specs, op.getContext())
+ }, op.getContext()));
+ transformedFilterAttrs.push_back(transformedFilterLayoutAttr);
+ }
// set source_layout attribute.
auto filterSrcLayoutAttr = rewriter.getNamedAttr("source_layout", filterLayoutAttr);
diff --git a/mlir/test/Dialect/MIOpen/lowering_kcyx_nchw_nkhw.mlir b/mlir/test/Dialect/MIOpen/lowering_kcyx_nchw_nkhw.mlir
new file mode 100644
index 000000000000..1f880ee37ef1
--- /dev/null
+++ b/mlir/test/Dialect/MIOpen/lowering_kcyx_nchw_nkhw.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-opt -miopen-lowering %s | FileCheck %s
+
+func @miopen_conv2d_kcyx_nchw_nkhw(%filter : memref<?x?x?x?xf32>, %input : memref<?x?x?x?xf32>, %output : memref<?x?x?x?xf32>) {
+ miopen.conv2d(%filter, %input, %output) {
+ filter_layout = ["k", "c", "y", "x"],
+ input_layout = ["ni", "ci", "hi", "wi"],
+ output_layout = ["no", "ko", "ho", "wo"],
+ dilations = [1, 1],
+ strides = [1, 1],
+ padding = [0, 0]
+ } : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>
+ return
+}
+// CHECK-LABEL: func @miopen_conv2d
+// CHECK-NOT: miopen.conv2d
+// CHECK-NEXT: miopen.transform
+// CHECK-NEXT: miopen.transform
+// CHECK-NEXT: miopen.transform
+// CHECK-NEXT: miopen.transform
+// CHECK-NEXT: miopen.transform
+// CHECK-NEXT: miopen.gridwise_gemm
diff --git a/mlir/test/Dialect/MIOpen/lowering_kyxc_nhwc_nhwk.mlir b/mlir/test/Dialect/MIOpen/lowering_kyxc_nhwc_nhwk.mlir
new file mode 100644
index 000000000000..8450b72ec78a
--- /dev/null
+++ b/mlir/test/Dialect/MIOpen/lowering_kyxc_nhwc_nhwk.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-opt -miopen-lowering -split-input-file %s | FileCheck %s
+
+func @miopen_conv2d_kyxc_nhwc_nhwk(%filter : memref<?x?x?x?xf32>, %input : memref<?x?x?x?xf32>, %output : memref<?x?x?x?xf32>) {
+ miopen.conv2d(%filter, %input, %output) {
+ filter_layout = ["k", "y", "x", "c"],
+ input_layout = ["ni", "hi", "wi", "ci"],
+ output_layout = ["no", "ho", "wo", "ko"],
+ dilations = [1, 1],
+ strides = [1, 1],
+ padding = [0, 0]
+ } : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>
+ return
+}
+// CHECK-LABEL: func @miopen_conv2d
+// CHECK-NOT: miopen.conv2d
+// CHECK-NEXT: miopen.transform
+// CHECK-NEXT: miopen.transform
+// CHECK-NEXT: miopen.transform
+// CHECK-NEXT: miopen.transform
+// CHECK-NEXT: miopen.transform
+// CHECK-NEXT: miopen.gridwise_gemm
More information about the llvm-branch-commits
mailing list