[SnitchDMA] Introduce SnitchDMA dialect (#128)

The DMA dialect represents DMA operations in their purest forms: Moving
data from one MemRef to another.
The real world is unfortunately not as pretty: Snitch's DMA only
supports two-dimensional transfers, the zero region is of a specific
size and DMA configuration is relatively expensive.

This PR introduces the `SnitchDMA` dialect which is meant to handle any
Snitch specific legalizations but also optimizations.

The intention for the future is to have lowering of DMA operations split
into multiple passes and phases performing 1) legalization of DMA
transfers, 2) Optimization of configuration, 3) optimization of number
of waits and 4) lowering to LLVM.

As a first step, this PR only introduces the dialect and the equivalent
of the `dmstati` instruction which is used in the current one-shot to
LLVM lowering.

codegen/compiler/src/Quidditch/Conversion/CMakeLists.txt

@@ -46,6 +46,7 @@ iree_cc_library(
         "ConvertDMAToLLVM.cpp"
         DEPS
         Quidditch::Dialect::DMA::IR::DMADialect
+        Quidditch::Dialect::SnitchDMA::IR::SnitchDMADialect
         MLIRAnalysis
         MLIRIR
         MLIRLLVMCommonConversion

codegen/compiler/src/Quidditch/Conversion/ConvertDMAToLLVM.cpp

@@ -9,8 +9,10 @@
 #include "mlir/Transforms/DialectConversion.h"
 
 #include "Quidditch/Dialect/DMA/IR/DMAOps.h"
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMAOps.h"
 
 using namespace mlir;
+using namespace quidditch;
 using namespace quidditch::dma;
 
 /// Returns the number of potentially non-contiguous outer dimensions of
@@ -414,19 +416,7 @@ struct WaitForTransfersOpLowering : ConvertOpToLLVMPattern<WaitForTransfersOp> {
     rewriter.create<LLVM::BrOp>(op->getLoc(), body);
 
     rewriter.setInsertionPointToEnd(body);
-    Value lastCompleted =
-        rewriter
-            .create<LLVM::InlineAsmOp>(
-                op->getLoc(), /*res=*/rewriter.getI32Type(),
-                /*operands=*/ValueRange(),
-                // dmstati $0, 0
-                // opcode6=0x2b, func3=0, func7=0b100, rd=$0, rs1=zero,
-                // rs2=imm5(0)
-                ".insn r 0x2b, 0, 0b100, $0, zero, zero\n",
-                /*constraints=*/"=r",
-                /*has_side_effects=*/true, /*is_align_stack=*/false,
-                /*asm_dialect=*/nullptr, /*operand_attrs=*/nullptr)
-            .getRes();
+    Value lastCompleted = rewriter.create<SnitchDMA::StatOp>(op->getLoc());
     Value notDone = rewriter.create<LLVM::ICmpOp>(
         op->getLoc(), LLVM::ICmpPredicate::ult, lastCompleted, current);
     rewriter.create<LLVM::CondBrOp>(op->getLoc(), notDone, body, after);
@@ -450,6 +440,26 @@ struct CompletedTokenOpLowering : ConvertOpToLLVMPattern<CompletedTokenOp> {
   }
 };
 
+struct StatOpLowering : ConvertOpToLLVMPattern<SnitchDMA::StatOp> {
+  using ConvertOpToLLVMPattern<SnitchDMA::StatOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(SnitchDMA::StatOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<LLVM::InlineAsmOp>(
+        op, /*res=*/rewriter.getI32Type(),
+        /*operands=*/ValueRange(),
+        // dmstati $0, 0
+        // opcode6=0x2b, func3=0, func7=0b100, rd=$0, rs1=zero,
+        // rs2=imm5(0)
+        ".insn r 0x2b, 0, 0b100, $0, zero, zero\n",
+        /*constraints=*/"=r",
+        /*has_side_effects=*/true, /*is_align_stack=*/false,
+        /*asm_dialect=*/nullptr, /*operand_attrs=*/nullptr);
+    return success();
+  }
+};
+
 } // namespace
 
 void quidditch::populateDMAToLLVMConversionPatterns(
@@ -476,7 +486,8 @@ void quidditch::populateDMAToLLVMConversionPatterns(
   dmaStart2D->setAttr("hal.import.bitcode", builder.getUnitAttr());
 
   patterns.insert<CompletedTokenOpLowering, WaitForTransfersOpLowering,
-                  StartZeroMemTransferOpOpLowering>(typeConverter);
+                  StartZeroMemTransferOpOpLowering, StatOpLowering>(
+      typeConverter);
   patterns.insert<StartTransferOp1DLowering>(dmaStart1D, typeConverter);
   patterns.insert<StartTransferOp2DLowering>(dmaStart2D, typeConverter);
   patterns.insert<StartContiguousZeroMemTransferOpOpLowering>(

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/CMakeLists.txt

@@ -0,0 +1 @@
+iree_add_all_subdirs()

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/CMakeLists.txt

@@ -0,0 +1,73 @@
+iree_add_all_subdirs()
+
+iree_cc_library(
+        NAME
+        SnitchDMADialect
+        HDRS
+        "SnitchDMADialect.h"
+        "SnitchDMAOps.h"
+        TEXTUAL_HDRS
+        "SnitchDMAAttrs.cpp.inc"
+        "SnitchDMAAttrs.h.inc"
+        "SnitchDMADialect.cpp.inc"
+        "SnitchDMADialect.h.inc"
+        "SnitchDMAOps.cpp.inc"
+        "SnitchDMAOps.h.inc"
+        "SnitchDMATypes.cpp.inc"
+        "SnitchDMATypes.h.inc"
+        SRCS
+        "SnitchDMAAttrs.cpp"
+        "SnitchDMADialect.cpp"
+        "SnitchDMAOps.cpp"
+        "SnitchDMATypes.cpp"
+        DEPS
+        ::SnitchDMAAttrsGen
+        ::SnitchDMADialectGen
+        ::SnitchDMAOpsGen
+        ::SnitchDMATypesGen
+        LLVMSupport
+        MLIRIR
+        MLIRInferTypeOpInterface
+        MLIRSupport
+        PUBLIC
+)
+
+iree_tablegen_library(
+        NAME
+        SnitchDMAOpsGen
+        TD_FILE
+        "SnitchDMAOps.td"
+        OUTS
+        --gen-op-decls SnitchDMAOps.h.inc
+        --gen-op-defs SnitchDMAOps.cpp.inc
+)
+
+iree_tablegen_library(
+        NAME
+        SnitchDMADialectGen
+        TD_FILE
+        "SnitchDMADialect.td"
+        OUTS
+        --gen-dialect-decls SnitchDMADialect.h.inc
+        --gen-dialect-defs SnitchDMADialect.cpp.inc
+)
+
+iree_tablegen_library(
+        NAME
+        SnitchDMAAttrsGen
+        TD_FILE
+        "SnitchDMAAttrs.td"
+        OUTS
+        --gen-attrdef-decls SnitchDMAAttrs.h.inc
+        --gen-attrdef-defs SnitchDMAAttrs.cpp.inc
+)
+
+iree_tablegen_library(
+        NAME
+        SnitchDMATypesGen
+        TD_FILE
+        "SnitchDMATypes.td"
+        OUTS
+        --gen-typedef-decls SnitchDMATypes.h.inc
+        --gen-typedef-defs SnitchDMATypes.cpp.inc
+)

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAAttrs.cpp

@@ -0,0 +1 @@
+#include "SnitchDMAAttrs.h"

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAAttrs.h

@@ -0,0 +1,7 @@
+
+#pragma once
+
+#include "mlir/IR/Attributes.h"
+
+#define GET_ATTRDEF_CLASSES
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMAAttrs.h.inc"

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAAttrs.td

@@ -0,0 +1,11 @@
+#ifndef QUIDDITCH_DIALECT_SNITCHDMA_IR_SNITCHDMAATTRS
+#define QUIDDITCH_DIALECT_SNITCHDMA_IR_SNITCHDMAATTRS
+
+include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.td"
+include "mlir/IR/AttrTypeBase.td"
+
+class SnitchDMA_Attr<string name, list<Trait> traits = []> :
+  AttrDef<SnitchDMA_Dialect, name, traits>;
+
+
+#endif

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.cpp

@@ -0,0 +1,36 @@
+#include "SnitchDMADialect.h"
+
+#include "SnitchDMAAttrs.h"
+#include "SnitchDMAOps.h"
+#include "SnitchDMATypes.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+
+#define GET_ATTRDEF_CLASSES
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMAAttrs.cpp.inc"
+
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.cpp.inc"
+
+using namespace mlir;
+using namespace quidditch::SnitchDMA;
+
+//===----------------------------------------------------------------------===//
+// DMADialect
+//===----------------------------------------------------------------------===//
+
+void SnitchDMADialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMAOps.cpp.inc"
+      >();
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMAAttrs.cpp.inc"
+      >();
+  addTypes<
+#define GET_TYPEDEF_LIST
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMATypes.cpp.inc"
+      >();
+}

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.h

@@ -0,0 +1,7 @@
+
+#pragma once
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/Operation.h"
+
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.h.inc"

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.td

@@ -0,0 +1,21 @@
+#ifndef QUIDDITCH_DIALECT_SNITCHDMA_IR_SNITCHDMADIALECT
+#define QUIDDITCH_DIALECT_SNITCHDMA_IR_SNITCHDMADIALECT
+
+include "mlir/IR/DialectBase.td"
+
+def SnitchDMA_Dialect : Dialect {
+  let name = "snitch_dma";
+  let cppNamespace = "::quidditch::SnitchDMA";
+
+  let description = [{
+    Dialect dealing with all implementation details specific to Snitch's DMA
+    engine.
+    Used to progressively lower and optimize the `dma` dialect.
+  }];
+
+  let useDefaultAttributePrinterParser = 0;
+  let useDefaultTypePrinterParser = 0;
+  let hasConstantMaterializer = 0;
+}
+
+#endif

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAOps.cpp

@@ -0,0 +1,11 @@
+#include "SnitchDMAOps.h"
+
+#define GET_OP_CLASSES
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMAOps.cpp.inc"
+
+using namespace mlir;
+using namespace quidditch::SnitchDMA;
+
+StringRef QueueResource::getName() {
+  return "queue";
+}

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAOps.h

@@ -0,0 +1,20 @@
+
+#pragma once
+
+#include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#include "SnitchDMATypes.h"
+
+#define GET_OP_CLASSES
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMAOps.h.inc"
+
+namespace quidditch::SnitchDMA {
+class QueueResource : public mlir::SideEffects::Resource::Base<QueueResource> {
+public:
+  llvm::StringRef getName() override;
+};
+} // namespace quidditch::SnitchDMA

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAOps.td

@@ -0,0 +1,31 @@
+#ifndef QUIDDITCH_DIALECT_SNITCHDMA_IR_SNITCHDMAOPS
+#define QUIDDITCH_DIALECT_SNITCHDMA_IR_SNITCHDMAOPS
+
+include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.td"
+include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMATypes.td"
+include "mlir/IR/CommonTypeConstraints.td"
+include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+class SnitchDMA_Op<string mnemonic, list<Trait> traits = []> :
+  Op<SnitchDMA_Dialect, mnemonic, traits>;
+
+def SnitchDMA_QueueResource
+  : Resource<"quidditch::SnitchDMA::QueueResource">;
+
+def SnitchDMA_StatOp : SnitchDMA_Op<"stat",
+  [MemoryEffects<[MemRead<SnitchDMA_QueueResource>]>]> {
+
+  let description = [{
+    Returns the id of the last DMA transfer that has been completed.
+  }];
+
+  let results = (outs I32:$completed_id);
+
+  let assemblyFormat = [{
+    attr-dict
+  }];
+}
+
+#endif

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMATypes.cpp

@@ -0,0 +1,11 @@
+#include "SnitchDMATypes.h"
+
+#include "llvm/ADT/TypeSwitch.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+
+#include "SnitchDMADialect.h"
+
+#define GET_TYPEDEF_CLASSES
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMATypes.cpp.inc"

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMATypes.h

@@ -0,0 +1,7 @@
+
+#pragma once
+
+#include "mlir/IR/Types.h"
+
+#define GET_TYPEDEF_CLASSES
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMATypes.h.inc"

codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMATypes.td

@@ -0,0 +1,10 @@
+#ifndef QUIDDITCH_DIALECT_SNITCHDMA_IR_SNITCHDMATYPES
+#define QUIDDITCH_DIALECT_SNITCHDMA_IR_SNITCHDMATYPES
+
+include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.td"
+include "mlir/IR/AttrTypeBase.td"
+
+class SnitchDMA_Type<string name, list<Trait> traits = []> :
+  TypeDef<SnitchDMA_Dialect, name, traits>;
+
+#endif

codegen/compiler/src/Quidditch/Target/ConvertToLLVM.cpp

@@ -7,6 +7,7 @@
 #include "Quidditch/Conversion/ConvertDMAToLLVM.h"
 #include "Quidditch/Conversion/ConvertSnitchToLLVM.h"
 #include "Quidditch/Dialect/Snitch/IR/QuidditchSnitchDialect.h"
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.h"
 #include "iree/compiler/Codegen/LLVMCPU/DispatchABI.h"
 #include "iree/compiler/Codegen/LLVMCPU/PassDetail.h"
 #include "iree/compiler/Codegen/LLVMCPU/Passes.h"

codegen/compiler/src/Quidditch/Target/Passes.td

@@ -56,6 +56,7 @@ def ConvertToLLVMPass : Pass<"quidditch-convert-to-llvm", "mlir::ModuleOp"> {
     "mlir::scf::SCFDialect",
     "mlir::memref::MemRefDialect",
     "mlir::affine::AffineDialect",
+    "quidditch::SnitchDMA::SnitchDMADialect",
   ];
 }
 

codegen/compiler/src/Quidditch/Target/QuidditchTarget.cpp

@@ -31,6 +31,7 @@
 #include "Quidditch/Dialect/Snitch/IR/QuidditchSnitchDialect.h"
 #include "Quidditch/Dialect/Snitch/IR/QuidditchSnitchOps.h"
 #include "Quidditch/Dialect/Snitch/Transforms/Passes.h"
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.h"
 
 #include "compiler/plugins/target/LLVMCPU/LinkerTool.h"
 #include "compiler/plugins/target/LLVMCPU/StaticLibraryGenerator.h"
@@ -136,7 +137,8 @@ class QuidditchTargetBackend final : public IREE::HAL::TargetBackend {
 
     registry.insert<arm_neon::ArmNeonDialect, arm_sme::ArmSMEDialect,
                     quidditch::Snitch::QuidditchSnitchDialect,
-                    quidditch::dma::DMADialect>();
+                    quidditch::dma::DMADialect,
+                    quidditch::SnitchDMA::SnitchDMADialect>();
   }
 
   void getDefaultExecutableTargets(

codegen/tools/CMakeLists.txt

@@ -5,6 +5,7 @@ target_link_libraries(quidditch-opt
         Quidditch::Conversion::ConvertSnitchToLLVM
         Quidditch::Conversion::ConvertToRISCV
         Quidditch::Dialect::DMA::Extensions::DMACoreSpecializationOpInterfaceImpl
+        Quidditch::Dialect::SnitchDMA::IR::SnitchDMADialect
         Quidditch::Dialect::Snitch::IR::QuidditchSnitchDialect
         Quidditch::Dialect::Snitch::Transforms::Passes
         Quidditch::Target::Passes

codegen/tools/quidditch-opt.cpp

@@ -4,6 +4,7 @@
 #include "Quidditch/Conversion/Passes.h"
 #include "Quidditch/Dialect/DMA/Extensions/DMACoreSpecializationOpInterfaceImpl.h"
 #include "Quidditch/Dialect/DMA/IR/DMADialect.h"
+#include "Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.h"
 #include "Quidditch/Dialect/Snitch/IR/QuidditchSnitchDialect.h"
 #include "Quidditch/Dialect/Snitch/Transforms/Passes.h"
 #include "Quidditch/Target/Passes.h"
@@ -31,7 +32,8 @@ int main(int argc, char **argv) {
   quidditch::dma::registerDMACoreSpecializationOpInterface(registry);
   iree_compiler::registerAllDialects(registry);
   registry.insert<quidditch::Snitch::QuidditchSnitchDialect,
-                  quidditch::dma::DMADialect>();
+                  quidditch::dma::DMADialect,
+                  quidditch::SnitchDMA::SnitchDMADialect>();
 
   quidditch::registerPasses();
   quidditch::registerConversionPasses();