opencompl · zero9178 · Sep 1, 2024 · Sep 1, 2024
diff --git a/codegen/compiler/src/Quidditch/Conversion/CMakeLists.txt b/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

diff --git a/codegen/compiler/src/Quidditch/Conversion/ConvertDMAToLLVM.cpp b/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>(

diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/CMakeLists.txt b/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/CMakeLists.txt
@@ -0,0 +1 @@
+iree_add_all_subdirs()
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/CMakeLists.txt b/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
+)
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAAttrs.cpp b/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAAttrs.cpp
@@ -0,0 +1 @@
+#include "SnitchDMAAttrs.h"
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAAttrs.h b/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"
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAAttrs.td b/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
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.cpp b/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"
+      >();
+}
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.h b/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"
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMADialect.td b/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
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAOps.cpp b/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";
+}
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAOps.h b/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
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMAOps.td b/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
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMATypes.cpp b/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"
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMATypes.h b/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"
diff --git a/codegen/compiler/src/Quidditch/Dialect/SnitchDMA/IR/SnitchDMATypes.td b/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
diff --git a/codegen/compiler/src/Quidditch/Target/ConvertToLLVM.cpp b/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"

diff --git a/codegen/compiler/src/Quidditch/Target/Passes.td b/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",
   ];
 }
 

diff --git a/codegen/compiler/src/Quidditch/Target/QuidditchTarget.cpp b/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(

diff --git a/codegen/tools/CMakeLists.txt b/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

diff --git a/codegen/tools/quidditch-opt.cpp b/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();