[quidditch_snitch] Rename cluster_index to compute_core_index (#108)

The op as it is currently used and specified returns the index of compute cores (in our current simulations 0 to exclusive 8) and has unspecified behavior on the DMA core. This will be used in a later PR to properly lower it when doing DMA code specialization.
opencompl · Aug 15, 2024 · 60fa26e · 60fa26e
1 parent edaccab
commit 60fa26e
Show file tree

Hide file tree

Showing 7 changed files with 26 additions and 18 deletions.
diff --git a/codegen/compiler/src/Quidditch/Conversion/ConvertSnitchToLLVM.cpp b/codegen/compiler/src/Quidditch/Conversion/ConvertSnitchToLLVM.cpp
@@ -482,18 +482,19 @@ struct CallMicrokernelOpLowering : ConvertOpToLLVMPattern<CallMicrokernelOp> {
   }
 };
 
-struct ClusterIndexOpLowering : ConvertOpToLLVMPattern<ClusterIndexOp> {
+struct ComputeCoreIndexOpLowering : ConvertOpToLLVMPattern<ComputeCoreIndexOp> {
 
-  LLVM::LLVMFuncOp clusterIndexFunc;
+  LLVM::LLVMFuncOp computeCoreIndexFunc;
 
-  ClusterIndexOpLowering(LLVM::LLVMFuncOp clusterIndexFunc,
-                         const LLVMTypeConverter &converter)
-      : ConvertOpToLLVMPattern(converter), clusterIndexFunc(clusterIndexFunc) {}
+  ComputeCoreIndexOpLowering(LLVM::LLVMFuncOp computeCoreIndexFunc,
+                             const LLVMTypeConverter &converter)
+      : ConvertOpToLLVMPattern(converter),
+        computeCoreIndexFunc(computeCoreIndexFunc) {}
 
   LogicalResult
-  matchAndRewrite(ClusterIndexOp op, OpAdaptor adaptor,
+  matchAndRewrite(ComputeCoreIndexOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, clusterIndexFunc,
+    rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, computeCoreIndexFunc,
                                               ValueRange());
     return success();
   }
@@ -525,17 +526,17 @@ void quidditch::populateSnitchToLLVMConversionPatterns(
           i32, ArrayRef<Type>{ptrType, ptrType, sizeT, sizeT, sizeT, sizeT}));
   dmaStart2D->setAttr("hal.import.bitcode", builder.getUnitAttr());
 
-  auto clusterCoreIndex = builder.create<LLVM::LLVMFuncOp>(
+  auto computeCoreIndex = builder.create<LLVM::LLVMFuncOp>(
       builder.getUnknownLoc(), "snrt_cluster_core_idx",
       LLVM::LLVMFunctionType::get(i32, ArrayRef<Type>{}));
-  clusterCoreIndex->setAttr("hal.import.bitcode", builder.getUnitAttr());
+  computeCoreIndex->setAttr("hal.import.bitcode", builder.getUnitAttr());
 
   patterns.insert<L1MemoryViewOpLowering, CompletedTokenOpLowering,
                   BarrierOpLowering, MicrokernelFenceOpLowering,
                   WaitForDMATransfersOpLowering>(typeConverter);
   patterns.insert<StartDMATransferOp1DLowering>(dmaStart1D, typeConverter);
   patterns.insert<StartDMATransferOp2DLowering>(dmaStart2D, typeConverter);
-  patterns.insert<ClusterIndexOpLowering>(clusterCoreIndex, typeConverter);
+  patterns.insert<ComputeCoreIndexOpLowering>(computeCoreIndex, typeConverter);
   patterns.insert<CallMicrokernelOpLowering>(SymbolTable(moduleOp),
                                              typeConverter);
 }
diff --git a/codegen/compiler/src/Quidditch/Dialect/Snitch/IR/QuidditchSnitchOps.td b/codegen/compiler/src/Quidditch/Dialect/Snitch/IR/QuidditchSnitchOps.td
@@ -235,8 +235,15 @@ def QuidditchSnitch_BarrierOp : QuidditchSnitch_Op<"barrier"> {
   }];
 }
 
-def QuidditchSnitch_ClusterIndexOp
-  : QuidditchSnitch_Op<"cluster_index", [Pure]> {
+def QuidditchSnitch_ComputeCoreIndexOp
+  : QuidditchSnitch_Op<"compute_core_index", [Pure]> {
+
+  let description = [{
+    Returns the index of the compute core within a given cluster.
+    This is guaranteed to return a number between 0 and exclusive
+    `compute_cores` where `compute_cores` is an `IntegerAttr` in
+    the surrounding `ExecutableTargetAttr`.
+  }];
 
   let results = (outs Index:$result);
 

diff --git a/codegen/compiler/src/Quidditch/Dialect/Snitch/Transforms/LowerForallOp.cpp b/codegen/compiler/src/Quidditch/Dialect/Snitch/Transforms/LowerForallOp.cpp
@@ -44,7 +44,7 @@ void LowerForallOp::runOnOperation() {
     forallOp.getTerminator().erase();
 
     OpBuilder builder(forallOp);
-    Value id = builder.create<ClusterIndexOp>(forallOp.getLoc());
+    Value id = builder.create<ComputeCoreIndexOp>(forallOp.getLoc());
 
     Value lb = forallOp.getLowerBound(builder).front();
     Value ub = forallOp.getUpperBound(builder).front();

diff --git a/codegen/compiler/src/Quidditch/Target/RemoveTrivialLoops.cpp b/codegen/compiler/src/Quidditch/Target/RemoveTrivialLoops.cpp
@@ -30,8 +30,8 @@ static std::optional<std::pair<AffineExpr, AffineExpr>>
 getWorkgroupRange(Value processorValue,
                   std::optional<IntegerAttr> numComputeCores,
                   ArrayRef<int64_t> workgroupCount) {
-  if (auto idOp =
-          processorValue.getDefiningOp<quidditch::Snitch::ClusterIndexOp>()) {
+  if (auto idOp = processorValue
+                      .getDefiningOp<quidditch::Snitch::ComputeCoreIndexOp>()) {
     if (!numComputeCores)
       return std::nullopt;
 

diff --git a/...on/ConvertSnitchToLLVM/cluster_index.mlir → ...nvertSnitchToLLVM/compute_core_index.mlir b/...on/ConvertSnitchToLLVM/cluster_index.mlir → ...nvertSnitchToLLVM/compute_core_index.mlir
@@ -3,6 +3,6 @@
 // CHECK-LABEL: @test
 func.func private @test() -> index {
   // CHECK: call @snrt_cluster_core_idx()
-  %0 = quidditch_snitch.cluster_index
+  %0 = quidditch_snitch.compute_core_index
   return %0 : index
 }
diff --git a/codegen/tests/Dialect/Snitch/Transforms/lower-forall.mlir b/codegen/tests/Dialect/Snitch/Transforms/lower-forall.mlir
@@ -13,7 +13,7 @@ hal.executable @test {
       // CHECK-SAME: %[[UB:[[:alnum:]]+]]
       // CHECK-SAME: %[[STEP:[[:alnum:]]+]]
       func.func @test(%lb : index, %ub : index, %step : index) {
-        // CHECK: %[[ID:.*]] = quidditch_snitch.cluster_index
+        // CHECK: %[[ID:.*]] = quidditch_snitch.compute_core_index
         // CHECK: %[[NLB:.*]] = affine.apply #[[$MAP1]](%[[LB]], %[[ID]], %[[STEP]])
         // CHECK: %[[NSTEP:.*]] = affine.apply #[[$MAP2]](%[[STEP]])
         // CHECK: scf.for %[[IV:.*]] = %[[NLB]] to %[[UB]] step %[[NSTEP]]

diff --git a/codegen/tests/Target/remove-trivial-loops.mlir b/codegen/tests/Target/remove-trivial-loops.mlir
@@ -7,7 +7,7 @@ hal.executable @test {
     builtin.module {
       // CHECK-LABEL: func @test
       func.func @test() {
-        %0 = quidditch_snitch.cluster_index
+        %0 = quidditch_snitch.compute_core_index
         %1 = arith.constant 8 : index
         // CHECK-NOT: scf.for
         scf.for %arg3 = %0 to %1 step %1 {