Ferdinand.cherry picks from fused ops into ljfitz.fxml 1172 memory schedule

include/CMakeLists.txt

@@ -5,4 +5,4 @@
 #
 # (c) Copyright 2021 Xilinx Inc.
 
-add_subdirectory(xten)
+add_subdirectory(xten)

include/xten/CMakeLists.txt

@@ -7,4 +7,4 @@
 
 add_subdirectory(Dialect)
 add_subdirectory(Conversion)
-add_subdirectory(Transform)
+add_subdirectory(Transform)

include/xten/Dialect/CMakeLists.txt

@@ -4,5 +4,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 #
 # (c) Copyright 2021 Xilinx Inc.
+# (c) Copyright 2022 Advanced Micro Devices, Inc.
 
-add_subdirectory(XTen)
+add_subdirectory(XTen)
+add_subdirectory(XTenNN)

include/xten/Dialect/XTen/XTenDataflowUtils.h

@@ -17,7 +17,7 @@
 
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 
 using namespace mlir;

include/xten/Dialect/XTen/XTenOps.h

@@ -20,6 +20,7 @@
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #include "torch-mlir/Dialect/Torch/IR/TorchDialect.h"
+#include "torch-mlir/Dialect/Torch/IR/TorchTypes.h"
 
 #define GET_OP_CLASSES
 #include "xten/Dialect/XTen/XTenOps.h.inc"

include/xten/Dialect/XTen/XTenOps.td

@@ -50,7 +50,7 @@ def XTen_AddOp: XTen_Op<"add", []>,
   }];
 }
 
-def XTen_MMOp: XTen_Op<"mm", [NoSideEffect]>,
+def XTen_MMOp: XTen_Op<"mm", [Pure]>,
               Results<(outs AnyTorchTensorType)> {
   let arguments = (
     ins AnyTorchTensorType:$x,
@@ -91,7 +91,7 @@ def XTen_SoftmaxOp: XTen_Op<"softmax", []>,
   }];
 }
 
-def XTen_GlobalAveragePool2D: XTen_Op<"globalaveragepool2d", [NoSideEffect]>,
+def XTen_GlobalAveragePool2D: XTen_Op<"globalaveragepool2d", [Pure]>,
                Results<(outs AnyTorchTensorType:$output)> {
   let arguments = (
     ins AnyTorchTensorType:$input
@@ -116,7 +116,7 @@ def XTen_NoOp: XTen_Op<"noop", []>,
   }];
 }
 
-def XTen_Conv2dOp: XTen_Op<"conv2d", [NoSideEffect]>,
+def XTen_Conv2dOp: XTen_Op<"conv2d", [Pure]>,
                                 Results<(outs AnyTorchTensorType:$result)> {
   let arguments = (
     ins AnyTorchTensorType:$input,
@@ -140,7 +140,7 @@ def XTen_Conv2dOp: XTen_Op<"conv2d", [NoSideEffect]>,
 }
 
 // TODO what happens when we have both?
-def XTen_PartialConv2dOp: XTen_Op<"partialconv2d", [NoSideEffect]>{
+def XTen_PartialConv2dOp: XTen_Op<"partialconv2d", [Pure]>{
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchOptionalTensorType:$PartialIn,
@@ -169,7 +169,7 @@ def XTen_PartialConv2dOp: XTen_Op<"partialconv2d", [NoSideEffect]>{
 }
 
 
-def XTen_Conv2dReLUOp: XTen_Op<"conv2d_relu", [NoSideEffect]>,
+def XTen_Conv2dReLUOp: XTen_Op<"conv2d_relu", [Pure]>,
                                    Results<(outs AnyTorchTensorType)> {
   let arguments = (
     ins AnyTorchTensorType:$input,
@@ -193,7 +193,7 @@ def XTen_Conv2dReLUOp: XTen_Op<"conv2d_relu", [NoSideEffect]>,
 }
 
 def XTen_PartialConv2dReLUOp: XTen_Op<"partialconv2d_relu",
-                                [NoSideEffect]> {
+                                [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchOptionalTensorType:$PartialIn,
@@ -221,7 +221,7 @@ def XTen_PartialConv2dReLUOp: XTen_Op<"partialconv2d_relu",
 	}];
 }
 
-def XTen_Conv2dBatchNormReLUOp: XTen_Op<"conv2d_bn_relu", [NoSideEffect]>,
+def XTen_Conv2dBatchNormReLUOp: XTen_Op<"conv2d_bn_relu", [Pure]>,
                                             Results<(outs AnyTorchTensorType)> {
   let arguments = (
     ins AnyTorchTensorType:$input,
@@ -253,7 +253,7 @@ def XTen_Conv2dBatchNormReLUOp: XTen_Op<"conv2d_bn_relu", [NoSideEffect]>,
 }
 
 def XTen_PartialConv2dBatchNormReLUOp: XTen_Op<"partialconv2d_bn_relu",
-                                        [NoSideEffect]> {
+                                        [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchOptionalTensorType:$PartialIn,
@@ -288,7 +288,7 @@ def XTen_PartialConv2dBatchNormReLUOp: XTen_Op<"partialconv2d_bn_relu",
 	}];
 }
 
-def XTen_ConcatOp: XTen_Op<"concat", [NoSideEffect]>,
+def XTen_ConcatOp: XTen_Op<"concat", [Pure]>,
                                 Results<(outs AnyTorchTensorType)> {
   let arguments = (
     ins Variadic<AnyTorchTensorType>:$inputs,
@@ -307,7 +307,7 @@ def XTen_ConcatOp: XTen_Op<"concat", [NoSideEffect]>,
 }
 
 // TODO Proper verifier for this operation?
-def XTen_SplitOp: XTen_Op<"split", [NoSideEffect]> {
+def XTen_SplitOp: XTen_Op<"split", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         XTen_AnyScalar:$dim
@@ -328,7 +328,7 @@ def XTen_SplitOp: XTen_Op<"split", [NoSideEffect]> {
 	}];
 }
 
-def XTen_Conv2dLReLUOp: XTen_Op<"conv2d_lrelu", [NoSideEffect]>,
+def XTen_Conv2dLReLUOp: XTen_Op<"conv2d_lrelu", [Pure]>,
                                    Results<(outs AnyTorchTensorType)> {
   let arguments = (
     ins AnyTorchTensorType:$input,
@@ -352,7 +352,7 @@ def XTen_Conv2dLReLUOp: XTen_Op<"conv2d_lrelu", [NoSideEffect]>,
 	}];
 }
 
-def XTen_Conv2dLReLUPadOp: XTen_Op<"conv2d_lrelu_pad", [NoSideEffect]>,
+def XTen_Conv2dLReLUPadOp: XTen_Op<"conv2d_lrelu_pad", [Pure]>,
                                    Results<(outs AnyTorchTensorType)> {
   let arguments = (
     ins AnyTorchTensorType:$input,
@@ -379,7 +379,7 @@ def XTen_Conv2dLReLUPadOp: XTen_Op<"conv2d_lrelu_pad", [NoSideEffect]>,
 }
 
 
-def XTen_Conv2dReLUPadOp: XTen_Op<"conv2d_relu_pad", [NoSideEffect]>,
+def XTen_Conv2dReLUPadOp: XTen_Op<"conv2d_relu_pad", [Pure]>,
                                    Results<(outs AnyTorchTensorType)> {
   let arguments = (
     ins AnyTorchTensorType:$input,
@@ -405,7 +405,7 @@ def XTen_Conv2dReLUPadOp: XTen_Op<"conv2d_relu_pad", [NoSideEffect]>,
 }
 
 
-def XTen_Conv2dReLUMaxPoolOp: XTen_Op<"conv2d_relu_maxpool", [NoSideEffect]> {
+def XTen_Conv2dReLUMaxPoolOp: XTen_Op<"conv2d_relu_maxpool", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchTensorType:$weight,
@@ -436,7 +436,7 @@ def XTen_Conv2dReLUMaxPoolOp: XTen_Op<"conv2d_relu_maxpool", [NoSideEffect]> {
 	}];
 }
 
-def XTen_Conv2dReLUPadMaxPoolOp: XTen_Op<"conv2d_relu_pad_maxpool", [NoSideEffect]> {
+def XTen_Conv2dReLUPadMaxPoolOp: XTen_Op<"conv2d_relu_pad_maxpool", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchTensorType:$weight,
@@ -470,7 +470,7 @@ def XTen_Conv2dReLUPadMaxPoolOp: XTen_Op<"conv2d_relu_pad_maxpool", [NoSideEffec
 }
 
 
-def XTen_Conv2dLReLUMaxPoolOp: XTen_Op<"conv2d_lrelu_maxpool", [NoSideEffect]> {
+def XTen_Conv2dLReLUMaxPoolOp: XTen_Op<"conv2d_lrelu_maxpool", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchTensorType:$weight,
@@ -502,7 +502,7 @@ def XTen_Conv2dLReLUMaxPoolOp: XTen_Op<"conv2d_lrelu_maxpool", [NoSideEffect]> {
 	}];
 }
 
-def XTen_Conv2dLReLUPadMaxPoolOp: XTen_Op<"conv2d_lrelu_pad_maxpool", [NoSideEffect]> {
+def XTen_Conv2dLReLUPadMaxPoolOp: XTen_Op<"conv2d_lrelu_pad_maxpool", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchTensorType:$weight,
@@ -536,7 +536,7 @@ def XTen_Conv2dLReLUPadMaxPoolOp: XTen_Op<"conv2d_lrelu_pad_maxpool", [NoSideEff
 	}];
 }
 
-def XTen_Conv2dTensorAddOp: XTen_Op<"conv2d_tensoradd", [NoSideEffect]> {
+def XTen_Conv2dTensorAddOp: XTen_Op<"conv2d_tensoradd", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchTensorType:$weight,
@@ -559,7 +559,7 @@ def XTen_Conv2dTensorAddOp: XTen_Op<"conv2d_tensoradd", [NoSideEffect]> {
   }];
 }
 
-def XTen_Conv2dTensorAddReLUOp: XTen_Op<"conv2d_tensoradd_relu", [NoSideEffect]> {
+def XTen_Conv2dTensorAddReLUOp: XTen_Op<"conv2d_tensoradd_relu", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchTensorType:$weight,
@@ -582,7 +582,7 @@ def XTen_Conv2dTensorAddReLUOp: XTen_Op<"conv2d_tensoradd_relu", [NoSideEffect]>
   }];
 }
 
-def XTen_Conv2dTensorAddLReLUOp: XTen_Op<"conv2d_tensoradd_lrelu", [NoSideEffect]> {
+def XTen_Conv2dTensorAddLReLUOp: XTen_Op<"conv2d_tensoradd_lrelu", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchTensorType:$weight,
@@ -607,7 +607,7 @@ def XTen_Conv2dTensorAddLReLUOp: XTen_Op<"conv2d_tensoradd_lrelu", [NoSideEffect
   }];
 }
 
-def XTen_Conv2dTensorAddGlobalAveragePoolOp: XTen_Op<"conv2d_tensoradd_globalaveragepool", [NoSideEffect]> {
+def XTen_Conv2dTensorAddGlobalAveragePoolOp: XTen_Op<"conv2d_tensoradd_globalaveragepool", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchTensorType:$weight,
@@ -631,7 +631,7 @@ def XTen_Conv2dTensorAddGlobalAveragePoolOp: XTen_Op<"conv2d_tensoradd_globalave
   }];
 }
 
-def XTen_Conv2dTensorAddReLUGlobalAveragePoolOp: XTen_Op<"conv2d_tensoradd_relu_globalaveragepool", [NoSideEffect]> {
+def XTen_Conv2dTensorAddReLUGlobalAveragePoolOp: XTen_Op<"conv2d_tensoradd_relu_globalaveragepool", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchTensorType:$weight,
@@ -655,7 +655,7 @@ def XTen_Conv2dTensorAddReLUGlobalAveragePoolOp: XTen_Op<"conv2d_tensoradd_relu_
   }];
 }
 
-def XTen_Conv2dTensorAddLReLUGlobalAveragePoolOp: XTen_Op<"conv2d_tensoradd_lrelu_globalaveragepool", [NoSideEffect]> {
+def XTen_Conv2dTensorAddLReLUGlobalAveragePoolOp: XTen_Op<"conv2d_tensoradd_lrelu_globalaveragepool", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchTensorType:$weight,
@@ -681,7 +681,7 @@ def XTen_Conv2dTensorAddLReLUGlobalAveragePoolOp: XTen_Op<"conv2d_tensoradd_lrel
   }];
 }
 
-def XTen_LinearOp: XTen_Op<"linear", [NoSideEffect]> {
+def XTen_LinearOp: XTen_Op<"linear", [Pure]> {
   let arguments = (
     ins AnyTorchTensorType:$input,
         AnyTorchTensorType:$weight,

include/xten/Dialect/XTenNN/CMakeLists.txt

@@ -0,0 +1,17 @@
+#
+# This file is licensed 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
+#
+# (c) Copyright 2022 Advanced Micro Devices, Inc.
+
+add_custom_target(XTenNNIncGen)
+
+add_subdirectory(Interfaces)
+
+add_subdirectory(IR)
+
+add_subdirectory(Transforms)
+
+
+

include/xten/Dialect/XTenNN/IR/CMakeLists.txt

@@ -0,0 +1,18 @@
+#
+# This file is licensed 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
+#
+# (c) Copyright 2022 Advanced Micro Devices, Inc.
+
+set(LLVM_TARGET_DEFINITIONS XTenNNOps.td)
+
+mlir_tablegen(XTenNNBase.h.inc -gen-dialect-decls)
+mlir_tablegen(XTenNNBase.cpp.inc -gen-dialect-defs)
+mlir_tablegen(XTenNNOps.h.inc -gen-op-decls)
+mlir_tablegen(XTenNNOps.cpp.inc -gen-op-defs)
+
+add_public_tablegen_target(XTenNNDialectIncGen)
+add_dependencies(XTenNNIncGen XTenNNDialectIncGen)
+
+add_mlir_doc(XTenNNOps XTenNN Dialects/ -gen-dialect-doc)

include/xten/Dialect/XTenNN/IR/XTenNN.h

@@ -0,0 +1,17 @@
+//===- XTenNN.h ----------------------------------------------*- C++ -*-===//
+//
+// This file is licensed 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
+//
+// (c) Copyright 2022 Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+//
+// Convenience include for the XTenNN dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "xten/Dialect/XTenNN/IR/XTenNNOps.h"

include/xten/Dialect/XTenNN/IR/XTenNNBase.h

@@ -0,0 +1,30 @@
+//===- XTenNNBase.h --------------------------------------------*- C++ -*-===//
+//
+// This file is licensed 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
+//
+// (c) Copyright 2022 Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+//
+// Declaration of the XTenNN dialect base.
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+#include <cstddef>
+#include <cstdint>
+
+//===- Generated includes -------------------------------------------------===//
+
+#include "xten/Dialect/XTenNN/IR/XTenNNBase.h.inc"
+
+//===----------------------------------------------------------------------===//

include/xten/Dialect/XTenNN/IR/XTenNNBase.td

@@ -0,0 +1,39 @@
+//===-- XTenNNBase.td - XTenNN dialect definitions *- tablegen -*------===//
+//
+// 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
+//
+// (c) Copyright 2022 Advanced Micro Devices, Inc.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef XTENNN_BASE
+#define XTENNN_BASE
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/OpBase.td"
+
+include "xten/Dialect/XTenNN/Interfaces/EnclaveOpInterface.td"
+
+def XTenNN_Dialect : Dialect {
+  let name = "xten_nn";
+  let cppNamespace = "::amd::xten_nn";
+  let summary = "The XTen Neural Network dialect";
+  let description = [{
+    The `xten_nn` dialect provides basic graph-like semantics. It implements a subgraph 
+    operation that can wrap operations into regions and interfaces to handle those.
+  }];
+
+  code extraClassDeclaration = [{
+    private:
+        void registerOps();
+
+    }];
+
+  let useDefaultTypePrinterParser = 0;
+
+  let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
+}
+
+#endif // XTENNN_BASE