Clean up krnl.dim, krnl.shape, and krnl.getref (onnx#2733)

* Remove krnl.dim and krnl.shape

Signed-off-by: Tung D. Le <[email protected]>

* Clean up krnl.getref

Signed-off-by: Tung D. Le <[email protected]>

---------

Signed-off-by: Tung D. Le <[email protected]>

Author: tungld

docs/Dialects/krnl.md

@@ -313,37 +313,6 @@ intend to optimize.
 | :----: | ----------- |
 «unnamed» | variadic of any type
 
-### `krnl.dim` (KrnlDimOp)
-
-_Krnl dimensions operation._
-
-Emits the dimension of a MemRef independent of the MemRef alloc:
-
-```
-"krnl.dim"(%memref, %index)
-```
-
-The index identifies the dimension within the shape which is going to be emitted.
-Initially the krnl.dim operation depends on the alloc of the MemRef.
-Unlike the std.dim operation which maintains a dependency on the alloc of the MemRef, the dimension emitted by krnl.dim will not depend on the alloc operation of the MemRef once the krnl.dim operation is lowered.
-
-Any changes to the original MemRef size after the krnl.dim has been lowered will not be picked up by the emitted dimension. This allows the original MemRef to be safely modified via code transformations or affine map normalization without the risk of changing the value already emitted via krnl.dim.
-
-Traits: MemRefsNormalizable
-
-#### Operands:
-
-| Operand | Description |
-| :-----: | ----------- |
-| `alloc` | memref of any type values
-| `index` | index
-
-#### Results:
-
-| Result | Description |
-| :----: | ----------- |
-| `dimension` | index
-
 ### `krnl.entry_point` (KrnlEntryPointOp)
 
 _Indicate ONNX entry point_
@@ -429,34 +398,6 @@ current tile being iterated over.
 | :----: | ----------- |
 | `ind_var_vals` | variadic of any type
 
-### `krnl.getref` (KrnlGetRefOp)
-
-_Krnl a MemRef from within another MemRef starting at a specific offset._
-
-    Retrieves a MemRef from within another MemRef:
-
-```
-    "krnl.getref"(%memref, %offset)
-```
-    The offset is an integer which is used as an index into the input MemRef. It works
-    just like an array index.
-
-Traits: MemRefsNormalizable
-
-#### Operands:
-
-| Operand | Description |
-| :-----: | ----------- |
-| `mempool` | memref of any type values
-| `offset` | integer
-| `value` | variadic of index
-
-#### Results:
-
-| Result | Description |
-| :----: | ----------- |
-| `output` | memref of any type values
-
 ### `krnl.global` (KrnlGlobalOp)
 
 _Krnl global operation_
@@ -1234,30 +1175,6 @@ Traits: MemRefsNormalizable
 | `seq` | memref of any type values
 | `index` | index
 
-### `krnl.shape` (KrnlShapeOp)
-
-_Krnl operation to retrieve the shape of a MemRef._
-
-Extracts the shape of a MemRef:
-```
-  "krnl.shape"(%memref)
-```
-The return result is of `shape.type`.
-
-Traits: MemRefsNormalizable
-
-#### Operands:
-
-| Operand | Description |
-| :-----: | ----------- |
-| `alloc` | memref of any type values
-
-#### Results:
-
-| Result | Description |
-| :----: | ----------- |
-| `shape` | memref of any type values
-
 ### `krnl.specialized_kernel` (KrnlSpecializedKernel)
 
 _Krnl specialized kernel op_

src/Compiler/CompilerPasses.cpp

@@ -193,9 +193,6 @@ void addONNXToKrnlPasses(mlir::PassManager &pm, int optLevel, bool enableCSE,
   // from ONNX dialect to Standard dialect exposes additional canonicalization
   // opportunities.
   pm.addPass(mlir::createCanonicalizerPass());
-  pm.addNestedPass<func::FuncOp>(
-      onnx_mlir::createDisconnectKrnlDimFromAllocPass());
-  pm.addPass(mlir::createCanonicalizerPass());
 }
 
 void addKrnlToAffinePasses(mlir::PassManager &pm) {
@@ -315,4 +312,4 @@ void addPasses(mlir::OwningOpRef<ModuleOp> &module, mlir::PassManager &pm,
     addKrnlToLLVMPasses(pm, outputNameNoExt, /*enableCSE=*/true);
 }
 
-} // namespace onnx_mlir
+} // namespace onnx_mlir

src/Conversion/KrnlToAffine/ConvertKrnlToAffine.cpp

@@ -773,8 +773,6 @@ void ConvertKrnlToAffinePass::runOnOperation() {
   ConversionTarget target(*ctx);
   // Legal/illegal ops.
   target.addIllegalOp<KrnlTerminatorOp>();
-  // krnl.dim operations must be lowered prior to this pass.
-  target.addIllegalOp<KrnlDimOp>();
   target.addIllegalOp<KrnlMatMulOp>();
   target.addIllegalOp<KrnlCopyToBufferOp>();
   target.addIllegalOp<KrnlCopyFromBufferOp>();

src/Conversion/KrnlToLLVM/CMakeLists.txt

@@ -5,7 +5,6 @@ add_onnx_mlir_library(OMKrnlToLLVM
   KrnlFindIndex.cpp
   KrnlCall.cpp
   KrnlEntryPoint.cpp
-  KrnlGetRef.cpp  
   KrnlGlobal.cpp
   KrnlInstrument.cpp
   KrnlMemcpy.cpp

src/Conversion/KrnlToLLVM/ConvertKrnlToLLVM.cpp

@@ -956,7 +956,6 @@ void populateKrnlToLLVMConversion(LLVMTypeConverter &typeConverter,
   krnl::populateLoweringKrnlCallOpPattern(typeConverter, patterns, ctx);
   krnl::populateLoweringKrnlFindIndexOpPattern(typeConverter, patterns, ctx);
   krnl::populateLoweringKrnlGlobalOpPattern(typeConverter, patterns, ctx);
-  krnl::populateLoweringKrnlGetRefOpPattern(typeConverter, patterns, ctx);
   krnl::populateLoweringKrnlInstrumentOpPattern(typeConverter, patterns, ctx);
   krnl::populateLoweringKrnlMemcpyOpPattern(typeConverter, patterns, ctx);
   krnl::populateLoweringKrnlPrintOpPattern(typeConverter, patterns, ctx);

src/Conversion/KrnlToLLVM/ConvertKrnlToLLVM.hpp

@@ -67,9 +67,6 @@ void populateLoweringKrnlFindIndexOpPattern(
     mlir::LLVMTypeConverter &typeConverter, mlir::RewritePatternSet &patterns,
     mlir::MLIRContext *ctx);
 
-void populateLoweringKrnlGetRefOpPattern(mlir::LLVMTypeConverter &typeConverter,
-    mlir::RewritePatternSet &patterns, mlir::MLIRContext *ctx);
-
 void populateLoweringKrnlGlobalOpPattern(mlir::LLVMTypeConverter &typeConverter,
     mlir::RewritePatternSet &patterns, mlir::MLIRContext *ctx);
 

src/Conversion/KrnlToLLVM/KrnlGetRef.cpp

@@ -257,20 +257,37 @@ struct ONNXCategoryMapperOpLowering
         .Case<IntegerType>(
             [&](IntegerType) { inputElem = createKrnl.load(memref, loopInd); })
         .Case<krnl::StringType>([&](krnl::StringType stringType) {
-          MathBuilder createMath(createKrnl);
-          Value zero = createMath.constant(
-              createMath.getBuilder().getIntegerType(64), 0);
           ArrayRef<int64_t> shape =
               memref.getType().cast<ShapedType>().getShape();
           SmallVector<int64_t, 4> newShape;
-          for (uint64_t i = 0; i < shape.size(); i++)
-            newShape.emplace_back(
-                (shape[i] == ShapedType::kDynamic) ? 1 : shape[i]);
-          auto memRefType = MemRefType::get(
-              newShape, krnl::StringType::get(elementType.getContext()));
-          // Sole use of krnl.getRef.
-          Value stringMemRef = createKrnl.getRef(memRefType, memref, zero);
-          inputElem = createKrnl.load(stringMemRef, loopInd);
+          bool hasDynamicDim = false;
+          for (uint64_t i = 0; i < shape.size(); i++) {
+            if (shape[i] == ShapedType::kDynamic) {
+              newShape.emplace_back(1);
+              hasDynamicDim = true;
+            } else {
+              newShape.emplace_back(shape[i]);
+            }
+          }
+          if (!hasDynamicDim) {
+            inputElem = createKrnl.load(memref, loopInd);
+          } else {
+            MemRefBuilder createMemRef(createKrnl);
+            MemRefType memRefType = MemRefType::get(
+                newShape, krnl::StringType::get(elementType.getContext()));
+            SmallVector<int64_t, 4> offsets(shape.size(), 0);
+            SmallVector<int64_t, 4> strides;
+            int64_t alignmentOffset; // not used, just to make the function call
+                                     // completed.
+            if (getStridesAndOffset(memRefType, strides, alignmentOffset)
+                    .failed())
+              llvm_unreachable("Failed to get strides");
+            Value stringMemRef =
+                createMemRef
+                    .subView(memRefType, memref, offsets, newShape, strides)
+                    .getResult();
+            inputElem = createKrnl.load(stringMemRef, loopInd);
+          }
         })
         .Default([&](Type type) {
           llvm::errs() << "type: " << type << "\n";

src/Conversion/ONNXToKrnl/ML/CategoryMapper.cpp

@@ -211,19 +211,10 @@ void KrnlBuilder::matmul(Value A, ValueRange aStart, Value B, ValueRange bStart,
       globalUBs[1], globalUBs[2], simdize, unroll, overCompute);
 }
 
-Value KrnlBuilder::dim(Type type, Value alloc, Value index) const {
-  return b().create<KrnlDimOp>(loc(), type, alloc, index);
-}
-
 KrnlMovableOp KrnlBuilder::movable() const {
   return b().create<KrnlMovableOp>(loc());
 }
 
-KrnlGetRefOp KrnlBuilder::getRef(
-    Type type, Value memref, Value offset, ValueRange indices) const {
-  return b().create<KrnlGetRefOp>(loc(), type, memref, offset, indices);
-}
-
 Value KrnlBuilder::constant(MemRefType type, StringRef name,
     std::optional<Attribute> value, std::optional<IntegerAttr> offset,
     std::optional<IntegerAttr> alignment) const {